app.py

import os
import re
import numpy as np
import gradio as gr
from ase.io import read
from ase.io.trajectory import Trajectory

# ========= 3Dmol.js (para ver trayectorias MD/Relax) =========
THREE_D_MOL_SOURCES = [
    "https://3dmol.org/build/3Dmol-min.js",
    "https://cdn.jsdelivr.net/npm/3dmol/build/3Dmol-min.js",
    "https://unpkg.com/3dmol/build/3Dmol-min.js",
]

def _loader_js():
    srcs = "[" + ",".join([f"'{u}'" for u in THREE_D_MOL_SOURCES]) + "]"
    return f"""
function _ensure3Dmol(cb){{
  if(typeof window.$3Dmol!=='undefined') return cb();
  const srcs={srcs}; let i=0;
  function tryNext(){{
    if(i>=srcs.length) return;
    const s=document.createElement('script'); s.src=srcs[i++]; s.onload=cb; s.onerror=tryNext; document.head.appendChild(s);
  }}
  tryNext();
}}
"""

def _atoms_to_xyz_block(atoms):
    syms = atoms.get_chemical_symbols()
    pos  = atoms.get_positions()
    out = [str(len(syms)), "frame"]
    for s,(x,y,z) in zip(syms,pos):
        out.append(f"{s} {x:.6f} {y:.6f} {z:.6f}")
    return "\n".join(out)

def traj_to_html(traj_path, width=520, height=520, interval_ms=200):
    if not traj_path or not os.path.exists(traj_path):
        return "<div style='color:#b00;padding:20px;'>No trajectory file found</div>"
    try:
        traj = Trajectory(traj_path)
        if len(traj)==0:
            return "<div style='color:#555;padding:20px;'>Empty trajectory</div>"
    except Exception as e:
        return f"<div style='color:#b00;padding:20px;'>Error reading trajectory: {e}</div>"

    frames = [_atoms_to_xyz_block(at) for at in traj]
    frames_json = str(frames).replace("'", '"')
    viewer_id = f"viewer_{abs(hash(traj_path))%100000}"
    loader = _loader_js()
    return f"""
<div style="margin-bottom:10px;padding:10px;background:#f5f5f5;border-radius:5px;">
  <strong>🧬 3D Molecular Viewer</strong> — {len(frames)} frames
</div>
<div id="{viewer_id}" style="width:{width}px;height:{height}px;border:2px solid #ddd;border-radius:8px;background:#fafafa;"></div>
<script>
{loader}
_ensure3Dmol(function(){{
  var el=document.getElementById("{viewer_id}"); if(!el||typeof $3Dmol==='undefined') return;
  var v=$3Dmol.createViewer(el, {{backgroundColor:'white'}});
  var frames={frames_json}; var i=0;
  function draw(k){{ v.clear(); v.addModel(frames[k], "xyz"); v.setStyle({{}}, {{stick:{{}}, sphere:{{}}}}); v.zoomTo(); v.render(); }}
  draw(0);
  if(frames.length>1) setInterval(function(){{ i=(i+1)%frames.length; draw(i); }}, {interval_ms});
}});
</script>
"""

# ================= OrbMol (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(structure_file, charge=0, spin_multiplicity=1):
    """
    Single Point Energy + fuerzas (OrbMol). Solo se ejecuta al pulsar el botón.
    """
    try:
        calc = _load_orbmol_calc()
        if not structure_file:
            return "Error: Please upload a structure file", "Error"

        file_path = structure_file  # gr.File(type='filepath') -> str
        if not os.path.exists(file_path):
            return f"Error: File not found: {file_path}", "Error"
        if os.path.getsize(file_path) == 0:
            return f"Error: Empty file: {file_path}", "Error"

        atoms = read(file_path)
        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, fc in enumerate(forces):
            lines.append(f"Atom {i+1}: [{fc[0]:.4f}, {fc[1]:.4f}, {fc[2]:.4f}] eV/Å")
        max_force = float(np.max(np.linalg.norm(forces, axis=1)))
        lines += ["", f"Max Force: {max_force:.4f} eV/Å"]

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

# ================= Simulaciones (tus helpers) =================
from simulation_scripts_orbmol import (
    run_md_simulation,
    run_relaxation_simulation,
)

# ========== Wrappers MD / Relax (con firma correcta) ==========
def md_wrapper(structure_file, charge, spin, steps, tempK, timestep_fs, ensemble):
    """
    Llama a tu run_md_simulation con el ORDEN correcto usando keywords:
    (structure_file, num_steps, num_prerelax_steps, md_timestep, temperature_k, md_ensemble,
     task_name='OMol', total_charge=..., spin_multiplicity=...)
    """
    try:
        if not structure_file:
            return ("Error: Please upload a structure file", None, "", "", "", "", None)
        file_path = structure_file
        if not os.path.exists(file_path):
            return ("Error: File not found: " + str(file_path), None, "", "", "", "", None)
        if os.path.getsize(file_path) == 0:
            return ("Error: Empty file: " + str(file_path), None, "", "", "", "", None)

        traj_path, log_text, script_text, explanation = run_md_simulation(
            file_path,
            int(steps),
            20,
            float(timestep_fs),
            float(tempK),
            "NVT" if ensemble == "NVT" else "NVE",
            total_charge=int(charge),
            spin_multiplicity=int(spin),
        )
        status = f"MD completed: {int(steps)} steps at {int(tempK)} K ({ensemble})"
        html_value = traj_to_html(traj_path)
        return (status, traj_path, log_text, script_text, explanation, html_value, None)
    except Exception as e:
        return (f"Error: {e}", None, "", "", "", "", None)

