import Base
import ./geom.bend as G
import ./protein.bend as P
import ./topology.bend as T
import ./force.bend as F
import ./bonded.bend as B
import ./rng.bend as Rng

# --- Simulation engine v0: velocity Verlet NVT (Berendsen) + steepest
# descent with convergence driver, over explicit DynAtom state. ---
# Units (AMBER-consistent): lengths A, masses amu, energies kcal/mol, time
# fs, charges e, ke(Coulomb) = 332.0637. Accelerations need ACC_F = 4.184e-4
# A/fs^2 per kcal/mol/A/amu; velocities are A/fs; temperature from KE via
# kB = 0.001987 kcal/mol/K. Dynamics forces here are uniform-LJ with
# 1-2/1-3 exclusions + Coulomb + bonds (zero charges and no angle forces
# in v0; per-element reporting lives in pdbdemo).
# State is one Data list (never affine pairs), so drivers reuse it freely;
# loops are single self-recursions with computed flags passed as params.

type DynAtom is Data:
  DA{serial: U32, elem: U32, pos: G.Vec3, vel: G.Vec3, frc: G.Vec3, mass: F32}

type SimParams is Data:
  SP{eps: F32, sig2: F32, c2: F32, ke: F32, kbond: F32}

def acc_conv() -> F32:
  0.0004184

def boltzmann() -> F32:
  0.001987

def ke_coul() -> F32:
  332.0637

def dyn_of_atom(h: P.Atom) -> DynAtom:
  match h:
    case P.Atom{serial, +elem, pos}:
      DA{serial, elem, pos, G.V3{0.0, 0.0, 0.0}, G.V3{0.0, 0.0, 0.0}, F.mass(elem)}

def dyn_of_atoms(xs: List<&2, P.Atom>) -> List<&2, DynAtom>:
  match xs:
    case Nil{}:
      Nil{}
    case h <> t:
      dyn_of_atom(h) <> dyn_of_atoms(t)

def dyn_atoms(ds: List<&2, DynAtom>) -> List<&2, P.Atom>:
  match ds:
    case Nil{}:
      Nil{}
    case h <> t:
      match h:
        case DA{serial, elem, pos, vel, frc, mass}:
          P.Atom{serial, elem, pos} <> dyn_atoms(t)

# Zero charges aligned with atoms (no charge model yet; Coulomb kernels
# take external charges and are exercised by laws).
def zero_charges(xs: List<&2, P.Atom>) -> List<&2, F32>:
  match xs:
    case Nil{}:
      Nil{}
    case h <> t:
      0.0 <> zero_charges(t)

# Uniform kicks in [-0.5, 0.5) per component, seeded per serial.
def kick_vec(+s: U32) -> G.Vec3:
  G.V3{(Rng.rng01(s) - 0.5 : F32), (Rng.rng01(U32.add(s, 99991)) - 0.5 : F32), (Rng.rng01(U32.add(s, 199999)) - 0.5 : F32)}

def init_kick_atom(h: DynAtom) -> DynAtom:
  match h:
    case DA{+serial, elem, pos, vel, frc, mass}:
      DA{serial, elem, pos, kick_vec(serial), frc, mass}

def init_kicks(ds: List<&2, DynAtom>) -> List<&2, DynAtom>:
  match ds:
    case Nil{}:
      Nil{}
    case h <> t:
      init_kick_atom(h) <> init_kicks(t)

def sum_vels_ds(ds: List<&2, DynAtom>) -> G.Vec3:
  match ds:
    case Nil{}:
      G.V3{0.0, 0.0, 0.0}
    case h <> t:
      match h:
        case DA{serial, elem, pos, vel, frc, mass}:
          G.Vec3.add(vel, sum_vels_ds(t))

def sub_com_ds(ds: List<&2, DynAtom>, +com: G.Vec3) -> List<&2, DynAtom>:
  match ds:
    case Nil{}:
      Nil{}
    case h <> t:
      match h:
        case DA{serial, elem, pos, vel, frc, mass}:
          DA{serial, elem, pos, G.Vec3.sub(vel, com), frc, mass} <> sub_com_ds(t, com)

def ke_head(h: DynAtom) -> F32:
  match h:
    case DA{serial, elem, pos, vel, frc, mass}:
      ((0.5 * mass : F32) * G.Vec3.norm2(vel) : F32)

def ke_sim_of(ds: List<&2, DynAtom>) -> F32:
  match ds:
    case Nil{}:
      0.0
    case h <> t:
      (ke_head(h) + ke_sim_of(t) : F32)

