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

# --- RMSD v0: centroid-free aligned RMSD (no rotation) + ensemble matrix.
# Lockstep pairs zip: a longer list's extra elements are ignored, so
# `rmsd` of unequal lists silently answers 0. Use `rmsd_checked` when the
# mismatch must be visible instead.

def pair_head_go(+p1: G.Vec3, y: P.Atom, rest: F32 & Nat) -> F32 & Nat:
  match y:
    case P.Atom{s2, e2, +p2}:
      T.sumd2_extend(G.Vec3.dist2(p1, p2), rest)

def pair_head(x: P.Atom, y: P.Atom, rest: F32 & Nat) -> F32 & Nat:
  match x:
    case P.Atom{s1, e1, +p1}:
      pair_head_go(p1, y, rest)

def pair_sumd2(xs: List<&2, P.Atom>, ys: List<&2, P.Atom>) -> F32 & Nat:
  match xs ys:
    case Nil{} Nil{}:
      (0.0, 0n)
    case Nil{} qh <> qt:
      (0.0, 0n)
    case h <> t Nil{}:
      (0.0, 0n)
    case x <> xt y <> yt:
      pair_head(x, y, pair_sumd2(xt, yt))

def rmsd(xs: List<&2, P.Atom>, ys: List<&2, P.Atom>) -> F32:
  T.rg_of(pair_sumd2(xs, ys))

# Length-checked RMSD: None{} on unequal lengths, Some{...} otherwise.
def rmsd_checked_go(
  ok: Bool, +xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>
) -> Maybe<&2, F32>:
  match ok:
    case True{}:
      Some{rmsd(xs, ys)}
    case False{}:
      None{}

def rmsd_checked(
  +xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>
) -> Maybe<&2, F32>:
  rmsd_checked_go(Nat.is_eq(List.length(&2, P.Atom, xs), List.length(&2, P.Atom, ys)), xs, ys)

def rmsd_against(+q: List<&2, P.Atom>, h: List<&2, P.Atom>) -> F32:
  T.rg_of(pair_sumd2(q, h))

def rmsd_to_all(
  +q: List<&2, P.Atom>, ys: List<&2, List<&2, P.Atom>>
) -> List<&2, F32>:
  match ys:
    case Nil{}:
      Nil{}
    case h <> t:
      rmsd_against(q, h) <> rmsd_to_all(q, t)

def rmsd_matrix(
  +xs: List<&2, List<&2, P.Atom>>, +all: List<&2, List<&2, P.Atom>>
) -> List<&2, List<&2, F32>>:
  match xs:
    case Nil{}:
      Nil{}
    case h <> t:
      rmsd_to_all(h, all) <> rmsd_matrix(t, all)

def rmsd_matrix_of(+ens: List<&2, List<&2, P.Atom>>) -> List<&2, List<&2, F32>>:
  rmsd_matrix(ens, ens)