app.py

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

from ase.io import read
from ase.io.trajectory import Trajectory

# ==== Visualizador 3D ====
try:
    from gradio_molecule3d import Molecule3D
    HAVE_MOL3D = True
except Exception:
    HAVE_MOL3D = False

def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
    """Fallback de visualización con 3Dmol.js si no hay Molecule3D."""
    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 SPE directo (rápido) ====
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):
    """Single Point Energy/Forces con OrbMol (input en Textbox XYZ)."""
    try:
        calc = _load_orbmol_calc()
        if not xyz_content or 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"


# ==== Simulación (estilo UMA) vía helpers locales ====
from simulation_scripts_orbmol import (
    run_md_simulation,
    run_relaxation_simulation,
    last_frame_xyz_from_traj,
)

def md_wrapper(xyz_content, charge, spin, steps, tempK, timestep_fs, ensemble):
    """Conecta Gradio con run_md_simulation y adapta outputs (viewer/HTML/file/log/script)."""
    try:
        traj_path, log_text, script_text, explanation = run_md_simulation(
            xyz_content,            # acepta string XYZ o ruta
            int(steps),
            20,                     # prerelax por defecto, como 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})"

        if HAVE_MOL3D:
            viewer_value = last_frame_xyz_from_traj(traj_path)  # string XYZ como value
            html_value = None
        else:
            viewer_value = None
            html_value = traj_to_html(traj_path)

        return (
            status,         # md_status
            viewer_value,   # md_viewer (Molecule3D value) o None
            html_value,     # md_html (fallback) o None
            traj_path,      # md_traj file
            log_text,       # md_log (Textbox)
            script_text,    # md_script (Code py)
            explanation,    # md_explain
        )
    except Exception as e:
        return (f"Error: {e}", None, None, None, "", "", "")

def relax_wrapper(xyz_content, steps, fmax, charge, spin, relax_cell):
    """Conecta Gradio con run_relaxation_simulation y adapta outputs."""
    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,         # rlx_status
            viewer_value,   # rlx_viewer (Molecule3D value) o None
            html_value,     # rlx_html (fallback) o None
            traj_path,      # rlx_traj file
            log_text,       # rlx_log (Textbox)
            script_text,    # rlx_script (Code py)
            explanation,    # rlx_explain
        )
    except Exception as e:
        return (f"Error: {e}", None, None, None, "", "", "")


# ==== Ejemplos ====
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():
        # --- Tab 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 (trained on **OMol25**, "
                        "ωB97M-V/def2-TZVPD). 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(
                        "* Strong results on **GMTKN55** y **Wiggle150**\n"
                        "* Accurate **protein–ligand** energies (PLA15)\n"
                        "* Stable MD en biomoléculas grandes"
                    )
                with gr.Accordion("Disclaimers", open=False):
                    gr.Markdown(
                        "* Verifica resultados para tu caso\n"
                        "* Considera el **nivel de teoría** de entrenamiento"
                    )

        # --- Tab 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=2000, 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 layout consistente
                    else:
                        md_viewer = gr.Textbox(visible=False)  # placeholder
                        md_html = gr.HTML()

                    md_traj = gr.File(label="Trajectory (.traj)", interactive=False)
                    md_log = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)  # <- FIX Code->Textbox
                    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],
            )

        # --- Tab 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=2000, 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.Textbox(label="Log", interactive=False, lines=15, max_lines=25)  # <- FIX Code->Textbox
                    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 = _load_orbmol_calc()

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