def temp_of_ke(ke: F32, n: Nat) -> F32:
  ((2.0 * (ke / acc_conv() : F32) : F32) / ((3.0 * F32.from_nat(n) : F32) * boltzmann() : F32) : F32)

def temp_of_ds(ds: List<&2, DynAtom>, n: Nat) -> F32:
  temp_of_ke(ke_sim_of(ds), n)

def vv_scale_head(h: DynAtom, lam: F32) -> DynAtom:
  match h:
    case DA{serial, elem, pos, vel, frc, mass}:
      DA{serial, elem, pos, G.Vec3.scale(lam, vel), frc, mass}

def vv_scale(ds: List<&2, DynAtom>, +lam: F32) -> List<&2, DynAtom>:
  match ds:
    case Nil{}:
      Nil{}
    case h <> t:
      vv_scale_head(h, lam) <> vv_scale(t, lam)

def scale_to_temp(ds: List<&2, DynAtom>, +now: F32, want: F32) -> List<&2, DynAtom>:
  vv_scale(ds, Bool.pick(F32, F32.is_eq(now, 0.0), 1.0, F32.sqrt((want / now : F32))))

def com_of_n(+kicked: List<&2, DynAtom>, n: Nat) -> G.Vec3:
  G.Vec3.scale((1.0 / F32.from_nat(n) : F32), sum_vels_ds(kicked))

def init_scaled(+nocom: List<&2, DynAtom>, +n: Nat, temp0: F32) -> List<&2, DynAtom>:
  scale_to_temp(nocom, temp_of_ds(nocom, n), temp0)

def init_nocom(+kicked: List<&2, DynAtom>, +n: Nat, temp0: F32) -> List<&2, DynAtom>:
  init_scaled(sub_com_ds(kicked, com_of_n(kicked, n)), n, temp0)

def init_vels(+ds: List<&2, DynAtom>, +temp0: F32) -> List<&2, DynAtom>:
  init_nocom(init_kicks(ds), List.length(&2, DynAtom, ds), temp0)

# Set per-atom forces from a force list (truncates to the shorter).
def dyn_set_force(ds: List<&2, DynAtom>, fs: List<&2, G.Vec3>) -> List<&2, DynAtom>:
  match ds:
    case Nil{}:
      Nil{}
    case h <> t:
      match h:
        case DA{serial, elem, pos, vel, frc, mass}:
          match fs:
            case Nil{}:
              DA{serial, elem, pos, vel, G.V3{0.0, 0.0, 0.0}, mass} <> t
            case f <> ft:
              DA{serial, elem, pos, vel, f, mass} <> dyn_set_force(t, ft)

def add_forces(a: List<&2, G.Vec3>, b: List<&2, G.Vec3>) -> List<&2, G.Vec3>:
  match a b:
    case Nil{} Nil{}:
      Nil{}
    case Nil{} bh <> bt:
      Nil{}
    case ah <> at Nil{}:
      Nil{}
    case ah <> at bh <> bt:
      G.Vec3.add(ah, bh) <> add_forces(at, bt)

# Fresh forces (uniform-LJ with exclusions + Coulomb + bonds) into state.
def set_forces(+ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams) -> List<&2, DynAtom>:
  match p:
    case SP{+eps, +sig2, +c2, +ke, +kbond}:
      dyn_set_force(ds, add_forces(add_forces(B.lj_forces_excl(dyn_atoms(ds), dyn_atoms(ds), eps, sig2, c2, bonds), F.coul_forces(dyn_atoms(ds), qs, dyn_atoms(ds), qs, ke, c2)), B.bond_forces(dyn_atoms(ds), bonds, kbond)))

# Velocity Verlet halves over explicit state (forces must be fresh).
def vv_half_head(h: DynAtom, +dt: F32) -> DynAtom:
  match h:
    case DA{serial, elem, pos, vel, +frc, +mass}:
      acc = G.Vec3.scale(((1.0 / mass : F32) * acc_conv() : F32), frc)
      +vh = G.Vec3.add(vel, G.Vec3.scale((0.5 * dt : F32), acc))
      DA{serial, elem, G.Vec3.add(pos, G.Vec3.scale(dt, vh)), vh, frc, mass}

def vv_half(ds: List<&2, DynAtom>, +dt: F32) -> List<&2, DynAtom>:
  match ds:
    case Nil{}:
      Nil{}
    case h <> t:
      vv_half_head(h, dt) <> vv_half(t, dt)

