import os
import tempfile
import numpy as np
import gradio as gr

from ase.io import read
from ase.optimize import LBFGS
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from ase.md.verlet import VelocityVerlet
from ase.io.trajectory import Trajectory
from ase import units

# Visualizador 3D “pro” si está disponible
try:
    from gradio_molecule3d import Molecule3D
    HAVE_MOL3D = True
except Exception:
    HAVE_MOL3D = False

# --- Helpers de visualización: fallback 3Dmol.js si no hay Molecule3D ---
def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
    traj = Trajectory(traj_path)
    xyz_frames = []
    for atoms in traj:
        symbols = atoms.get_chemical_symbols()
        coords = atoms.get_positions()
        parts = [str(len(symbols)), "frame"]
        for s, (x, y, z) in zip(symbols, coords):
            parts.append(f"{s} {x:.6f} {y:.6f} {z:.6f}")
        xyz_frames.append("\n".join(parts))

    html = f"""
<div id="viewer_md" style="width:{width}px; height:{height}px;"></div>
<script src="https://3dmol.org/build/3Dmol-min.js"></script>
<script>
(function() {{
  var viewer = $3Dmol.createViewer("viewer_md", {{backgroundColor: 'white'}});
  var frames = {xyz_frames!r};
  var i = 0;
  function show(i) {{
    viewer.clear();
    viewer.addModel(frames[i], "xyz");
    viewer.setStyle({{}}, {{stick: {{}}}});
    viewer.zoomTo();
    viewer.render();
  }}
  if(frames.length>0) show(0);
  if(frames.length>1) setInterval(function(){{
    i=(i+1)%frames.length; show(i);
  }}, {int(interval_ms)});
}})();
</script>
"""
    return html

# --- OrbMol directo para SPE ---
from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator

_model_calc = None
def _load_orbmol_calc():
    global _model_calc
    if _model_calc is None:
        orbff = pretrained.orb_v3_conservative_inf_omat(device="cpu", precision="float32-high")
        _model_calc = ORBCalculator(orbff, device="cpu")
    return _model_calc

def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
    try:
        calc = _load_orbmol_calc()
        if not xyz_content.strip():
            return "Error: Please enter XYZ coordinates", "Error"

        with tempfile.NamedTemporaryFile(mode="w", suffix=".xyz", delete=False) as f:
            f.write(xyz_content)
            xyz_file = f.name

        atoms = read(xyz_file)
        atoms.info = {"charge": int(charge), "spin": int(spin_multiplicity)}
        atoms.calc = calc

        energy = atoms.get_potential_energy()  # eV
        forces = atoms.get_forces()            # eV/Å

        lines = [f"Total Energy: {energy:.6f} eV", "", "Atomic Forces:"]
        for i, f in enumerate(forces):
            lines.append(f"Atom {i+1}: [{f[0]:.4f}, {f[1]:.4f}, {f[2]:.4f}] eV/Å")
        max_force = float(np.max(np.linalg.norm(forces, axis=1)))
        lines += ["", f"Max Force: {max_force:.4f} eV/Å"]

        try:
            os.unlink(xyz_file)
        except Exception:
            pass

        return "\n".join(lines), "Calculation completed with OrbMol"
    except Exception as e:
        return f"Error during calculation: {e}", "Error"

# --- Importa las rutinas FAIRChem-like para MD/Relax (ya soportan string XYZ o ruta) ---
from simulation_scripts_orbmol import (
    run_md_simulation,
    run_relaxation_simulation,
    last_frame_xyz_from_traj,
)

