#!/usr/bin/env python3
"""ML potential {{ node_type }} generated by CatGo workflow engine."""

# --- metadata preamble (C1 reproducibility plumbing) ---
import hashlib as __hashlib
import json as __json
import os as __os
import socket as __socket
import time as __time
from datetime import datetime as __datetime, timezone as __timezone

__t_start = __time.time()
__mace_metadata = {
    "timestamp": __datetime.now(__timezone.utc).isoformat(),
    "host": __socket.gethostname(),
    "pid": __os.getpid(),
}
try:
    import torch as __torch
    __mace_metadata["torch_version"] = __torch.__version__
    __mace_metadata["cuda_available"] = bool(__torch.cuda.is_available())
    if __torch.cuda.is_available():
        try:
            __mace_metadata["gpu_name"] = __torch.cuda.get_device_name(0)
        except Exception:
            __mace_metadata["gpu_name"] = None
except Exception:
    __mace_metadata["torch_version"] = None
try:
    import mace as __mace_pkg
    __mace_metadata["mace_torch_version"] = getattr(__mace_pkg, "__version__", None)
except Exception:
    __mace_metadata["mace_torch_version"] = None

from ase.io import read, write

atoms = read("POSCAR", format="vasp")

{% set model = params.get("model", "mace-mp-0") %}
{% set device = params.get("device", "auto") %}
{% set model_path = params.get("model_path", "") %}
{% if device == "cpu" %}
__mace_device = "cpu"
{% elif device == "cuda" or device == "gpu" %}
__mace_device = "cuda"
{% else %}
__mace_device = "cuda" if __import__("torch").cuda.is_available() else "cpu"
{% endif %}
{% if model == "mace-mp-0" %}
{% if model_path %}
from mace.calculators import MACECalculator
__mace_model_path = {{ model_path|tojson }}
__mace_model_name = __mace_model_path
calc = MACECalculator(model_paths=[__mace_model_path], device=__mace_device, default_dtype="float64")
{% else %}
from mace.calculators import mace_mp
__mace_model_name = "mace-mp-0-medium"
# Default foundation model: leave path empty. mace_mp does not expose the
# cached checkpoint file path, and calc.models[0] is the loaded torch
# module (not a path) — stringifying it dumps ~4 KB of architecture repr
# into metadata.json. Pin a local checkpoint via model_path to get hashed.
__mace_model_path = ""
calc = mace_mp(model="medium", default_dtype="float64", device=__mace_device)
{% endif %}
{% elif model in ("chgnet", "m3gnet", "CHGNet", "M3GNet") %}
# Both CHGNet and M3GNet go through matgl to match mlp.py::_build_calculator_block.
# Previously the template used the standalone `chgnet` package for `chgnet`,
# which is a different API and can be installed without matgl (or vice-versa),
# causing imperative vs HPC results to silently diverge.
from matgl.ext.ase import M3GNetCalculator
import matgl
__mace_model_name = {{ model|tojson }}
__mace_model_path = ""
pot = matgl.load_model({{ model|tojson }})
calc = M3GNetCalculator(pot)
{% else %}
raise RuntimeError("Unknown ML model: {{ model }}")
{% endif %}

atoms.calc = calc

{% if node_type == "mlp_relax" %}
{% set optimizer = params.get("optimizer", "BFGS") %}
{% set fmax = params.get("fmax", 0.01) %}
{% set max_steps = params.get("max_steps", 500) %}
{% set relax_cell = params.get("relax_cell", False) %}
{% if relax_cell %}
from ase.filters import ExpCellFilter
target = ExpCellFilter(atoms)
{% else %}
target = atoms
{% endif %}
from ase.optimize import {{ optimizer }}
opt = {{ optimizer }}(target, trajectory="opt.traj", logfile="opt.log")
opt.run(fmax={{ fmax }}, steps={{ max_steps }})

write("CONTCAR", atoms, format="vasp")
print(f"Final energy: {atoms.get_potential_energy():.6f} eV")
{% elif node_type == "mlp_md" %}
{% set temp = params.get("temp", 300) %}
{% set steps = params.get("steps", 1000) %}
{% set timestep = params.get("timestep", 1.0) %}
from ase.md.langevin import Langevin
from ase import units

dyn = Langevin(atoms, timestep={{ timestep }}*units.fs, temperature_K={{ temp }}, friction=0.01)

def write_frame(a=atoms):
    write("trajectory.xyz", a, format="extxyz", append=True)

dyn.attach(write_frame, interval=10)
dyn.run({{ steps }})

write("CONTCAR", atoms, format="vasp")
print(f"Final energy: {atoms.get_potential_energy():.6f} eV")
{% elif node_type == "mlp_single_point" %}
import json
import numpy as np

energy = atoms.get_potential_energy()
forces = atoms.get_forces()
max_force = float(np.max(np.linalg.norm(forces, axis=1)))