def vv_full_head(h: DynAtom, dt: F32) -> DynAtom:
  match h:
    case DA{serial, elem, pos, vel, +frc, +mass}:
      acc = G.Vec3.scale(((1.0 / mass : F32) * acc_conv() : F32), frc)
      DA{serial, elem, pos, G.Vec3.add(vel, G.Vec3.scale((0.5 * dt : F32), acc)), frc, mass}

def vv_full(ds: List<&2, DynAtom>, +dt: F32) -> List<&2, DynAtom>:
  match ds:
    case Nil{}:
      Nil{}
    case h <> t:
      vv_full_head(h, dt) <> vv_full(t, dt)

# Berendsen thermostat folded into the step (T == 0 keeps velocities).
def md_thermo(+ds: List<&2, DynAtom>, +dt: F32, +tau: F32, +temp0: F32, +n: Nat) -> List<&2, DynAtom>:
  +now = temp_of_ds(ds, n)
  vv_scale(ds, Bool.pick(F32, F32.is_eq(now, 0.0), 1.0, F32.sqrt((1.0 + (dt / tau : F32) * ((temp0 / now : F32) - 1.0 : F32) : F32))))

def md_step(+ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams, +dt: F32, +tau: F32, +temp0: F32, +n: Nat) -> List<&2, DynAtom>:
  md_thermo(vv_full(set_forces(vv_half(ds, dt), qs, bonds, p), dt), dt, tau, temp0, n)

def md_run(steps: Nat, ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams, +dt: F32, +tau: F32, +temp0: F32, +n: Nat) -> List<&2, DynAtom>:
  match steps:
    case 0n:
      ds
    case 1n+f:
      md_run(f, md_step(ds, qs, bonds, p, dt, tau, temp0, n), qs, bonds, p, dt, tau, temp0, n)

# Steepest descent over state (fixed fuel; CLI reports convergedness).
# Normalized steps (fixed distance along each force) relax clashes that
# raw F*dt would overshoot; raw forces still feed MD and max-force.
def sd_norm_head(h: DynAtom, +step: F32) -> DynAtom:
  match h:
    case DA{serial, elem, pos, vel, +frc, mass}:
      +n2 = G.Vec3.norm2(frc)
      DA{serial, elem, G.Vec3.add(pos, G.Vec3.scale(step, Bool.pick(G.Vec3, F32.is_eq(n2, 0.0), G.V3{0.0, 0.0, 0.0}, G.Vec3.scale((1.0 / F32.sqrt(n2) : F32), frc)))), vel, frc, mass}

def sd_norm_sweep(ds: List<&2, DynAtom>, +step: F32) -> List<&2, DynAtom>:
  match ds:
    case Nil{}:
      Nil{}
    case h <> t:
      sd_norm_head(h, step) <> sd_norm_sweep(t, step)

def min_run(fuel: Nat, ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams, +dt: F32) -> List<&2, DynAtom>:
  match fuel:
    case 0n:
      ds
    case 1n+f:
      fds = set_forces(ds, qs, bonds, p)
      min_run(f, sd_norm_sweep(fds, dt), qs, bonds, p, dt)

def maxforce2_of(ds: List<&2, DynAtom>) -> F32:
  match ds:
    case Nil{}:
      0.0
    case h <> t:
      match h:
        case DA{serial, elem, pos, vel, frc, mass}:
          F32.max(G.Vec3.norm2(frc), maxforce2_of(t))

def converged_ds(ds: List<&2, DynAtom>, +tol: F32) -> Bool:
  F32.is_lt(maxforce2_of(ds), (tol * tol : F32))

# Flat bonds back to pairs for the pair-list energies.
def pairs_of_flat(bs: List<&2, U32>) -> List<&1, U32 & U32>:
  match bs:
    case Nil{}:
      Nil{}
    case x <> Nil{}:
      Nil{}
    case x <> y <> t:
      (x, y) <> pairs_of_flat(t)

# Reportable potential: excluded uniform LJ + Coulomb + bonds.
def energy_sum(+xs: List<&2, P.Atom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams) -> F32:
  match p:
    case SP{+eps, +sig2, +c2, +ke, +kbond}:
      ((B.lj_total_excl(xs, xs, eps, sig2, c2, bonds) + B.coul_total_excl(xs, qs, xs, qs, ke, c2, bonds) : F32) + B.bond_total(xs, pairs_of_flat(bonds), kbond) : F32)

def total_energy(+ds: List<&2, DynAtom>, +qs: List<&2, F32>, +bonds: List<&2, U32>, +p: SimParams) -> F32:
  energy_sum(dyn_atoms(ds), qs, bonds, p)