# Wrappers para conectar outputs de Gradio correctamente (string XYZ / HTML, file, logs...)
def md_wrapper(xyz_content, charge, spin, steps, tempK, timestep_fs, ensemble):
    try:
        traj_path, log_text, script_text, explanation = run_md_simulation(
            xyz_content,            # acepta string XYZ
            int(steps),
            20,                     # pre-relax steps como en la UI de UMA
            float(timestep_fs),
            float(tempK),
            "NVT" if ensemble == "NVT" else "NVE",
            int(charge),
            int(spin),
        )
        status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
        # Viewer
        if HAVE_MOL3D:
            xyz_final = last_frame_xyz_from_traj(traj_path)  # value para Molecule3D
            viewer_value = xyz_final
            html_value = None
        else:
            viewer_value = None
            html_value = traj_to_html(traj_path)

        return (
            status,                # MD Status
            viewer_value,          # Molecule3D value (o None)
            html_value,            # HTML fallback (o None)
            traj_path,             # File download
            log_text,              # Log
            script_text,           # Script
            explanation,           # Explanation
        )
    except Exception as e:
        return (f"Error: {e}", None, None, None, "", "", "")

def relax_wrapper(xyz_content, steps, fmax, charge, spin, relax_cell):
    try:
        traj_path, log_text, script_text, explanation = run_relaxation_simulation(
            xyz_content,
            int(steps),
            float(fmax),
            int(charge),
            int(spin),
            bool(relax_cell),
        )
        status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)"
        if HAVE_MOL3D:
            viewer_value = last_frame_xyz_from_traj(traj_path)
            html_value = None
        else:
            viewer_value = None
            html_value = traj_to_html(traj_path)

        return (
            status,
            viewer_value,
            html_value,
            traj_path,
            log_text,
            script_text,
            explanation,
        )
    except Exception as e:
        return (f"Error: {e}", None, None, None, "", "", "")

# ------------------------
#   Ejemplos rápidos
# ------------------------
examples = [
    ["""2
Hydrogen molecule
H 0.0 0.0 0.0
H 0.0 0.0 0.74""", 0, 1],
    ["""3
Water molecule
O 0.0000 0.0000 0.0000
H 0.7571 0.0000 0.5864
H -0.7571 0.0000 0.5864""", 0, 1],
    ["""5
Methane
C 0.0000 0.0000 0.0000
H 1.0890 0.0000 0.0000
H -0.3630 1.0267 0.0000
H -0.3630 -0.5133 0.8887
H -0.3630 -0.5133 -0.8887""", 0, 1],
]

