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

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

# Intentar importar Molecule3D para vista 3D nativa
try:
    from gradio_molecule3d import Molecule3D
    HAVE_MOL3D = True
except Exception:
    HAVE_MOL3D = False

from orb_models.forcefield import pretrained
from orb_models.forcefield.calculator import ORBCalculator


# =========================
#   OrbMol global model
# =========================
model_calc = None

def load_orbmol_model():
    global model_calc
    if model_calc is None:
        try:
            print("Loading OrbMol model...")
            orbff = pretrained.orb_v3_conservative_inf_omat(
                device="cpu",
                precision="float32-high"
            )
            model_calc = ORBCalculator(orbff, device="cpu")
            print("OrbMol model loaded successfully")
        except Exception as e:
            print(f"Error loading model: {e}")
            model_calc = None
    return model_calc


# =========================
#   Single-point (SPE)
# =========================
def predict_molecule(xyz_content, charge=0, spin_multiplicity=1):
    try:
        calc = load_orbmol_model()
        if calc is None:
            return "Error: Could not load OrbMol model", ""

        if not xyz_content.strip():
            return "Error: Please enter XYZ coordinates", ""

        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 = []
        lines.append(f"Total Energy: {energy:.6f} eV\n")
        lines.append("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.append(f"\nMax Force: {max_force:.4f} eV/Å")

        os.unlink(xyz_file)
        return "\n".join(lines), "Calculation completed with OrbMol"

    except Exception as e:
        return f"Error during calculation: {str(e)}", "Error"


# =========================
#   Trajectory → HTML 3D fallback
# =========================
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


# =========================
#   MD with OrbMol
# =========================
def run_md(xyz_content, charge=0, spin_multiplicity=1,
           steps=100, temperature=300, timestep=1.0):
    try:
        calc = load_orbmol_model()
        if calc is None:
            return "Error: Could not load OrbMol model", ""

        if not xyz_content.strip():
            return "Error: Please enter XYZ coordinates", ""

        # Leer estructura
        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

        # Pre-relajación ligera
        opt = LBFGS(atoms, logfile=None)
        opt.run(fmax=0.05, steps=20)

        MaxwellBoltzmannDistribution(atoms, temperature_K=2*float(temperature))

        dyn = VelocityVerlet(atoms, timestep=float(timestep) * units.fs)

        tf = tempfile.NamedTemporaryFile(suffix=".traj", delete=False)
        tf.close()
        traj = Trajectory(tf.name, "w", atoms)
        dyn.attach(traj.write, interval=1)
        dyn.run(int(steps))

        if HAVE_MOL3D:
            # Mostrar último frame en Molecule3D
            last = traj[-1]
            mol_xyz = f"{len(last)}\nFinal frame\n"
            for s, (x, y, z) in zip(last.get_chemical_symbols(), last.get_positions()):
                mol_xyz += f"{s} {x:.6f} {y:.6f} {z:.6f}\n"
            view = Molecule3D(value=mol_xyz, label="Final Frame (XYZ)")
        else:
            view = traj_to_html(tf.name)

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

        return f"MD completed: {int(steps)} steps at {int(temperature)} K", view

    except Exception as e:
        return f"Error during MD simulation: {str(e)}", ""


# =========================
#   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],
]


# =========================
#   Gradio UI
# =========================
with gr.Blocks(theme=gr.themes.Ocean(), title="OrbMol Demo") as demo:
    with gr.Tabs():
        # -------- Tab 1: Single Point --------
        with gr.Tab("Single Point Energy"):
            with gr.Row():
                with gr.Column(scale=2):
                    with gr.Column(variant="panel"):
                        gr.Markdown("# OrbMol Demo - Quantum-Accurate Predictions")
                        gr.Markdown("OrbMol is a neural network potential trained on the OMol25 dataset.")

                        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")

                        predict_btn = gr.Button("Run OrbMol Prediction", variant="primary")

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

            gr.Examples(examples=examples, inputs=[xyz_input, charge_input, spin_input])

            predict_btn.click(
                predict_molecule,
                inputs=[xyz_input, charge_input, spin_input],
                outputs=[results_output, status_output],
            )

            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\n* Supports charge and spin")
                with gr.Accordion("Benchmarks", open=False):
                    gr.Markdown("* <1 kcal/mol error on Wiggle150\n* Accurate protein–ligand binding energies\n* Stable MD on biomolecules >20k atoms")
                with gr.Accordion("Disclaimers", open=False):
                    gr.Markdown("* Validate results for your use case\n* Training level of theory may limit accuracy")

        # -------- Tab 2: MD --------
        with gr.Tab("Molecular Dynamics"):
            with gr.Row():
                with gr.Column(scale=2):
                    xyz_input_md = gr.Textbox(label="XYZ Coordinates", lines=12)
                    charge_input_md = gr.Slider(value=0, minimum=-10, maximum=10, step=1, label="Charge")
                    spin_input_md = gr.Slider(value=1, minimum=1, maximum=11, step=1, label="Spin Multiplicity")
                    steps_input = gr.Slider(value=100, minimum=10, maximum=1000, step=10, label="Steps")
                    temp_input = gr.Slider(value=300, minimum=10, maximum=1000, step=10, label="Temperature (K)")
                    timestep_input = gr.Slider(value=1.0, minimum=0.1, maximum=5.0, step=0.1, label="Timestep (fs)")
                    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", lines=2, interactive=False)
                    md_view = gr.HTML() if not HAVE_MOL3D else Molecule3D(label="Trajectory Viewer")

            run_md_btn.click(
                run_md,
                inputs=[xyz_input_md, charge_input_md, spin_input_md, steps_input, temp_input, timestep_input],
                outputs=[md_status, md_view],
            )


print("Starting OrbMol model loading...")
load_orbmol_model()

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