"""AceFF-2: interactive machine-learning potential for small molecules.

Enter a molecule (XYZ block, PDB, or SMILES), pick a net charge, and run a
single-point calculation (energy + per-atom forces) and optionally an ASE LBFGS
geometry relaxation with the AceFF-2 TensorNet-2 potential
(https://huggingface.co/papers/2601.00581, Acellera, Apache-2.0).
"""

import json
import os
import time
import traceback

# Keep warp JIT kernels out of the home dir
os.environ.setdefault("WARP_CACHE_DIR", "/tmp/warp_cache")

import spaces  # must be first (before any CUDA-touching import)
import gradio as gr

MODELS = {
    "AceFF-2.0 (H B C N O F Si P S Cl Br I)": ("Acellera/AceFF-2.0", "aceff_v2.0.ckpt"),
    "AceFF-2-OMOL (adds Na, K; OMol25 data)": ("Acellera/AceFF-2-OMOL", "aceff_v2_omol.ckpt"),
}
SUPPORTED_Z = {1, 5, 6, 7, 8, 9, 14, 15, 16, 17, 35, 53}  # AceFF-2 elements
SUPPORTED_Z_OMOL = SUPPORTED_Z | {11, 19}
EV_TO_KCAL_MOL = 23.0605
MAX_ATOMS = 60
DEFAULT_MODEL = list(MODELS)[0]

_state = {}  # per-GPU-worker cache: models + ASE calculators


def _sym_z(symbol):
    from ase.data import atomic_numbers

    return int(atomic_numbers.get(symbol.title(), 0))


def _get_model(model_key):
    """Load (and cache) an AceFF checkpoint on the GPU worker via load_model."""
    if "model" in _state.get(model_key, {}):
        return _state[model_key]["model"]
    import torch
    from huggingface_hub import hf_hub_download
    from torchmdnet.models.model import load_model

    repo_id, filename = MODELS[model_key]
    path = hf_hub_download(repo_id=repo_id, filename=filename)
    model = load_model(path, derivative=False, remove_ref_energy=True)
    model.eval()
    for p in model.parameters():
        p.requires_grad = False
    model.to("cuda")
    _state.setdefault(model_key, {})["model"] = model
    return model


def _get_calculator(model_key):
    """Load (and cache) the TorchMD-Net ASE calculator on the GPU worker."""
    calcs = _state.setdefault(model_key, {})
    if "calc" in calcs:
        return calcs["calc"]
    from huggingface_hub import hf_hub_download
    from torchmdnet.calculators import TMDNETCalculator

    repo_id, filename = MODELS[model_key]
    calc = TMDNETCalculator(
        hf_hub_download(repo_id=repo_id, filename=filename), device="cuda"
    )
    calcs["calc"] = calc
    return calc


def _energy_forces(model, z, pos, q):
    """Single-point energy + forces: fresh graph, autograd.grad each call.

    The checkpoint is loaded with derivative=False so the model returns only
    the energy; forces come from torch.autograd.grad, which is repeatable
    across calls (a derivative=True model frees its graph on .backward()).
    """
    import torch

    pos = pos.detach().clone().requires_grad_(True)
    energy, _ = model(z, pos, q=q)
    grad = torch.autograd.grad(energy, pos)[0]
    return float(energy.item()), (-grad).detach().cpu().numpy()


