import gradio as gr
import os
import pandas as pd
import tempfile
from pathlib import Path
from inference_wrapper import RNAFoldingPredictor
from huggingface_hub import hf_hub_download  # ✅ NEW IMPORT

# 🔽 Download the model file from Hugging Face Hub (ptope/rna-model-weights repo)
ckpt_file = hf_hub_download(
    repo_id="ptope/rna-model-weights",
    filename="best_val_model.pt"
)

# Set this path in case RNAFoldingPredictor expects a path to the model file
MODEL_PATH = ckpt_file

# Static directory for temporary visualizations
STATIC_DIR = "static"
os.makedirs(STATIC_DIR, exist_ok=True)

# ✅ Initialize the predictor with the downloaded checkpoint path
predictor = RNAFoldingPredictor(MODEL_PATH)

# Load test sequences
test_sequences = pd.read_csv("test_sequences.csv")
test_seq_list = test_sequences["sequence"].dropna().unique().tolist()

def create_pdb_from_prediction(prediction_df):
    pdb_lines = ["HEADER    RNA STRUCTURE PREDICTION"]
    atom_index = 1
    prev_c1_index = None
    for _, row in prediction_df.iterrows():
        resname = row['resname']
        resid = row['resid']
        x, y, z = row['x_1'], row['y_1'], row['z_1']
        c1_index = atom_index
        pdb_lines.append(f"ATOM  {atom_index:5d}  C1'  {resname} A{resid:4d}    {x:8.3f}{y:8.3f}{z:8.3f}  1.00  0.00           C")
        atom_index += 1
        base_offset = {'A': (1.5, 0.0, 0.5), 'C': (1.2, 0.3, 0.3), 'G': (1.7, -0.2, 0.7), 'U': (1.0, 0.5, 0.0)}.get(resname, (1.5, 0.0, 0.0))
        base_x, base_y, base_z = x + base_offset[0], y + base_offset[1], z + base_offset[2]
        base_index = atom_index
        pdb_lines.append(f"ATOM  {atom_index:5d}  N9   {resname} A{resid:4d}    {base_x:8.3f}{base_y:8.3f}{base_z:8.3f}  1.00  0.00           N")
        atom_index += 1
        pdb_lines.append(f"CONECT{c1_index:5d}{base_index:5d}")
        if prev_c1_index is not None:
            pdb_lines.append(f"CONECT{prev_c1_index:5d}{c1_index:5d}")
        prev_c1_index = c1_index
    pdb_lines.append("END")
    return "\n".join(pdb_lines)

def predict_and_serve(sequence, description=""):
    print("===> ENTERED predict_and_serve")
    print(f"Input sequence: {sequence}")
    sequence = sequence.strip().upper()
    if not sequence or not all(b in "ACGU" for b in sequence):
        return "Invalid input", None, ""
    df = predictor.predict(sequence, description)
    pdb_text = create_pdb_from_prediction(df)
    filename = f"{next(tempfile._get_candidate_names())}.pdb"
    path = Path(STATIC_DIR) / filename

    print(f"[DEBUG] Writing PDB to: {path}")
    print(f"[DEBUG] PDB content:\n{pdb_text[:500]}")  # preview the first 500 chars
    
    with open(path, "w") as f:
        f.write(pdb_text)
    
    space_id = os.environ.get("SPACE_ID", "your-username--space-name")  # fallback if run locally
    iframe_url = f"https://molstar.org/viewer/?loadFromUrl=https://{space_id}.hf.space/file=static/{filename}"
    iframe_html = f'<iframe src="{iframe_url}" width="100%" height="600px" style="border:1px solid #ccc;"></iframe>'
    return f"Predicted {len(df)} residues.", str(path), iframe_html

# Interface
with gr.Blocks() as demo:
    gr.Markdown("## 🧬 RNA 3D Viewer (Hugging Face + Mol*)")

    with gr.Row():
        with gr.Column():
            dropdown = gr.Dropdown(choices=test_seq_list, label="Select a Test Sequence")
            seq_input = gr.Textbox(label="RNA Sequence", lines=4)
            desc_input = gr.Textbox(label="Description (optional)", lines=1)
            dropdown.change(fn=lambda s: s, inputs=dropdown, outputs=seq_input)
            submit_btn = gr.Button("Submit")

        with gr.Column():
            status_output = gr.Textbox(label="Status")
            file_output = gr.File(label="Download .pdb")

    viewer_html = gr.HTML(label="Mol* Viewer")

    submit_btn.click(
        fn=predict_and_serve,
        inputs=[seq_input, desc_input],
        outputs=[status_output, file_output, viewer_html]
    )

demo.launch()