import gradio as gr
from evodiff.pretrained import OA_DM_38M, D3PM_UNIFORM_38M
from evodiff.generate import generate_oaardm, generate_d3pm



def make_seq(seq_len, model_type):
    if model_type == "EvoDiff-Seq-OADM 38M":
        checkpoint = OA_DM_38M()
        model, collater, tokenizer, scheme = checkpoint
        tokeinzed_sample, generated_sequence = generate_oaardm(model, tokenizer, seq_len, batch_size=1, device='cpu')
    
    if model_type == "EvoDiff-D3PM-Uniform 38M":
        checkpoint = D3PM_UNIFORM_38M(return_all=True)
        model, collater, tokenizer, scheme, timestep, Q_bar, Q = checkpoint
        tokeinzed_sample, generated_sequence = generate_d3pm(model, tokenizer, Q, Q_bar, timestep, seq_len, batch_size=1, device='cpu')

    return generated_sequence

# iface = gr.Interface(
#     fn=make_seq,
#     inputs=gr.Slider(10, 100),
#     outputs="text"
#     )

# iface.launch()


with gr.Blocks() as edapp:
    with gr.Row():
        gr.Markdown(
            """
            # EvoDiff
            ## Generation of protein sequences and evolutionary alignments via discrete diffusion models

            Created By: Microsoft Research [Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Alex X. Lu, Nicolo Fusi, ProfileAva P. Amini, and Kevin K. Yang]
            
            Spaces App By: Colby T. Ford
            """
        )
    with gr.Row():
        gr.Markdown(
            """
            ## Unconditional sequence generation
            Generate a sequence with EvoDiff-Seq-OADM 38M
            """)
        gr.Interface(
            fn=make_seq,
            inputs=[
                gr.Slider(10, 100, label = "Sequence Length")
                gr.Dropdown(["EvoDiff-Seq-OADM 38M", "EvoDiff-D3PM-Uniform 38M"], type="value"),
                gr.()
                ],
            outputs="text"
            )

if __name__ == "__main__":
    edapp.launch()