def _parse_molecule(mol_text):
    """Parse XYZ / PDB / SMILES input into (symbols, positions Nx3 ndarray)."""
    import numpy as np
    from ase.io import read as ase_read

    raw = mol_text or ""
    if not raw.strip():
        raise ValueError(
            "Please provide a molecule: paste an XYZ block or PDB, or type a SMILES string."
        )
    nonempty = [l for l in raw.splitlines() if l.strip()]
    all_lines = raw.splitlines()

    is_pdb = any(l.startswith(("ATOM", "HETATM")) for l in nonempty)
    n_first = None
    try:
        n_first = int(all_lines[0].strip().split()[0])
    except (ValueError, IndexError):
        pass
    counts_ok = False
    if n_first is not None and n_first > 0:
        coord_lines = [
            l for l in all_lines[2:] if l.strip() and len(l.split()) >= 4
        ]
        counts_ok = len(coord_lines) >= n_first and len(all_lines) >= n_first + 2
    is_xyz = counts_ok

    if not is_xyz and not is_pdb:
        if len(nonempty) == 1 and len(nonempty[0]) < 120:
            return _smiles_to_atoms(nonempty[0])
        raise ValueError(
            "Input must be an XYZ block (atom count on line 1, comment on line 2), "
            "a PDB, or a SMILES string on a single line."
        )

    import tempfile

    with tempfile.NamedTemporaryFile(
        "w", suffix=".pdb" if is_pdb else ".xyz", delete=False
    ) as f:
        f.write(raw + "\n")
        path = f.name
    try:
        atoms = ase_read(path)
    finally:
        os.unlink(path)
    symbols = list(atoms.get_chemical_symbols())
    pos = np.array(atoms.get_positions(), dtype=np.float32)
    if len(symbols) > 1 and float(pos.std()) < 1e-6:
        raise ValueError(
            "All atoms sit at the same position — is this a valid 3D structure?"
        )
    return symbols, pos


def _smiles_to_atoms(smiles):
    """Embed a 3D geometry from SMILES with RDKit (ETKDG)."""
    import numpy as np
    from rdkit import Chem
    from rdkit.Chem import AllChem

    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"Could not parse input as SMILES: {smiles!r}")
    mol = Chem.AddHs(mol)
    status = AllChem.EmbedMultipleConfs(
        mol, numConfs=1, useRandomCoords=True, randomSeed=42
    )
    confs = list(status)
    if not confs:
        raise ValueError(f"RDKit could not generate a 3D geometry for: {smiles}")
    conf = mol.GetConformer(confs[0])
    pos = np.array(conf.GetPositions(), dtype=np.float32)
    symbols = [a.GetSymbol() for a in mol.GetAtoms()]
    return symbols, pos


def _xyz_block(symbols, pos, comment=""):
    import numpy as np

    lines = [str(len(symbols)), comment]
    for s, (x, y, z) in zip(symbols, np.asarray(pos)):
        lines.append(f"{s:<2s} {x:>12.6f} {y:>12.6f} {z:>12.6f}")
    return "\n".join(lines)


def _viewer_html(title, xyz_in, xyz_out, arrows, arrow_scale_label):
    """Build the 3Dmol.js viewer as an escaped srcdoc iframe.

    Gradio injects gr.HTML values via innerHTML, and browsers refuse to run
    <script> tags inserted that way (gradio warns about it at runtime), so a
    plain inline script would leave a blank box. An <iframe srcdoc="...">
    is a real document parsed by the browser: its scripts execute normally,
    and the iframe reloads whenever the value updates, so the viewer re-renders
    after every calculation.
    """
    import html as html_mod

    has_out = bool(xyz_out)
    doc = f"""<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<script src="https://3Dmol.org/build/3Dmol-min.js"></script>
<style>
  html, body {{ margin: 0; padding: 0; height: 100%; background: #17191c; overflow: hidden;
               font-family: system-ui, -apple-system, sans-serif; }}
  #label {{ position: absolute; top: 8px; left: 10px; right: 10px; color: #e8eaed;
            font-size: 13px; z-index: 10; pointer-events: none; }}
  #label b {{ font-size: 14px; }}
  #label span {{ opacity: 0.7; }}
  #v {{ position: absolute; inset: 0; }}
</style>
</head>
<body>
<div id="label"><b>{title}</b>
  <span> &nbsp;· rotate: drag &nbsp;· zoom: scroll &nbsp;· translate: right-drag &nbsp;· red arrows = per-atom force vectors (&times;{arrow_scale_label})</span></div>
<div id="v"></div>
<script>
var XYZ_IN = {json.dumps(xyz_in)};
var XYZ_OUT = {json.dumps(xyz_out) if has_out else "null"};
var ARROWS = {json.dumps(arrows)};
function boot(attempts) {{
  var el = document.getElementById('v');
  if (typeof $3Dmol === 'undefined') {{
    if (attempts > 0) setTimeout(function () {{ boot(attempts - 1); }}, 400);
    return;
  }}
  var viewer = $3Dmol.createViewer({{el: el, backgroundColor: '#17191c'}});
  var mIn = viewer.addModel(XYZ_IN, 'xyz');
  mIn.setStyle({{}}, {{stick: {{radius: 0.14, colorscheme: 'cyanCarbon'}},
                      sphere: {{scale: 0.2, colorscheme: 'cyanCarbon'}}}});
  if (XYZ_OUT) {{
    var mOut = viewer.addModel(XYZ_OUT, 'xyz');
    mOut.setStyle({{}}, {{stick: {{radius: 0.10, opacity: 0.35}}}});
  }}
  for (var i = 0; i < ARROWS.length; i++) {{
    var a = ARROWS[i];
    if (a.s <= 0.01) continue;
    viewer.addArrow({{start: {{x: a.p0[0], y: a.p0[1], z: a.p0[2]}},
                      end:   {{x: a.p1[0], y: a.p1[1], z: a.p1[2]}},
                      radius: 0.05, color: 'red'}});
  }}
  viewer.zoomTo();
  viewer.render();
}}
boot(25);
</script>
</body>
</html>"""
    escaped = html_mod.escape(doc, quote=True)
    return (
        f'<iframe srcdoc="{escaped}" title="3Dmol viewer" '
        f'style="width:100%;height:450px;border:0;border-radius:8px;display:block;" '
        f'allowfullscreen></iframe>'
    )