def relax_wrapper(structure_file, steps, fmax, charge, spin, relax_cell):
    """
    Firma correcta usando keywords:
    (structure_file, num_steps, fmax, task_name='OMol', total_charge=..., spin_multiplicity=..., relax_unit_cell=...)
    """
    try:
        if not structure_file:
            return ("Error: Please upload a structure file", None, "", "", "", "", None)
        file_path = structure_file
        if not os.path.exists(file_path):
            return ("Error: File not found: " + str(file_path), None, "", "", "", "", None)
        if os.path.getsize(file_path) == 0:
            return ("Error: Empty file: " + str(file_path), None, "", "", "", "", None)

        traj_path, log_text, script_text, explanation = run_relaxation_simulation(
            file_path,
            int(steps),
            float(fmax),
            total_charge=int(charge),
            spin_multiplicity=int(spin),
            relax_unit_cell=bool(relax_cell),
        )
        status = f"Relaxation finished (≤ {int(steps)} steps, fmax={float(fmax)} eV/Å)"
        html_value = traj_to_html(traj_path)
        return (status, traj_path, log_text, script_text, explanation, html_value, None)
    except Exception as e:
        return (f"Error: {e}", None, "", "", "", "", None)

# ===================== UI (solo calcula al pulsar botón) =====================
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 — Single Point Energy")
                    xyz_input = gr.File(
                        label="Upload Structure File (.xyz/.pdb/.cif/.traj/.mol/.sdf)",
                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                        file_count="single",
                        type="filepath",
                    )
                    with gr.Row():
                        charge_input = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                        spin_input   = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
                    run_spe = gr.Button("Run OrbMol Prediction", variant="primary")
                with gr.Column(variant="panel", min_width=520):
                    spe_out    = gr.Textbox(label="Energy & Forces", lines=18, interactive=False)
                    spe_status = gr.Textbox(label="Status", interactive=False, max_lines=1)
            run_spe.click(predict_molecule, [xyz_input, charge_input, spin_input], [spe_out, spe_status])

        # ===== MD =====
        with gr.Tab("Molecular Dynamics"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("## Molecular Dynamics Simulation")
                    xyz_md = gr.File(
                        label="Upload Structure File (.xyz/.pdb/.cif/.traj/.mol/.sdf)",
                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                        file_count="single",
                        type="filepath",
                    )
                    with gr.Row():
                        charge_md = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                        spin_md   = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin Multiplicity")
                    with gr.Row():
                        steps_md  = gr.Slider(minimum=10, maximum=2000, value=100, step=10, label="Steps")
                        temp_md   = gr.Slider(minimum=10, maximum=1500, value=300, step=10, label="Temperature (K)")
                    with gr.Row():
                        timestep_md = gr.Slider(minimum=0.1, maximum=5.0, value=1.0, step=0.1, label="Timestep (fs)")
                        ensemble_md = gr.Radio(["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)
                    md_traj   = gr.File(label="Trajectory (.traj)", interactive=False)
                    md_html   = gr.HTML(label="Trajectory Viewer", sanitize=False)
                    md_log    = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)
                    md_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=18, max_lines=28)
                    md_explain= gr.Markdown()
                    md_plot   = gr.Plot(label="(optional)")
            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_traj, md_log, md_script, md_explain, md_html, md_plot],
            )

        # ===== Relax =====
        with gr.Tab("Relaxation / Optimization"):
            with gr.Row():
                with gr.Column(scale=2):
                    gr.Markdown("## Structure Relaxation/Optimization")
                    xyz_rlx = gr.File(
                        label="Upload Structure File (.xyz/.pdb/.cif/.traj/.mol/.sdf)",
                        file_types=[".xyz", ".pdb", ".cif", ".traj", ".mol", ".sdf"],
                        file_count="single",
                        type="filepath",
                    )
                    steps_rlx = gr.Slider(minimum=1, maximum=2000, value=300, step=1, label="Max Steps")
                    fmax_rlx  = gr.Slider(minimum=0.001, maximum=0.5, value=0.05, step=0.001, label="Fmax (eV/Å)")
                    with gr.Row():
                        charge_rlx = gr.Slider(minimum=-10, maximum=10, value=0, step=1, label="Charge")
                        spin_rlx   = gr.Slider(minimum=1, maximum=11, value=1, step=1, label="Spin")
                    relax_cell = gr.Checkbox(False, label="Relax Unit Cell")
                    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)
                    rlx_traj   = gr.File(label="Trajectory (.traj)", interactive=False)
                    rlx_html   = gr.HTML(label="Final Structure / Trajectory", sanitize=False)
                    rlx_log    = gr.Textbox(label="Log", interactive=False, lines=15, max_lines=25)
                    rlx_script = gr.Code(label="Reproduction Script", language="python", interactive=False, lines=18, max_lines=28)
                    rlx_explain= gr.Markdown()
                    rlx_plot   = gr.Plot(label="(optional)")
            run_rlx_btn.click(
                relax_wrapper,
                inputs=[xyz_rlx, steps_rlx, fmax_rlx, charge_rlx, spin_rlx, relax_cell],
                outputs=[rlx_status, rlx_traj, rlx_log, rlx_script, rlx_explain, rlx_html, rlx_plot],
            )

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)