# ------------------------
#   UI Gradio
# ------------------------
with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
    with gr.Tabs():
        # ========== SPE ==========
        with gr.Tab("Single Point Energy"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("# OrbMol Demo — Quantum-Accurate Molecular Predictions")
                    gr.Markdown(
                        "Predict **energies** and **forces** with OrbMol (OMol25). "
                        "Supports **charge** and **spin multiplicity**."
                    )

                    xyz_input = gr.Textbox(
                        label="XYZ Coordinates",
                        placeholder="Paste XYZ here...",
                        lines=12,
                    )
                    with gr.Row():
                        charge_input = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
                        spin_input = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity")
                    run_spe = gr.Button("Run OrbMol Prediction", variant="primary")

                with gr.Column(variant="panel", min_width=500):
                    spe_out = gr.Textbox(label="Energy & Forces", lines=15, interactive=False)
                    spe_status = gr.Textbox(label="Status", interactive=False, max_lines=1)

            gr.Examples(examples=examples, inputs=[xyz_input, charge_input, spin_input], label="Examples")
            run_spe.click(
                predict_molecule,
                inputs=[xyz_input, charge_input, spin_input],
                outputs=[spe_out, spe_status],
            )

            with gr.Sidebar(open=True):
                gr.Markdown("## Learn more about OrbMol")
                with gr.Accordion("What is OrbMol?", open=False):
                    gr.Markdown(
                        "* Neural network potential for molecules\n"
                        "* Built on Orb-v3, trained on OMol25 (ωB97M-V/def2-TZVPD)\n"
                        "* Supports charge and spin multiplicity"
                    )
                with gr.Accordion("Benchmarks", open=False):
                    gr.Markdown(
                        "* Near-DFT accuracy on **GMTKN55** and **Wiggle150**\n"
                        "* Accurate **protein–ligand** interaction energies (PLA15)\n"
                        "* Stable long MD on biomolecules grandes"
                    )
                with gr.Accordion("Disclaimers", open=False):
                    gr.Markdown(
                        "* Validate for your use case\n"
                        "* Consider training **level of theory** and intended domain"
                    )

        # ========== MD ==========
        with gr.Tab("Molecular Dynamics"):
            with gr.Row():
                with gr.Column(scale=2):
                    xyz_md = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here...")
                    with gr.Row():
                        charge_md = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
                        spin_md = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity")
                    with gr.Row():
                        steps_md = gr.Slider(value=100, minimum=10, maximum=1000, step=10, label="Steps")
                        temp_md = gr.Slider(value=300, minimum=10, maximum=1500, step=10, label="Temperature (K)")
                    with gr.Row():
                        timestep_md = gr.Slider(value=1.0, minimum=0.1, maximum=5.0, step=0.1, label="Timestep (fs)")
                        ensemble_md = gr.Radio(choices=["NVE", "NVT"], value="NVE", label="Ensemble")
                    run_md_btn = gr.Button("Run MD Simulation", variant="primary")

                with gr.Column(variant="panel", min_width=520):
                    md_status = gr.Textbox(label="MD Status", interactive=False)
                    if HAVE_MOL3D:
                        md_viewer = Molecule3D(label="Trajectory Viewer")
                        md_html = gr.HTML(visible=False)  # placeholder para consistencia de outputs
                    else:
                        md_viewer = gr.Textbox(visible=False)  # placeholder
                        md_html = gr.HTML()

                    md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                    md_log = gr.Code(label="Log", language="text", interactive=False, lines=15, max_lines=25)
                    md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                    md_explain = gr.Markdown()

            run_md_btn.click(
                md_wrapper,
                inputs=[xyz_md, charge_md, spin_md, steps_md, temp_md, timestep_md, ensemble_md],
                outputs=[md_status, md_viewer, md_html, md_traj, md_log, md_script, md_explain],
            )

        # ========== Relaxation ==========
        with gr.Tab("Relaxation / Optimization"):
            with gr.Row():
                with gr.Column(scale=2):
                    xyz_rlx = gr.Textbox(label="XYZ Coordinates", lines=12, placeholder="Paste XYZ here...")
                    steps_rlx = gr.Slider(value=300, minimum=1, maximum=1000, step=1, label="Max Steps")
                    fmax_rlx = gr.Slider(value=0.05, minimum=0.001, maximum=0.5, step=0.001, label="Fmax (eV/Å)")
                    with gr.Row():
                        charge_rlx = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
                        spin_rlx = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin")
                    relax_cell = gr.Checkbox(label="Relax Unit Cell", value=False)
                    run_rlx_btn = gr.Button("Run Optimization", variant="primary")

                with gr.Column(variant="panel", min_width=520):
                    rlx_status = gr.Textbox(label="Status", interactive=False)
                    if HAVE_MOL3D:
                        rlx_viewer = Molecule3D(label="Final Structure")
                        rlx_html = gr.HTML(visible=False)
                    else:
                        rlx_viewer = gr.Textbox(visible=False)
                        rlx_html = gr.HTML()
                    rlx_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                    rlx_log = gr.Code(label="Log", language="text", interactive=False, lines=15, max_lines=25)
                    rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=20, max_lines=30)
                    rlx_explain = gr.Markdown()

            run_rlx_btn.click(
                relax_wrapper,
                inputs=[xyz_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
                outputs=[rlx_status, rlx_viewer, rlx_html, rlx_traj, rlx_log, rlx_script, rlx_explain],
            )

print("Starting OrbMol model loading…")
_load_orbmol_calc()

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)