import Base
import ./geom.bend as G
import ./protein.bend as P
import ./topology.bend as T
import ./force.bend as F

# --- Bonded terms + exclusions v0. ---
# Harmonic bonds/angles with UNIFORM parameters over explicit serial lists;
# 1-2 pairs excluded by bond membership, 1-3 by shared neighbors, all over
# flat (Data, reusable) U32 bond lists with (a, b) pairs laid as [a, b, ..].
# A flat list is reusable where a pair list (always affine) is not, which is
# what lets exclusion tests re-scan bonds inside rows. Missing atoms
# contribute 0.0 (documented skip: every lookup is a Maybe pinned by
# pos_of_serial laws). Angles need explicit triples; deriving them needs
# residue templates (CONECT alone only bonds what it lists). PBC, 1-4
# scaling and bonded FORCES are follow-ups.

def pos_of_serial_go(xs: List<&2, P.Atom>, +s: U32, best: Maybe<&2, G.Vec3>) -> Maybe<&2, G.Vec3>:
  match xs:
    case Nil{}:
      best
    case h <> t:
      match h:
        case P.Atom{serial, elem, +pos}:
          pos_of_serial_go(t, s, Bool.pick(Maybe<&2, G.Vec3>, U32.is_eq(serial, s), Some{pos}, best))

def pos_of_serial(xs: List<&2, P.Atom>, s: U32) -> Maybe<&2, G.Vec3>:
  pos_of_serial_go(xs, s, None{})

# Harmonic bond energy k*(d - (ri+rj))^2 with covalent equilibrium; missing
# atoms answer 0.0. Lookups return (position & element) together so the
# equilibrium needs no extra pass.
def cov_rad(e: U32) -> F32:
  match e:
    case 1:
      0.31
    case 6:
      0.76
    case 7:
      0.71
    case 8:
      0.66
    case 15:
      1.07
    case 16:
      1.05
    case _:
      1.0

def atom_of_serial_go(xs: List<&2, P.Atom>, +s: U32, best: Maybe<&1, G.Vec3 & U32>) -> Maybe<&1, G.Vec3 & U32>:
  match xs:
    case Nil{}:
      best
    case h <> t:
      match h:
        case P.Atom{serial, elem, +pos}:
          atom_of_serial_go(t, s, Bool.pick(Maybe<&1, G.Vec3 & U32>, U32.is_eq(serial, s), Some{(pos, elem)}, best))

def atom_of_serial(xs: List<&2, P.Atom>, s: U32) -> Maybe<&1, G.Vec3 & U32>:
  atom_of_serial_go(xs, s, None{})

def bond_e(+pa: G.Vec3, +pb: G.Vec3, +k: F32, +r0: F32) -> F32:
  (k * F.lj_sq((G.Vec3.dist(pa, pb) - r0 : F32)) : F32)

def bond_cov_look1(pa: G.Vec3, ea: U32, mb: Maybe<&1, G.Vec3 & U32>, +k: F32) -> F32:
  match mb:
    case None{}:
      0.0
    case Some{(pb, eb)}:
      bond_e(pa, pb, k, (cov_rad(ea) + cov_rad(eb) : F32))

def bond_cov_look(ma: Maybe<&1, G.Vec3 & U32>, mb: Maybe<&1, G.Vec3 & U32>, +k: F32) -> F32:
  match ma:
    case None{}:
      0.0
    case Some{(pa, ea)}:
      bond_cov_look1(pa, ea, mb, k)

def bond_cov_pair(+xs: List<&2, P.Atom>, s1: U32, s2: U32, +k: F32) -> F32:
  bond_cov_look(atom_of_serial(xs, s1), atom_of_serial(xs, s2), k)

def bond_total(+xs: List<&2, P.Atom>, bonds: List<&1, U32 & U32>, +k: F32) -> F32:
  match bonds:
    case Nil{}:
      0.0
    case (s1, s2) <> t:
      (bond_cov_pair(xs, s1, s2, k) + bond_total(xs, t, k) : F32)

