import Base
import ./geom.bend as G
import ./protein.bend as P

# --- Nonbonded kernels v0: uniform Lennard-Jones + Coulomb, squared cutoffs.
# Self pairs (d2 == 0) and pairs beyond c2 contribute nothing.
# Convention: for self energy/forces, pass the same reusable list twice,
# e.g. lj_total(+xs, xs, eps, sig2, c2) style call sites. That call is a
# DOUBLE sum: every unordered pair {i, j} is counted twice (i,j and j,i),
# so the physical total is half the answer. The diagonal is dropped by
# POSITION (d2 == 0.0), not by index: callers must give distinct atoms
# distinct positions, since a coincident distinct pair is silently dropped
# instead of diverging.

def lj_sq(+x: F32) -> F32:
  (x * x : F32)

def lj_cube(+inv: F32) -> F32:
  ((inv * inv : F32) * inv : F32)

def lj_eval(+u6: F32, +eps: F32) -> F32:
  ((4.0 * eps : F32) * ((lj_sq(u6) - u6 : F32)) : F32)

def lj_from_inv(inv: F32, +eps: F32) -> F32:
  lj_eval(lj_sq(lj_cube(inv)), eps)

# LJ energy from squared distance: 4e*((s2/d2)^12 - (s2/d2)^6), no sqrt.
def lj_of_d2(d2: F32, +eps: F32, +sig2: F32) -> F32:
  lj_from_inv((sig2 / d2 : F32), eps)

def lj_self_go(d2: F32, +eps: F32, +sig2: F32, self: Bool) -> F32:
  match self:
    case True{}:
      0.0
    case False{}:
      lj_of_d2(d2, eps, sig2)

def lj_pair_go(+d2: F32, +eps: F32, +sig2: F32, below: Bool) -> F32:
  match below:
    case True{}:
      lj_self_go(d2, eps, sig2, F32.is_eq(d2, 0.0))
    case False{}:
      0.0

def lj_pair_at(+pi: G.Vec3, +pj: G.Vec3, +eps: F32, +sig2: F32, +c2: F32) -> F32:
  +d2 = G.Vec3.dist2(pi, pj)
  lj_pair_go(d2, eps, sig2, F32.is_lt(d2, c2))

def lj_head(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, h: P.Atom) -> F32:
  match h:
    case P.Atom{serial, elem, +pj}:
      lj_pair_at(pi, pj, eps, sig2, c2)

def lj_row(
  +pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>
) -> F32:
  match ys:
    case Nil{}:
      0.0
    case h <> t:
      (lj_row(pi, eps, sig2, c2, t) + lj_head(pi, eps, sig2, c2, h) : F32)

def lj_self_atom(
  h: P.Atom, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32
) -> F32:
  match h:
    case P.Atom{serial, elem, +pi}:
      lj_row(pi, eps, sig2, c2, ys)

def lj_total(
  xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32
) -> F32:
  match xs:
    case Nil{}:
      0.0
    case h <> t:
      (lj_total(t, ys, eps, sig2, c2) + lj_self_atom(h, ys, eps, sig2, c2) : F32)

# --- Coulomb with per-atom charges in an aligned List<&2, F32>. ---
# The charge lists zip with the atom lists: a shorter charge list silently
# drops the remaining atoms (the Nil{} arms answer zero), so callers must
# keep charges aligned with atoms; see LAWS coul_trunc_charges.

def coul_formula(d2: F32, +qq: F32, +ke: F32) -> F32:
  ((ke * qq : F32) / F32.sqrt(d2) : F32)

def coul_self_go(d2: F32, +qq: F32, +ke: F32, self: Bool) -> F32:
  match self:
    case True{}:
      0.0
    case False{}:
      coul_formula(d2, qq, ke)

def coul_pair_go(+d2: F32, +qq: F32, +ke: F32, below: Bool) -> F32:
  match below:
    case True{}:
      coul_self_go(d2, qq, ke, F32.is_eq(d2, 0.0))
    case False{}:
      0.0

def coul_at(+pi: G.Vec3, +pj: G.Vec3, +qq: F32, +ke: F32, +c2: F32) -> F32:
  +d2 = G.Vec3.dist2(pi, pj)
  coul_pair_go(d2, qq, ke, F32.is_lt(d2, c2))

def coul_head(
  +pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, h: P.Atom, qj: F32
) -> F32:
  match h:
    case P.Atom{serial, elem, +pj}:
      coul_at(pi, pj, (qi * qj : F32), ke, c2)

def coul_row(
  +pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32,
  ys: List<&2, P.Atom>, qs: List<&2, F32>
) -> F32:
  match ys qs:
    case Nil{} Nil{}:
      0.0
    case Nil{} qh <> qt:
      0.0
    case h <> t Nil{}:
      0.0
    case h <> t qh <> qt:
      (coul_row(pi, qi, ke, c2, t, qt) + coul_head(pi, qi, ke, c2, h, qh) : F32)