def _fmt_metrics(e_ev, de_ev, fmax_ev, relaxed):
    import numpy as np

    fmax_kcal = fmax_ev * EV_TO_KCAL_MOL
    items = [
        ("Energy", f"{e_ev:.4f} eV ({e_ev * EV_TO_KCAL_MOL:,.1f} kcal/mol)"),
    ]
    if relaxed:
        items.append(("Δ after relaxation", f"{de_ev:+.4f} eV ({de_ev * EV_TO_KCAL_MOL:+,.1f} kcal/mol)"))
    items.append(("Max |F|", f"{fmax_ev:.3f} eV/Å ({fmax_kcal:.1f} kcal/mol/Å)"))
    return items


@spaces.GPU(duration=30)
def run(
    molecule: str,
    net_charge: float = 0,
    do_relax: bool = False,
    relax_fmax: float = 0.05,
    relax_steps: int = 200,
    model_choice: str = DEFAULT_MODEL,
    show_forces: bool = True,
):
    """Run an AceFF-2 single-point calculation (energy + forces) and optional
    geometry relaxation for one small molecule.

    Args:
        molecule: XYZ block, PDB text, or a SMILES string.
        net_charge: total molecular charge (-2..+2 supported by the model).
        do_relax: also relax the geometry with ASE LBFGS at this potential.
        relax_fmax: convergence threshold on max atomic force (eV/Å).
        relax_steps: max optimizer steps.
        model_choice: which AceFF-2 checkpoint to use.
        show_forces: draw per-atom force arrows on the 3D viewer.

    Returns:
        (status markdown, metric items, viewer HTML, relaxed XYZ text)
    """
    t_start = time.perf_counter()
    try:
        import numpy as np
        import torch
        from ase import Atoms
        from ase.optimize import LBFGS

        model_key = model_choice if model_choice in MODELS else DEFAULT_MODEL
        supported = SUPPORTED_Z_OMOL if "OMOL" in model_key else SUPPORTED_Z

        # ---- 1. parse + validate input -------------------------------------
        symbols, pos = _parse_molecule(molecule)
        if len(symbols) > MAX_ATOMS:
            raise ValueError(
                f"Molecule has {len(symbols)} atoms; this demo caps at {MAX_ATOMS}."
            )
        znums = [_sym_z(s) for s in symbols]
        bad = sorted({s for s, zn in zip(symbols, znums) if zn not in supported})
        if bad:
            elems = (
                "H, B, C, N, O, F, Na, Si, P, S, Cl, K, Br, I"
                if "OMOL" in model_key
                else "H, B, C, N, O, F, Si, P, S, Cl, Br, I"
            )
            raise ValueError(
                f"Unsupported element(s) {', '.join(bad)}. {model_key} was trained on: {elems}."
            )
        if not -5 <= float(net_charge) <= 5:
            raise ValueError("Net charge must be an integer between -2 and +2.")
        if abs(float(net_charge) - round(float(net_charge))) > 1e-9:
            raise ValueError("Net charge must be an integer (e.g. 0, -1, +2).")
        q_float = float(net_charge)
        if not -2 <= q_float <= 2:
            raise ValueError(
                "AceFF-2 was trained on total molecular charges from -2 to +2 only."
            )

        # ---- 2. single point E + F (reference path: forward + backward) ----
        model = _get_model(model_key)
        z = torch.tensor(znums, dtype=torch.long, device="cuda")
        p = torch.tensor(np.asarray(pos), dtype=torch.float32, device="cuda")
        q = torch.tensor([q_float], dtype=torch.float32, device="cuda")
        e_init_ev, f_init = _energy_forces(model, z, p, q)

        xyz_in = _xyz_block(symbols, pos, "input geometry")
        xyz_out = None
        steps_used = 0
        e_final_ev = e_init_ev
        fmax_final = float(np.abs(f_init).max()) if len(f_init) else 0.0

        # ---- 3. optional LBFGS relaxation via the ASE calculator -----------
        if do_relax:
            atoms = Atoms(symbols=symbols, positions=np.asarray(pos))
            atoms.info["charge"] = q_float
            atoms.calc = _get_calculator(model_key)
            dyn = LBFGS(atoms, logfile=None)
            dyn.run(fmax=float(relax_fmax), steps=int(relax_steps))
            steps_used = dyn.get_number_of_steps()
            e_final_ev = float(atoms.get_potential_energy())
            fmax_final = float(np.abs(atoms.get_forces()).max())
            xyz_out = _xyz_block(
                symbols, atoms.get_positions(), "relaxed geometry (AceFF-2, LBFGS)"
            )

        de_ev = e_final_ev - e_init_ev
        metrics = _fmt_metrics(e_final_ev, de_ev, fmax_final, do_relax)

        # ---- 4. viewer: arrows scale longest force to ~0.55 Å -------------
        fmax_init = float(np.abs(f_init).max()) if len(f_init) else 0.0
        scale = 0.55 / fmax_init if fmax_init > 1e-9 else 0.0
        disp_pos = np.asarray(pos)
        arrows = []
        if show_forces and scale > 0:
            for (x, y, zw), (fx, fy, fz) in zip(disp_pos, f_init):
                mag = float(np.linalg.norm([fx, fy, fz]))
                arrows.append(
                    {
                        "p0": [round(float(x), 3), round(float(y), 3), round(float(zw), 3)],
                        "p1": [
                            round(float(x + fx * scale), 3),
                            round(float(y + fy * scale), 3),
                            round(float(zw + fz * scale), 3),
                        ],
                        "s": mag,
                    }
                )
        title = (
            "Input geometry (solid) → relaxed geometry (faded)"
            if do_relax
            else "Input geometry"
        )
        viewer = _viewer_html(title, xyz_in, xyz_out, arrows, f"{1.0/scale:.2f} Å/(eV/Å)" if scale else "—")

        mode = "single point + relaxation" if do_relax else "single point"
        q_str = f"{int(q_float):+d}" if q_float == int(q_float) else f"{q_float:+g}"
        status = (
            f"**Done** — {mode} · {len(symbols)} atoms · charge {q_str}"
            + (f" · {steps_used} LBFGS steps" if do_relax else "")
            + f" · {time.perf_counter() - t_start:.1f}s"
        )
        return status, metrics, viewer, xyz_out or ""
    except Exception as e:
        print("run() error:\n", traceback.format_exc(limit=6))
        raise gr.Error(f"{type(e).__name__}: {e}")