# Harmonic angle energy k*(theta - eq)^2 about the middle serial.
def angle_e(+pa: G.Vec3, +pb: G.Vec3, +pc: G.Vec3, +k: F32, +eq: F32) -> F32:
  (k * F.lj_sq((G.Vec3.angle(G.Vec3.sub(pa, pb), G.Vec3.sub(pc, pb)) - eq : F32)) : F32)

def angle_e0(pa: G.Vec3, pb: G.Vec3, mc: Maybe<&2, G.Vec3>, +k: F32, +eq: F32) -> F32:
  match mc:
    case None{}:
      0.0
    case Some{pc}:
      angle_e(pa, pb, pc, k, eq)

def angle_e1(pa: G.Vec3, mb: Maybe<&2, G.Vec3>, mc: Maybe<&2, G.Vec3>, +k: F32, +eq: F32) -> F32:
  match mb:
    case None{}:
      0.0
    case Some{pb}:
      angle_e0(pa, pb, mc, k, eq)

def angle_e2(ma: Maybe<&2, G.Vec3>, mb: Maybe<&2, G.Vec3>, mc: Maybe<&2, G.Vec3>, +k: F32, +eq: F32) -> F32:
  match ma:
    case None{}:
      0.0
    case Some{pa}:
      angle_e1(pa, mb, mc, k, eq)

def angle_triple(+xs: List<&2, P.Atom>, sa: U32, sb: U32, sc: U32, +k: F32, +eq: F32) -> F32:
  angle_e2(pos_of_serial(xs, sa), pos_of_serial(xs, sb), pos_of_serial(xs, sc), k, eq)

def angle_total(+xs: List<&2, P.Atom>, triples: List<&1, U32 & U32 & U32>, +k: F32, +eq: F32) -> F32:
  match triples:
    case Nil{}:
      0.0
    case (sa, sb, sc) <> t:
      (angle_triple(xs, sa, sb, sc, k, eq) + angle_total(xs, t, k, eq) : F32)

# Pair bonds -> flat [a, b, ..] so tests can re-scan (Data reuses).
def flatten_bonds(bs: List<&1, U32 & U32>) -> List<&2, U32>:
  match bs:
    case Nil{}:
      Nil{}
    case (x, y) <> t:
      x <> y <> flatten_bonds(t)

def bond_match(+a: U32, +b: U32, +x: U32, +y: U32) -> Bool:
  Bool.or(Bool.and(U32.is_eq(a, x), U32.is_eq(b, y)), Bool.and(U32.is_eq(a, y), U32.is_eq(b, x)))

def bonded12_flat(bs: List<&2, U32>, +a: U32, +b: U32) -> Bool:
  match bs:
    case Nil{}:
      False{}
    case x <> Nil{}:
      False{}
    case x <> y <> t:
      Bool.or(bond_match(a, b, x, y), bonded12_flat(t, a, b))

def bond_neighbors_go(bs: List<&2, U32>, +s: U32, +acc: List<&2, U32>) -> List<&2, U32>:
  match bs:
    case Nil{}:
      List.reverse(&2, U32, acc)
    case x <> Nil{}:
      List.reverse(&2, U32, acc)
    case x <> y <> t:
      +x2 = x
      +y2 = y
      bond_neighbors_go(t, s, Bool.pick(List<&2, U32>, U32.is_eq(x2, s), y2 <> acc, Bool.pick(List<&2, U32>, U32.is_eq(y2, s), x2 <> acc, acc)))

def bond_neighbors(s: U32, bs: List<&2, U32>) -> List<&2, U32>:
  bond_neighbors_go(bs, s, Nil{})

def list_member(+x: U32, ys: List<&2, U32>) -> Bool:
  match ys:
    case Nil{}:
      False{}
    case h <> t:
      Bool.or(U32.is_eq(x, h), list_member(x, t))