def coul_self(
  h: P.Atom, +qh: F32, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>,
  +ke: F32, +c2: F32
) -> F32:
  match h:
    case P.Atom{serial, elem, +pi}:
      coul_row(pi, qh, ke, c2, ys, qs_e)

def coul_total(
  xs: List<&2, P.Atom>, +qs_o: List<&2, F32>, +ys: List<&2, P.Atom>,
  +qs_e: List<&2, F32>, +ke: F32, +c2: F32
) -> F32:
  match xs qs_o:
    case Nil{} Nil{}:
      0.0
    case Nil{} qh <> qt:
      0.0
    case h <> t Nil{}:
      0.0
    case h <> t qh <> qt:
      (coul_total(t, qt, ys, qs_e, ke, c2) + coul_self(h, qh, ys, qs_e, ke, c2) : F32)

# --- LJ forces: F_i = 24e*(2*u6^2 - u6)/d2 * (pi - pj), u6 = (sig2/d2)^3^2. ---

def lj_k(+u6: F32, +eps: F32, +d2: F32) -> F32:
  (((24.0 * eps : F32) * (((2.0 * lj_sq(u6) : F32) - u6 : F32)) : F32) / d2 : F32)

def lj_fvec(+d: G.Vec3, +k: F32) -> G.Vec3:
  G.Vec3.scale(k, d)

def lj_fpair(+d2: F32, +d: G.Vec3, +eps: F32, +sig2: F32) -> G.Vec3:
  lj_fvec(d, lj_k(lj_sq(lj_cube((sig2 / d2 : F32))), eps, d2))

def lj_fself(+d2: F32, +d: G.Vec3, +eps: F32, +sig2: F32, self: Bool) -> G.Vec3:
  match self:
    case True{}:
      G.V3{0.0, 0.0, 0.0}
    case False{}:
      lj_fpair(d2, d, eps, sig2)

def lj_fgo(+d2: F32, +d: G.Vec3, +eps: F32, +sig2: F32, below: Bool) -> G.Vec3:
  match below:
    case True{}:
      lj_fself(d2, d, eps, sig2, F32.is_eq(d2, 0.0))
    case False{}:
      G.V3{0.0, 0.0, 0.0}

def lj_fat(+pi: G.Vec3, +pj: G.Vec3, +eps: F32, +sig2: F32, +c2: F32) -> G.Vec3:
  +d2 = G.Vec3.dist2(pi, pj)
  lj_fgo(d2, G.Vec3.sub(pi, pj), eps, sig2, F32.is_lt(d2, c2))

def lj_fhead(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, h: P.Atom) -> G.Vec3:
  match h:
    case P.Atom{serial, elem, +pj}:
      lj_fat(pi, pj, eps, sig2, c2)

def lj_frow(
  +pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>
) -> G.Vec3:
  match ys:
    case Nil{}:
      G.V3{0.0, 0.0, 0.0}
    case h <> t:
      G.Vec3.add(lj_frow(pi, eps, sig2, c2, t), lj_fhead(pi, eps, sig2, c2, h))

def lj_fself_atom(
  h: P.Atom, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32
) -> G.Vec3:
  match h:
    case P.Atom{serial, elem, +pi}:
      lj_frow(pi, eps, sig2, c2, ys)

def lj_forces(
  xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32
) -> List<&2, G.Vec3>:
  match xs:
    case Nil{}:
      Nil{}
    case h <> t:
      lj_fself_atom(h, ys, eps, sig2, c2) <> lj_forces(t, ys, eps, sig2, c2)

# --- Steepest descent: x += F * dt, fuel-bounded. Equal-length xs/fs. ---

def sd_move(h: P.Atom, f: G.Vec3, +dt: F32) -> P.Atom:
  match h f:
    case P.Atom{serial, elem, +pos} G.V3{+fx, +fy, +fz}:
      P.Atom{serial, elem, G.Vec3.add(pos, G.Vec3.scale(dt, G.V3{fx, fy, fz}))}

def sd_sweep(
  xs: List<&2, P.Atom>, fs: List<&2, G.Vec3>, +dt: F32
) -> List<&2, P.Atom>:
  match xs fs:
    case Nil{} Nil{}:
      Nil{}
    case Nil{} f <> ft:
      Nil{}
    case h <> t Nil{}:
      h <> t
    case h <> t f <> ft:
      sd_move(h, f, dt) <> sd_sweep(t, ft, dt)

def minimize(
  fuel: Nat, +xs: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +dt: F32
) -> List<&2, P.Atom>:
  match fuel:
    case 0n:
      xs
    case 1n+f:
      minimize(f, sd_sweep(xs, lj_forces(xs, xs, eps, sig2, c2), dt), eps, sig2, c2, dt)

# --- Per-element Lennard-Jones: illustrative eps (kcal/mol) and sigma (A)
# by element number; unknown elements take a generic (0.20, 3.00). These
# are placeholders, not a force field. Lorentz-Berthelot combining:
# eps = sqrt(ei*ej), sig2 = ((si+sj)/2)^2.
def lj_eps(e: U32) -> F32:
  match e:
    case 1:
      0.03
    case 6:
      0.066
    case 7:
      0.17
    case 8:
      0.21
    case 16:
      0.25
    case _:
      0.2

