import gradio as gr
import warnings
import os
import subprocess
from pathlib import Path
import shutil
import spaces
from atomworks.io.utils.visualize import view
from lightning.fabric import seed_everything
from rfd3.engine import RFD3InferenceConfig, RFD3InferenceEngine
from utils import download_weights
from utils.pipelines import test_rfd3_from_notebook, unconditional_generation


download_weights()

    
# Gradio UI
with gr.Blocks(title="RFD3 Test") as demo:
    gr.Markdown("# RFdiffusion3 (RFD3) for Backbone generation")
    gr.Markdown("Models auto-downloaded on launch. Click to test.")

    
    test_btn = gr.Button("Run RFD3 Test")
    output = gr.Textbox(label="Test Result")
    
    test_btn.click(test_rfd3_from_notebook, outputs=output)

    gr.Markdown("Unconditional generation of backbones")
    with gr.Row():
        num_designs_per_batch = gr.Number(
            value=2,
            label="Number of Designs per Batch",
            precision=0,
            minimum=1,
            maximum=8
        )
        num_batches = gr.Number(
            value=5,
            label="Number of Batches",
            precision=0,
            minimum=1,
            maximum=10
        )
        length = gr.Number(
            value=40,
            label="Length of Protein (number of residues)",
            precision=0,
            minimum=10,
            maximum=200
        )

    gen_btn = gr.Button("Run Unconditional Generation")
    gen_output = gr.Textbox(label="Generation Result")
    gen_btn.click(unconditional_generation, inputs=[num_batches, num_designs_per_batch, length], outputs=gen_output)



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