import gradio as gr
import torch
from transformers import EsmForProteinFolding, EsmTokenizer
import py3Dmol

# Load model
model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1")
tokenizer = EsmTokenizer.from_pretrained("facebook/esmfold_v1")

def predict_structure(sequence):
    # Tokenize sequence
    inputs = tokenizer(sequence, return_tensors="pt", add_special_tokens=False)
    
    # Predict structure (ML happens here)
    with torch.no_grad():
        output = model(**inputs)
    
    # Extract PDB string
    pdb_str = output["pdb"]

    # Visualize structure
    viewer = py3Dmol.view(width=400, height=400)
    viewer.addModel(pdb_str, "pdb")
    viewer.setStyle({"cartoon": {"color": "spectrum"}})
    viewer.zoomTo()
    
    return viewer.show(), pdb_str

# Gradio UI
iface = gr.Interface(
    fn=predict_structure,
    inputs=gr.Textbox(label="Enter Amino Acid Sequence", placeholder="MKTAYIAKQRQISFVK..."),
    outputs=[
        gr.HTML(label="3D Structure"),
        gr.Textbox(label="PDB Output")
    ],
    title="Protein Structure Prediction",
    description="Enter a protein sequence to predict its 3D structure using ESMFold."
)

iface.launch()