with gr.Blocks(title="AceFF-2 · ML potential for small molecules") as demo:
    gr.Markdown(
        """
# AceFF-2 — Machine Learning Potential for Small Molecules

**AceFF** ([paper](https://huggingface.co/papers/2601.00581)) is a state-of-the-art neural
interatomic potential for drug-like molecules, built on the **TensorNet v2** architecture in
[TorchMD-Net](https://github.com/torchmd/torchmd-net) (Acellera, Apache-2.0).
Give it a molecule — **XYZ**, **PDB**, or **SMILES** — and it predicts the **potential energy**
and **per-atom forces**, and can **relax the geometry** with ASE.

Supported elements: H, B, C, N, O, F, Si, P, S, Cl, Br, I (Na, K with the OMOL checkpoint) ·
net charge −2…+2 · small molecules only (≤ 60 atoms in this demo).
        """
    )
    with gr.Row(elem_id="col-container"):
        with gr.Column(variant="panel", scale=5):
            molecule = gr.Textbox(
                label="Molecule (XYZ block, PDB, or SMILES)",
                placeholder="Paste an XYZ block or PDB, or type a SMILES like CCO…",
                lines=10,
                max_lines=18,
            )
            with gr.Row():
                net_charge = gr.Number(label="Net charge (e)", value=0, precision=0, minimum=-2, maximum=2)
                do_relax = gr.Checkbox(label="Relax geometry (LBFGS)", value=False)
                show_forces = gr.Checkbox(label="Show force arrows", value=True)
            model_choice = gr.Radio(
                choices=list(MODELS),
                value=DEFAULT_MODEL,
                label="Checkpoint",
                info="AceFF-2.0: Acellera training set · AceFF-2-OMOL: public OMol25 data, adds Na/K.",
            )
            with gr.Accordion("Relaxation settings", open=False):
                relax_fmax = gr.Number(label="Convergence: max force (eV/Å)", value=0.05)
                relax_steps = gr.Number(label="Max optimizer steps", value=200, precision=0, minimum=1, maximum=1000)
            run_btn = gr.Button("Compute energy & forces", variant="primary")

        with gr.Column(scale=7):
            status = gr.Markdown("Ready. Pick an example below or paste your own molecule.")
            metrics = gr.HighlightedText(
                label="Results",
                combine_adjacent=False,
            )
            viewer = gr.HTML(
                value="<div style='height:430px;background:#17191c;border-radius:8px;display:flex;"
                "align-items:center;justify-content:center;color:#9aa0a6;font-family:var(--font);'>"
                "3D viewer appears here after a calculation (force arrows in red).</div>"
            )
            with gr.Accordion("Relaxed geometry (XYZ) — copy or download", open=False):
                xyz_out = gr.Textbox(label="XYZ", lines=8, interactive=True)

    gr.Examples(
        examples=[
            ["examples/caffeine.xyz", 0, False, 0.05, 200, DEFAULT_MODEL, True],
            ["examples/aspirin.xyz", 0, True, 0.05, 200, DEFAULT_MODEL, True],
            ["examples/acetate.xyz", -1, True, 0.05, 200, DEFAULT_MODEL, True],
            ["examples/ethanol.xyz", 0, False, 0.05, 200, DEFAULT_MODEL, True],
            ["examples/water.xyz", 0, False, 0.05, 200, DEFAULT_MODEL, True],
        ],
        inputs=[molecule, net_charge, do_relax, relax_fmax, relax_steps, model_choice, show_forces],
        outputs=[status, metrics, viewer, xyz_out],
        fn=run,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples (molecules from the AceFF usage docs; geometries embedded with RDKit)",
    )

    run_btn.click(
        run,
        inputs=[molecule, net_charge, do_relax, relax_fmax, relax_steps, model_choice, show_forces],
        outputs=[status, metrics, viewer, xyz_out],
        api_name="compute",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(),
            css="""
#col-container { max-width: 1280px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
""")