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

# --- Topology v0: contacts, clashes, spread. Cutoffs are SQUARED (no sqrt). ---
# Parallel note: each pair test is independent; a future pass can evaluate
# `count_below` / `contacts_between` with balanced fork-join + GPU (`!`).

def hit_nat_go(hit: Bool) -> Nat:
  match hit:
    case True{}:
      1n
    case False{}:
      0n

def head_hit_nat(h: P.Atom, +qp: G.Vec3, +c2: F32) -> Nat:
  match h:
    case P.Atom{serial, elem, +pos}:
      hit_nat_go(F32.is_lt(G.Vec3.dist2(pos, qp), c2))

def count_below(xs: List<&2, P.Atom>, +qp: G.Vec3, +c2: F32) -> Nat:
  match xs:
    case Nil{}:
      0n
    case h <> t:
      Nat.add(count_below(t, qp, c2), head_hit_nat(h, qp, c2))

# Plain (non-squared) cutoff: squares it once, then counts.
def count_within(xs: List<&2, P.Atom>, +qp: G.Vec3, +cutoff: F32) -> Nat:
  count_below(xs, qp, (cutoff * cutoff : F32))

def head_is_close(h: P.Atom, +qp: G.Vec3, +min2: F32) -> Bool:
  match h:
    case P.Atom{serial, elem, +pos}:
      F32.is_lt(G.Vec3.dist2(pos, qp), min2)

def has_close(xs: List<&2, P.Atom>, +qp: G.Vec3, +min2: F32) -> Bool:
  match xs:
    case Nil{}:
      False{}
    case h <> t:
      Bool.or(has_close(t, qp, min2), head_is_close(h, qp, min2))

def head_dist2(h: P.Atom, +qp: G.Vec3) -> F32:
  match h:
    case P.Atom{serial, elem, +pos}:
      G.Vec3.dist2(pos, qp)

def sum_dist2(xs: List<&2, P.Atom>, +qp: G.Vec3) -> F32:
  match xs:
    case Nil{}:
      0.0
    case h <> t:
      (sum_dist2(t, qp) + head_dist2(h, qp) : F32)

def sumd2_extend(a: F32, rest: F32 & Nat) -> F32 & Nat:
  match rest:
    case (s, n):
      ((a + s : F32), 1n+n)

# One pass: summed squared distances about qp, plus atom count.
def sumd2_count(xs: List<&2, P.Atom>, +qp: G.Vec3) -> F32 & Nat:
  match xs:
    case Nil{}:
      (0.0, 0n)
    case h <> t:
      sumd2_extend(head_dist2(h, qp), sumd2_count(t, qp))

def rg_of_go(s: F32, n: Nat) -> F32:
  match n:
    case 0n:
      0.0
    case 1n+p:
      F32.sqrt((s / (F32.from_nat(n)) : F32))

def rg_of(pair: F32 & Nat) -> F32:
  match pair:
    case (s, n):
      rg_of_go(s, n)

# Radius of gyration of an atom list about a given center.
def rg_about(xs: List<&2, P.Atom>, +qp: G.Vec3) -> F32:
  rg_of(sumd2_count(xs, qp))

def contacts_of_head(h: P.Atom, +ys: List<&2, P.Atom>, +c2: F32) -> Nat:
  match h:
    case P.Atom{serial, elem, +pos}:
      count_below(ys, pos, c2)

# Interface size: pairs (x in xs, y in ys) with dist2 < c2. Consumes xs once;
# ys is reusable so every head reuses it (reference counted).
def contacts_between(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +c2: F32) -> Nat:
  match xs:
    case Nil{}:
      0n
    case h <> t:
      Nat.add(contacts_between(t, ys, c2), contacts_of_head(h, ys, c2))

def min_cons_go(d: F32, rest: Maybe<&2, F32>) -> Maybe<&2, F32>:
  match rest:
    case None{}:
      Some{d}
    case Some{best}:
      Some{F32.min(d, best)}

def min_cons(h: P.Atom, rest: Maybe<&2, F32>, +qp: G.Vec3) -> Maybe<&2, F32>:
  match h:
    case P.Atom{serial, elem, +pos}:
      min_cons_go(G.Vec3.dist2(pos, qp), rest)

# Closest squared distance from qp to any atom in xs. None{} when empty.
def min_dist2_to(xs: List<&2, P.Atom>, +qp: G.Vec3) -> Maybe<&2, F32>:
  match xs:
    case Nil{}:
      None{}
    case h <> t:
      min_cons(h, min_dist2_to(t, qp), qp)

def min_maybe_some(x: F32, b: Maybe<&2, F32>) -> Maybe<&2, F32>:
  match b:
    case None{}:
      Some{x}
    case Some{y}:
      Some{F32.min(x, y)}

def min_maybe(a: Maybe<&2, F32>, b: Maybe<&2, F32>) -> Maybe<&2, F32>:
  match a:
    case None{}:
      b
    case Some{x}:
      min_maybe_some(x, b)

def head_link(h: P.Atom, +ys: List<&2, P.Atom>) -> Maybe<&2, F32>:
  match h:
    case P.Atom{serial, elem, +pos}:
      min_dist2_to(ys, pos)

# Closest squared approach between two atom sets. None{} when either is empty.
def min_link2(xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>) -> Maybe<&2, F32>:
  match xs:
    case Nil{}:
      None{}
    case h <> t:
      min_maybe(min_link2(t, ys), head_link(h, ys))

def min_link(link: Maybe<&2, F32>) -> Maybe<&2, F32>:
  match link:
    case None{}:
      None{}
    case Some{d2}:
      Some{F32.sqrt(d2)}

# Per-atom neighbor counts: for each x in xs, atoms of ys within dist2 < c2.
def coord_counts(
  xs: List<&2, P.Atom>, +ys: List<&2, P.Atom>, +c2: F32
) -> List<&2, Nat>:
  match xs:
    case Nil{}:
      Nil{}
    case h <> t:
      contacts_of_head(h, ys, c2) <> coord_counts(t, ys, c2)