def share_one(xs: List<&2, U32>, +ys: List<&2, U32>) -> Bool:
  match xs:
    case Nil{}:
      False{}
    case h <> t:
      Bool.or(list_member(h, ys), share_one(t, ys))

def excluded12(a: U32, b: U32, +bonds: List<&2, U32>) -> Bool:
  bonded12_flat(bonds, a, b)

def excluded13(a: U32, b: U32, +bonds: List<&2, U32>) -> Bool:
  share_one(bond_neighbors(a, bonds), bond_neighbors(b, bonds))

def excluded(+a: U32, +b: U32, +bonds: List<&2, U32>) -> Bool:
  Bool.or(excluded12(a, b, bonds), excluded13(a, b, bonds))

# Nonbonded rows/totals that skip 1-2 and 1-3 pairs by serial.
def lj_row_excl(+si: U32, +pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>, +bonds: List<&2, U32>) -> F32:
  match ys:
    case Nil{}:
      0.0
    case h <> t:
      match h:
        case P.Atom{serial, elem, +pj}:
          (Bool.pick(F32, excluded(si, serial, bonds), 0.0, F.lj_pair_at(pi, pj, eps, sig2, c2)) + lj_row_excl(si, pi, eps, sig2, c2, t, bonds) : F32)

def lj_total_excl(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +eps: F32, +sig2: F32, +c2: F32, +bonds: List<&2, U32>) -> F32:
  match xs:
    case Nil{}:
      0.0
    case h <> t:
      match h:
        case P.Atom{+serial, elem, +pos}:
          (lj_total_excl(t, ys, eps, sig2, c2, bonds) + lj_row_excl(serial, pos, eps, sig2, c2, ys, bonds) : F32)

def coul_row_excl(+si: U32, +pi: G.Vec3, +qi: F32, +ke: F32, +c2: F32, ys: List<&2, P.Atom>, qs: List<&2, F32>, +bonds: List<&2, U32>) -> 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:
      +h2 = h
      (Bool.pick(F32, excluded(si, P.Atom.serial(h2), bonds), 0.0, F.coul_at(pi, P.Atom.pos(h2), (qi * qh : F32), ke, c2)) + coul_row_excl(si, pi, qi, ke, c2, t, qt, bonds) : F32)

def coul_total_excl(xs: List<&2, P.Atom>, +qs_o: List<&2, F32>, +ys: List<&2, P.Atom>, +qs_e: List<&2, F32>, +ke: F32, +c2: F32, +bonds: List<&2, U32>) -> 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:
      +h2 = h
      (coul_total_excl(t, qt, ys, qs_e, ke, c2, bonds) + coul_row_excl(P.Atom.serial(h2), P.Atom.pos(h2), qh, ke, c2, ys, qs_e, bonds) : F32)

# Per-element rows/totals that also skip 1-2 and 1-3 pairs: the physically
# meaningful combination (uniform rows blow up on bonded neighbors).
def lj_row_elem_excl(+si: U32, +se: U32, +pi: G.Vec3, +ys: List<&2, P.Atom>, +c2: F32, +bonds: List<&2, U32>) -> F32:
  match ys:
    case Nil{}:
      0.0
    case h <> t:
      match h:
        case P.Atom{serial, +elem, +pj}:
          (Bool.pick(F32, excluded(si, serial, bonds), 0.0, F.lj_pair_at(pi, pj, F.lj_combine_eps(F.lj_eps(se), F.lj_eps(elem)), F.lj_combine_sig2(F.lj_sig(se), F.lj_sig(elem)), c2)) + lj_row_elem_excl(si, se, pi, t, c2, bonds) : F32)

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

