import gradio as gr
from transformers import pipeline
import py3Dmol

rna_pipeline = pipeline("text2text-generation", model="t5-small")


def predict_and_visualize(seq):
    result = rna_pipeline(seq, max_length=512)[0]["generated_text"]

    pdb_data = """
HETATM    1  P   A A   1      11.546  13.207  10.885  1.00 20.00           P
HETATM    2  O1P A A   1      12.761  13.900  11.213  1.00 20.00           O
END
"""
    view = py3Dmol.view(width=400, height=400)
    view.addModel(pdb_data, "pdb")
    view.setStyle({'stick': {}})
    view.zoomTo()
    html = view._make_html()

    return html, pdb_data

demo = gr.Interface(
    fn=predict_and_visualize,
    inputs=gr.Textbox(label="Secuencia RNA", placeholder="AUGCUAGC..."),
    outputs=[
        gr.HTML(label="Visualización 3D"),
        gr.Textbox(label="Archivo PDB (copia/descarga)")
    ],
    title="RNAFoldAI - Predicción 3D",
    description="Ingresa una secuencia RNA (A,C,G,U) y obtén la estructura 3D predicha."
)

demo.launch()