import os
import gradio as gr

import torch
from foldingdiff import sampling
from foldingdiff import angles_and_coords as ac

def sample_at_length(l:int, seed:int):
    """
    Sample a single structure at the given length
    """
    torch.manual_seed(seed)
    l = int(l)
    s = sampling.sample_simple("wukevin/foldingdiff_cath", n=1, sweep_lengths=(l, l+1))[0]
    # Create a PDB file
    outdir = os.path.join(os.getcwd(), "output")
    os.makedirs(outdir, exist_ok=True)
    pdb_file = ac.create_new_chain_nerf(os.path.join(outdir, "generated.pdb"), s)
    return s, pdb_file

interface = gr.Interface(
    fn=sample_at_length,
    inputs=[
        gr.Number(value=55, label="Protein backbone length to generate", show_label=True, precision=0),
        gr.Number(value=8964, label="Random seed", show_label=True, precision=0),
    ],
    outputs=[
        gr.Dataframe(label="Generated angles defining structure", max_rows=8),
        gr.File(label="Generated structure in PDB format (cartesian coordinates)"),
    ],
)
interface.launch()