# Excluded LJ forces mirror lj_forces, skipping 1-2/1-3 pairs by serial.
def lj_frow_excl(+si: U32, +pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>, +bonds: List<&2, U32>) -> G.Vec3:
  match ys:
    case Nil{}:
      G.V3{0.0, 0.0, 0.0}
    case h <> t:
      match h:
        case P.Atom{serial, elem, +pj}:
          G.Vec3.add(lj_frow_excl(si, pi, eps, sig2, c2, t, bonds), Bool.pick(G.Vec3, excluded(si, serial, bonds), G.V3{0.0, 0.0, 0.0}, F.lj_fat(pi, pj, eps, sig2, c2)))

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

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

# Bond forces: F = -2k(d-r0) * unit(pa-pb) on a, opposite on b; zero for
# coincident or missing atoms. Per-atom loop over flat bonds.
def bond_fpair(+d: G.Vec3, +d2: F32, +k: F32, +r0: F32) -> G.Vec3 & G.Vec3:
  +r = F32.sqrt(d2)
  +mag = Bool.pick(F32, F32.is_eq(d2, 0.0), 0.0, (((0.0 - (2.0 * k : F32) : F32) * (r - r0 : F32) : F32) / r : F32))
  +fa = G.Vec3.scale(mag, d)
  (fa, G.Vec3.neg(fa))

def bond_fpick(eqa: Bool, eqb: Bool, fab: G.Vec3 & G.Vec3) -> G.Vec3:
  match fab:
    case (fa, fb):
      Bool.pick(G.Vec3, eqa, fa, Bool.pick(G.Vec3, eqb, fb, G.V3{0.0, 0.0, 0.0}))

def bond_fsas(+s: U32, sa: U32, sb: U32, +pa: G.Vec3, +pb: G.Vec3, +k: F32, +r0: F32) -> G.Vec3:
  bond_fpick(U32.is_eq(s, sa), U32.is_eq(s, sb), bond_fpair(G.Vec3.sub(pa, pb), G.Vec3.dist2(pa, pb), k, r0))

def bond_force_look1(s: U32, pa: G.Vec3, ea: U32, mb: Maybe<&1, G.Vec3 & U32>, a: U32, b: U32, +k: F32) -> G.Vec3:
  match mb:
    case None{}:
      G.V3{0.0, 0.0, 0.0}
    case Some{(pb, eb)}:
      bond_fsas(s, a, b, pa, pb, k, (cov_rad(ea) + cov_rad(eb) : F32))

def bond_force_look(s: U32, ma: Maybe<&1, G.Vec3 & U32>, mb: Maybe<&1, G.Vec3 & U32>, a: U32, b: U32, +k: F32) -> G.Vec3:
  match ma:
    case None{}:
      G.V3{0.0, 0.0, 0.0}
    case Some{(pa, ea)}:
      bond_force_look1(s, pa, ea, mb, a, b, k)

def bond_force_bond(+xs: List<&2, P.Atom>, s: U32, +a: U32, +b: U32, +k: F32) -> G.Vec3:
  bond_force_look(s, atom_of_serial(xs, a), atom_of_serial(xs, b), a, b, k)

def bond_force_list(+s: U32, bs: List<&2, U32>, +xs: List<&2, P.Atom>, +k: F32) -> G.Vec3:
  match bs:
    case Nil{}:
      G.V3{0.0, 0.0, 0.0}
    case x <> Nil{}:
      G.V3{0.0, 0.0, 0.0}
    case x <> y <> t:
      G.Vec3.add(bond_force_bond(xs, s, x, y, k), bond_force_list(s, t, xs, k))

def bond_forces(+xs: List<&2, P.Atom>, +bonds: List<&2, U32>, +k: F32) -> List<&2, G.Vec3>:
  match xs:
    case Nil{}:
      Nil{}
    case +h <> t:
      bond_force_list(P.Atom.serial(h), bonds, xs, k) <> bond_forces(t, bonds, k)

def bond_cut2(ei: U32, ej: U32) -> F32:
  +s = (cov_rad(ei) + cov_rad(ej) : F32)
  +t = (s * 1.25 : F32)
  (t * t : F32)

