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

# --- Periodic boundaries v0: orthorhombic box, minimum image. ---
# The box rides as a Vec3 of edge lengths (Lx, Ly, Lz). Minimum image per
# component: d - L*round(d/L). Vacuum (non-periodic) runs ignore this
# module; solvated follow-ups use these kernels. No PME: Coulomb under PBC
# would still be cutoff (like everything here), so only LJ ships PBC
# variants for now.

def mic1d(+d: F32, +len: F32) -> F32:
  (d - (len * F32.round((d / len : F32)) : F32) : F32)

def mic_vec(+d: G.Vec3, +box: G.Vec3) -> G.Vec3:
  match d box:
    case G.V3{+dx, +dy, +dz} G.V3{+lx, +ly, +lz}:
      G.V3{mic1d(dx, lx), mic1d(dy, ly), mic1d(dz, lz)}

def mic_dist2(+a: G.Vec3, +b: G.Vec3, +box: G.Vec3) -> F32:
  G.Vec3.norm2(mic_vec(G.Vec3.sub(a, b), box))

def lj_pair_pbc(+pi: G.Vec3, +pj: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3) -> F32:
  +d2 = mic_dist2(pi, pj, box)
  F.lj_pair_go(d2, eps, sig2, F32.is_lt(d2, c2))

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

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

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

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

def lj_fat_pbc(+pi: G.Vec3, +pj: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3) -> G.Vec3:
  +dv = mic_vec(G.Vec3.sub(pi, pj), box)
  +d2 = G.Vec3.norm2(dv)
  F.lj_fgo(d2, dv, eps, sig2, F32.is_lt(d2, c2))

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

def lj_frow_pbc(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, +box: G.Vec3, 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_pbc(pi, eps, sig2, c2, box, t), lj_fhead_pbc(pi, eps, sig2, c2, box, h))

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

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