def lj_sig(e: U32) -> F32:
  match e:
    case 1:
      2.5
    case 6:
      3.5
    case 7:
      3.25
    case 8:
      2.96
    case 16:
      3.55
    case _:
      3.0

def lj_combine_eps(ei: F32, ej: F32) -> F32:
  F32.sqrt((ei * ej : F32))

def lj_combine_sig2(si: F32, sj: F32) -> F32:
  lj_sq(((si + sj : F32) / 2.0 : F32))

def lj_head_elem(+si: U32, +pi: G.Vec3, +c2: F32, h: P.Atom) -> F32:
  match h:
    case P.Atom{serial, +elem, +pj}:
      lj_pair_at(pi, pj, lj_combine_eps(lj_eps(si), lj_eps(elem)), lj_combine_sig2(lj_sig(si), lj_sig(elem)), c2)

def lj_row_elem(+si: U32, +pi: G.Vec3, +ys: List<&2, P.Atom>, +c2: F32) -> F32:
  match ys:
    case Nil{}:
      0.0
    case h <> t:
      (lj_row_elem(si, pi, t, c2) + lj_head_elem(si, pi, c2, h) : F32)

def lj_total_elem(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +c2: F32) -> F32:
  match xs:
    case Nil{}:
      0.0
    case +h <> t:
      (lj_total_elem(t, ys, c2) + lj_row_elem(P.Atom.elem(h), P.Atom.pos(h), ys, c2) : F32)

# --- Coulomb forces: F_i = ke*qi*qj/d2^1.5 * (pi - pj), repulsive for
# like charges along +d. Self pairs (d2 == 0) and pairs beyond c2 give zero.
def coul_k(+qq: F32, +ke: F32, +d2: F32) -> F32:
  ((ke * qq : F32) / ((d2 * F32.sqrt(d2) : F32)) : F32)

def coul_fpair(+d2: F32, +d: G.Vec3, +qq: F32, +ke: F32) -> G.Vec3:
  G.Vec3.scale(coul_k(qq, ke, d2), d)

def coul_fself(+d2: F32, +d: G.Vec3, +qq: F32, +ke: F32, self: Bool) -> G.Vec3:
  match self:
    case True{}:
      G.V3{0.0, 0.0, 0.0}
    case False{}:
      coul_fpair(d2, d, qq, ke)

def coul_fgo(+d2: F32, +d: G.Vec3, +qq: F32, +ke: F32, below: Bool) -> G.Vec3:
  match below:
    case True{}:
      coul_fself(d2, d, qq, ke, F32.is_eq(d2, 0.0))
    case False{}:
      G.V3{0.0, 0.0, 0.0}

def coul_fat(+pi: G.Vec3, +pj: G.Vec3, +qq: F32, +ke: F32, +c2: F32) -> G.Vec3:
  +d2 = G.Vec3.dist2(pi, pj)
  coul_fgo(d2, G.Vec3.sub(pi, pj), qq, ke, F32.is_lt(d2, c2))

def coul_fhead(+pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, h: P.Atom, qj: F32) -> G.Vec3:
  match h:
    case P.Atom{serial, elem, +pj}:
      coul_fat(pi, pj, (qi * qj : F32), ke, c2)

def coul_frow(+pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, ys: List<&2, P.Atom>, qs: List<&2, F32>) -> G.Vec3:
  match ys qs:
    case Nil{} Nil{}:
      G.V3{0.0, 0.0, 0.0}
    case Nil{} qh <> qt:
      G.V3{0.0, 0.0, 0.0}
    case h <> t Nil{}:
      G.V3{0.0, 0.0, 0.0}
    case h <> t qh <> qt:
      G.Vec3.add(coul_frow(pi, qi, ke, c2, t, qt), coul_fhead(pi, qi, ke, c2, h, qh))

def coul_fself_atom(h: P.Atom, +qh: F32, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>, +ke: F32, +c2: F32) -> G.Vec3:
  match h:
    case P.Atom{serial, elem, +pi}:
      coul_frow(pi, qh, ke, c2, ys, qs_e)

def coul_forces(xs: List<&2, P.Atom>, +qs_o: List<&2, F32>, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>, +ke: F32, +c2: F32) -> List<&2, G.Vec3>:
  match xs qs_o:
    case Nil{} Nil{}:
      Nil{}
    case Nil{} qh <> qt:
      Nil{}
    case h <> t Nil{}:
      Nil{}
    case h <> t qh <> qt:
      coul_fself_atom(h, qh, ys, qs_e, ke, c2) <> coul_forces(t, qt, ys, qs_e, ke, c2)

# --- Atomic masses (amu) by element; unknown -> 12.0. ---
def mass(e: U32) -> F32:
  match e:
    case 1:
      1.008
    case 6:
      12.011
    case 7:
      14.007
    case 8:
      15.999
    case 15:
      30.974
    case 16:
      32.06
    case _:
      12.0