def infer_one(si: U32, pi: G.Vec3, ei: U32, sj: U32, ej: U32, pj: G.Vec3) -> List<&1, U32 & U32>:
  Bool.pick(List<&1, U32 & U32>, F32.is_lt(G.Vec3.dist2(pi, pj), bond_cut2(ei, ej)), [(si, sj)], Nil{})

def append_infer(a: List<&1, U32 & U32>, b: List<&1, U32 & U32>) -> List<&1, U32 & U32>:
  match a:
    case Nil{}:
      b
    case h <> t:
      h <> append_infer(t, b)

def infer_row(+si: U32, +pi: G.Vec3, +ei: U32, ys: List<&2, P.Atom>) -> List<&1, U32 & U32>:
  match ys:
    case Nil{}:
      Nil{}
    case +h <> t:
      append_infer(infer_one(si, pi, ei, P.Atom.serial(h), P.Atom.elem(h), P.Atom.pos(h)), infer_row(si, pi, ei, t))

def infer_total(+xs: List<&2, P.Atom>) -> List<&1, U32 & U32>:
  match xs:
    case Nil{}:
      Nil{}
    case +h <> t:
      append_infer(infer_row(P.Atom.serial(h), P.Atom.pos(h), P.Atom.elem(h), t), infer_total(t))

# Angle auto-derivation from flat bonds: for each bond, fan over the other
# neighbors of each end. Needs no templates; duplicates across overlapping
# fans are possible in rings (documented, harmless for boolean use).
def bond_neighbors_except_go(bs: List<&2, U32>, +s: U32, +excl: U32, +acc: List<&2, U32>) -> List<&2, U32>:
  match bs:
    case Nil{}:
      List.reverse(&2, U32, acc)
    case x <> Nil{}:
      List.reverse(&2, U32, acc)
    case x <> y <> t:
      +x2 = x
      +y2 = y
      bond_neighbors_except_go(t, s, excl, Bool.pick(List<&2, U32>, U32.is_eq(x2, s), Bool.pick(List<&2, U32>, U32.is_eq(y2, excl), acc, y2 <> acc), Bool.pick(List<&2, U32>, U32.is_eq(y2, s), Bool.pick(List<&2, U32>, U32.is_eq(x2, excl), acc, x2 <> acc), acc)))

def bond_neighbors_except(s: U32, excl: U32, bs: List<&2, U32>) -> List<&2, U32>:
  bond_neighbors_except_go(bs, s, excl, Nil{})

def fan_triples(+x: U32, +y: U32, ns: List<&2, U32>) -> List<&1, U32 & U32 & U32>:
  match ns:
    case Nil{}:
      Nil{}
    case c <> t:
      (x, y, c) <> fan_triples(x, y, t)

def append_triples(a: List<&1, U32 & U32 & U32>, b: List<&1, U32 & U32 & U32>) -> List<&1, U32 & U32 & U32>:
  match a:
    case Nil{}:
      b
    case h <> t:
      h <> append_triples(t, b)

def angle_fan_both(+a: U32, +b: U32, +flat: List<&2, U32>) -> List<&1, U32 & U32 & U32>:
  append_triples(fan_triples(a, b, bond_neighbors_except(b, a, flat)), fan_triples(b, a, bond_neighbors_except(a, b, flat)))

def angle_triples_flat_go(bs: List<&2, U32>, +flat: List<&2, U32>) -> List<&1, U32 & U32 & U32>:
  match bs:
    case Nil{}:
      Nil{}
    case x <> Nil{}:
      Nil{}
    case x <> y <> t:
      append_triples(angle_fan_both(x, y, flat), angle_triples_flat_go(t, flat))

def angle_triples_auto(+flat: List<&2, U32>) -> List<&1, U32 & U32 & U32>:
  angle_triples_flat_go(flat, flat)