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

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


def generate_structure(sequence, style):
    try:
        inputs = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)

        with torch.no_grad():
            outputs = model(**inputs)

        pdb = model.output_to_pdb(outputs)[0]

        # Save file
        file_path = "output.pdb"
        with open(file_path, "w") as f:
            f.write(pdb)

        # 🔥 SAFE visualization
        view = py3Dmol.view(width=800, height=500)
        view.addModel(pdb, "pdb")

        if style == "Cartoon":
            view.setStyle({"cartoon": {"color": "spectrum"}})
        elif style == "Stick":
            view.setStyle({"stick": {}})
        elif style == "Sphere":
            view.setStyle({"sphere": {}})

        view.zoomTo()

        # ✅ safer HTML export
        html = view._make_html()
        html = f"""
        <div style="width:100%; height:500px;">
            {html}
        </div>
        """

        preview = "\n".join(pdb.split("\n")[:20])

        return file_path, html, preview

    except Exception as e:
        return None, f"<h3>Error: {str(e)}</h3>", ""


# 🎛️ UI
iface = gr.Interface(
    fn=generate_structure,
    inputs=[
        gr.Textbox(label="Protein Sequence"),
        gr.Radio(["Cartoon", "Stick", "Sphere"], label="Visualization Style", value="Cartoon")
    ],
    outputs=[
        gr.File(label="Download PDB"),
        gr.HTML(label="3D Visualization"),
        gr.Textbox(label="File Preview")
    ],
    title="🧬 Protein Structure + Visualization (ESMFold)",
    description="Enter a protein sequence to generate 3D structure and visualize it."
)

iface.launch()