result = {
    "energy_ev": energy,
    "max_force_ev_ang": max_force,
    "n_atoms": len(atoms),
}
with open("single_point.json", "w") as fp:
    json.dump(result, fp, indent=2)

write("CONTCAR", atoms, format="vasp")
print(f"Final energy: {energy:.6f} eV")
print(f"Max force: {max_force:.6f} eV/A")
{% elif node_type == "mlp_vibrations" %}
{% set delta = params.get("delta", 0.01) %}
{% set nfree = params.get("nfree", 2) %}
{% set freeze_mode = params.get("freeze_mode", "none") %}
{% set freeze_invert = params.get("freeze_invert", False) %}
{% set freeze_elements = params.get("freeze_elements", "") %}
{% set freeze_indices_str = params.get("freeze_indices", "") %}
{% set freeze_layers = params.get("freeze_layers", 0) %}
{% set freeze_z_below = params.get("freeze_z_below", 0) %}
import json
from ase.vibrations import Vibrations

# Derive list[int] of atom indices to displace from freeze_mode settings.
# Must stay in sync with _build_vibration_indices_block() in mlp.py.
{% if freeze_mode == "none" %}
indices = None  # vibrate all atoms
{% else %}
freeze_mode = {{ freeze_mode|tojson }}
freeze_invert = {{ freeze_invert|tojson }}
frozen = set()
{% if freeze_mode in ("indices", "manual") %}
_idx_str = {{ freeze_indices_str|tojson }}
for part in _idx_str.split(','):
    part = part.strip()
    if not part:
        continue
    if '-' in part:
        a, b = part.split('-', 1)
        frozen.update(range(int(a), int(b) + 1))
    else:
        frozen.add(int(part))
{% elif freeze_mode == "element" %}
_raw_elems = {{ freeze_elements|tojson }}.split(',')
_elems = set()
for _e in _raw_elems:
    _e = _e.strip()
    if _e:
        _elems.add(_e)
