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


download_weights()



@spaces.GPU(duration=300)  
def test_rfd3_from_notebook():
    # Set seed for reproducibility
    seed_everything(0)

    # Configure RFD3 inference
    config = RFD3InferenceConfig(
        specification={
            'length': 40,  # Generate 80-residue proteins
        },
        diffusion_batch_size=2,  # Generate 2 structures per batch
    )

    # Initialize engine and run generation
    try: 
        model = RFD3InferenceEngine(**config)
        outputs = model.run(
            inputs=None,      # None for unconditional generation
            out_dir=None,     # None to return in memory (no file output)
            n_batches=1,      # Generate 1 batch
        )
        return_str = "RDF3 test passed! Generated structures:\n"

        for idx, data in outputs.items():
            return_str += f"Batch {idx}: {len(data)} structure(s)\n"
            for i, struct in enumerate(data):
                return_str += f"Structure {i+1}: {struct.atom_array.array_length()} Atoms\n"
                #return_str += struct.atom_array
        return return_str
    except Exception as e:
        return f"Error: {str(e)}"

    
# 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
        )
    # Configure RFD3 inference


    # Initialize engine and run generation
    @spaces.GPU(duration=300) 
    def unconditional_generation(num_batches, num_designs_per_batch, length):

        config = RFD3InferenceConfig(
        specification={
            'length': length,
        },
        diffusion_batch_size=num_designs_per_batch,  # Generate 2 structures per batch
        )
        try: 
            model = RFD3InferenceEngine(**config)
            outputs = model.run(
                inputs=None,      # None for unconditional generation
                out_dir=None,     # None to return in memory (no file output)
                n_batches=num_batches,      # Generate 1 batch
            )
            return_str = "RDF3 test passed! Generated structures:\n"

            for idx, data in outputs.items():
                return_str += f"Batch {idx}: {len(data)} structure(s)\n"
                for i, struct in enumerate(data):
                    return_str += f"Structure {i+1}: {struct.atom_array.array_length()} Atoms\n"
                    #return_str += struct.atom_array
            return return_str
        
        except Exception as e:
            return f"Error: {str(e)}"
    
    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()