for i, sym in enumerate(atoms.get_chemical_symbols()):
    if sym in _elems:
        frozen.add(i)
{% elif freeze_mode == "z_below" %}
_z_thresh = float({{ freeze_z_below }})
for i, atom in enumerate(atoms):
    if atom.position[2] < _z_thresh:
        frozen.add(i)
{% elif freeze_mode == "layers" %}
_n_layers = int({{ freeze_layers }})
if _n_layers > 0:
    _z = [a.position[2] for a in atoms]
    _sorted_idx = sorted(range(len(atoms)), key=lambda i: _z[i])
    _cut = _sorted_idx[:_n_layers * max(1, len(atoms) // max(1, _n_layers))]
    frozen.update(_cut)
{% endif %}
_all = set(range(len(atoms)))
if freeze_invert:
    indices = sorted(frozen)
else:
    indices = sorted(_all - frozen)
if not indices:
    indices = None  # fall back to all atoms if filter produced empty set
{% endif %}

vib = Vibrations(atoms, indices=indices, delta={{ delta }}, nfree={{ nfree }}, name="vib")
vib.clean()
vib.run()

raw_freqs = vib.get_frequencies()
freq_list = []
for f in raw_freqs:
    if abs(f.imag) > abs(f.real):
        freq_list.append(-abs(f.imag))
    else:
        freq_list.append(float(f.real))

zpe = float(vib.get_zero_point_energy())

# C2: validate TS — exactly one imaginary mode with |freq| >= 20 cm^-1.
imag_modes = [float(f) for f in freq_list if f < 0]
imag_modes.sort()
dominant_imag = imag_modes[0] if imag_modes else None
nontrivial_imag = [f for f in imag_modes if abs(f) >= 20.0]
is_valid_ts = len(nontrivial_imag) == 1

result = {
    "frequencies_cm": freq_list,
    "zpe_ev": zpe,
    "n_atoms_displaced": len(indices) if indices is not None else len(atoms),
    "n_frequencies": len(freq_list),
    "imag_modes_cm": imag_modes,
    "dominant_imag_freq_cm": dominant_imag,
    "is_valid_ts": is_valid_ts,
    "n_nontrivial_imag": len(nontrivial_imag),
}
with open("frequencies.json", "w") as fp:
    json.dump(result, fp, indent=2)

vib.summary()
print(f"ZPE: {zpe:.6f} eV")

write("CONTCAR", atoms, format="vasp")
print(f"Final energy: {atoms.get_potential_energy():.6f} eV")
{% elif node_type == "mlp_neb" %}
{% set nimages = params.get("nimages", 8) %}
{% set fmax = params.get("fmax", 0.05) %}
{% set max_steps = params.get("max_steps", 500) %}
{% set climb = params.get("climb", True) %}
{% set neb_optimizer = params.get("mlp_optimizer") or params.get("optimizer", "FIRE") %}
import json
import numpy as np
try:
    from ase.mep import DyNEB
except ImportError:
    from ase.neb import DyNEB
from ase.optimize import {{ neb_optimizer }}

# NEB uses two input files written by the engine (POSCAR_initial + POSCAR_final).
# The POSCAR read at the top of the script is ignored for this branch.
initial = read("POSCAR_initial", format="vasp")
final = read("POSCAR_final", format="vasp")

# Each NEB image needs its own calculator instance. make_calc() rebinds the
# three metadata globals (already defined by the top-level calculator block)
# so the footer records model/device/path even when multiple images exist.
def make_calc():
    global __mace_device, __mace_model_name, __mace_model_path
    {% if model == "mace-mp-0" %}
    {% if model_path %}
    from mace.calculators import MACECalculator
    __mace_model_path = {{ model_path|tojson }}
    __mace_model_name = __mace_model_path
    return MACECalculator(model_paths=[__mace_model_path], device=__mace_device, default_dtype="float64")
    {% else %}
    from mace.calculators import mace_mp
    __mace_model_name = "mace-mp-0-medium"
    __mace_model_path = ""
    return mace_mp(model="medium", default_dtype="float64", device=__mace_device)
    {% endif %}
    {% elif model == "chgnet" %}
    from chgnet.model.dynamics import CHGNetCalculator
    __mace_model_name = "chgnet"
    __mace_model_path = ""
    return CHGNetCalculator()
    {% elif model == "m3gnet" %}
    from matgl.ext.ase import M3GNetCalculator
    import matgl as _matgl
    __mace_model_name = "M3GNet-MP-2021.2.8-PES"
    __mace_model_path = ""
    return M3GNetCalculator(_matgl.load_model("M3GNet-MP-2021.2.8-PES"))
    {% else %}
    raise RuntimeError("Unknown ML model: {{ model }}")
    {% endif %}

initial.calc = make_calc()
final.calc = make_calc()

images = [initial]
for _ in range({{ nimages }}):
    img = initial.copy()
    img.calc = make_calc()
    images.append(img)
images.append(final)

neb = DyNEB(images, climb={{ climb }}, fmax={{ fmax }})
neb.interpolate("idpp", mic=True)

opt = {{ neb_optimizer }}(neb, trajectory="neb.traj", logfile="neb.log")
opt.run(fmax={{ fmax }}, steps={{ max_steps }})
converged = opt.nsteps < {{ max_steps }}

energies = [float(img.get_potential_energy()) for img in images]
e_ref = energies[0]
ts_idx = int(np.argmax(energies))
barrier_ev = energies[ts_idx] - e_ref
barrier_kcal = barrier_ev * 23.0609

path_images = []
for i, e in enumerate(energies):
    de = (e - e_ref) * 23.0609
    entry = {"image": str(i), "de_kcal_mol": round(de, 4)}
    if i == ts_idx:
        entry["is_ts"] = True
        entry["image"] = "TS"
    path_images.append(entry)

ts_atoms = images[ts_idx]
write("CONTCAR", ts_atoms, format="vasp")
write("ts_structure.xyz", ts_atoms, format="extxyz")

result = {
    "neb_converged": bool(converged),
    "activation_barrier_kcal_mol": round(barrier_kcal, 4),
    "path_summary": {"images": path_images},
    "energies_ev": [round(e, 6) for e in energies],
    "energy": round(energies[ts_idx], 6),
}
with open("neb_results.json", "w") as fp:
    json.dump(result, fp, indent=2)

print(f"Final energy: {energies[ts_idx]:.6f} eV")
print(f"Barrier: {barrier_kcal:.2f} kcal/mol")
print(f"NEB converged: {converged}")
{% else %}
raise RuntimeError("Unknown MLP node type: {{ node_type }}")
{% endif %}

# --- metadata footer (C1 reproducibility plumbing) ---
try:
    __mace_metadata["mace_model"] = __mace_model_name
    __mace_metadata["device"] = __mace_device
    __mace_metadata["model_path"] = __mace_model_path
    if __mace_model_path and __os.path.isfile(__mace_model_path):
        __h = __hashlib.sha256()
        with open(__mace_model_path, "rb") as __mf:
            for __chunk in iter(lambda: __mf.read(65536), b""):
                __h.update(__chunk)
        __mace_metadata["model_sha256"] = __h.hexdigest()
    else:
        __mace_metadata["model_sha256"] = None
except NameError:
    pass
__mace_metadata["wall_time_s"] = round(__time.time() - __t_start, 3)
try:
    with open("metadata.json", "w") as __mf:
        __json.dump(__mace_metadata, __mf, indent=2, default=str)
except Exception as __exc:
    print(f"[metadata] failed to write metadata.json: {__exc}")
