import Base
import ./geom.bend as G

# --- Protein model v0: Atom / Residue / Chain / Complex. All Data. ---

type Atom is Data:
  Atom{serial: U32, elem: U32, pos: G.Vec3}

type Residue is Data:
  Residue{seq: Nat, atoms: List<&2, Atom>}

type Chain is Data:
  Chain{cid: U32, residues: List<&2, Residue>}

type Complex is Data:
  Complex{chains: List<&2, Chain>}

# Atom position (drops serial/elem, keeps pos)
def Atom.pos(a: Atom) -> G.Vec3:
  match a:
    case Atom{serial, elem, pos}:
      pos

# Atom element number (drops serial/pos, keeps elem)
def Atom.elem(a: Atom) -> U32:
  match a:
    case Atom{serial, elem, pos}:
      elem

# Atom serial number (drops elem/pos, keeps serial)
def Atom.serial(a: Atom) -> U32:
  match a:
    case Atom{serial, elem, pos}:
      serial

def Atom.dist2(+a: Atom, +b: Atom) -> F32:
  match a b:
    case Atom{serial1, elem1, +pa} Atom{serial2, elem2, +pb}:
      G.Vec3.dist2(pa, pb)

def Atom.dist(+a: Atom, +b: Atom) -> F32:
  F32.sqrt(Atom.dist2(a, b))

def count_atoms(xs: List<&2, Atom>) -> Nat:
  match xs:
    case Nil{}:
      0n
    case h <> t:
      1n+count_atoms(t)

def count_residues(xs: List<&2, Residue>) -> Nat:
  match xs:
    case Nil{}:
      0n
    case h <> t:
      1n+count_residues(t)

def sum_count_extend_go(pos: G.Vec3, rest: G.Vec3 & Nat) -> G.Vec3 & Nat:
  match rest:
    case (rest_sum, rest_n):
      (G.Vec3.add(pos, rest_sum), 1n+rest_n)

def sum_count_extend(h: Atom, rest: G.Vec3 & Nat) -> G.Vec3 & Nat:
  match h:
    case Atom{serial, elem, pos}:
      sum_count_extend_go(pos, rest)

# Sum positions + count in one pass (list consumed once)
def sum_count_atoms(xs: List<&2, Atom>) -> G.Vec3 & Nat:
  match xs:
    case Nil{}:
      (G.V3{0.0, 0.0, 0.0}, 0n)
    case h <> t:
      sum_count_extend(h, sum_count_atoms(t))

def centroid_of_sum_go(sum: G.Vec3, n: Nat) -> G.Vec3:
  match n:
    case 0n:
      sum
    case 1n+p:
      G.Vec3.scale((1.0 / (F32.from_nat(n)) : F32), sum)

def centroid_of_sum(pair: G.Vec3 & Nat) -> G.Vec3:
  match pair:
    case (sum, n):
      centroid_of_sum_go(sum, n)

def Residue.centroid(r: Residue) -> G.Vec3:
  match r:
    case Residue{seq, atoms}:
      centroid_of_sum(sum_count_atoms(atoms))

def Residue.num_atoms(r: Residue) -> Nat:
  match r:
    case Residue{seq, atoms}:
      count_atoms(atoms)

def Chain.num_residues(c: Chain) -> Nat:
  match c:
    case Chain{cid, residues}:
      count_residues(residues)

def append_atoms(xs: List<&2, Atom>, ys: List<&2, Atom>) -> List<&2, Atom>:
  match xs:
    case Nil{}:
      ys
    case h <> t:
      h <> append_atoms(t, ys)

def flatten_res_extend(h: Residue, rest: List<&2, Atom>) -> List<&2, Atom>:
  match h:
    case Residue{seq, atoms}:
      append_atoms(atoms, rest)

def flatten_residues(xs: List<&2, Residue>) -> List<&2, Atom>:
  match xs:
    case Nil{}:
      Nil{}
    case h <> t:
      flatten_res_extend(h, flatten_residues(t))

def atoms_of(h: Residue) -> Nat:
  match h:
    case Residue{seq, atoms}:
      count_atoms(atoms)

def count_atoms_in(xs: List<&2, Residue>) -> Nat:
  match xs:
    case Nil{}:
      0n
    case h <> t:
      Nat.add(count_atoms_in(t), atoms_of(h))

def Chain.flatten(c: Chain) -> List<&2, Atom>:
  match c:
    case Chain{cid, residues}:
      flatten_residues(residues)

def Chain.num_atoms(c: Chain) -> Nat:
  match c:
    case Chain{cid, residues}:
      count_atoms_in(residues)

def Chain.centroid(c: Chain) -> G.Vec3:
  match c:
    case Chain{cid, residues}:
      centroid_of_sum(sum_count_atoms(flatten_residues(residues)))

def count_chains(xs: List<&2, Chain>) -> Nat:
  match xs:
    case Nil{}:
      0n
    case h <> t:
      1n+count_chains(t)

def Complex.num_chains(k: Complex) -> Nat:
  match k:
    case Complex{chains}:
      count_chains(chains)

def append_residues(
  xs: List<&2, Residue>, ys: List<&2, Residue>
) -> List<&2, Residue>:
  match xs:
    case Nil{}:
      ys
    case h <> t:
      h <> append_residues(t, ys)

def flatten_chains_extend(
  h: Chain, rest: List<&2, Residue>
) -> List<&2, Residue>:
  match h:
    case Chain{cid, residues}:
      append_residues(residues, rest)

def flatten_chains_residues(xs: List<&2, Chain>) -> List<&2, Residue>:
  match xs:
    case Nil{}:
      Nil{}
    case h <> t:
      flatten_chains_extend(h, flatten_chains_residues(t))

def Complex.num_residues(k: Complex) -> Nat:
  match k:
    case Complex{chains}:
      count_residues(flatten_chains_residues(chains))

def Complex.flatten(k: Complex) -> List<&2, Atom>:
  match k:
    case Complex{chains}:
      flatten_residues(flatten_chains_residues(chains))

def Complex.num_atoms(k: Complex) -> Nat:
  match k:
    case Complex{chains}:
      count_atoms_in(flatten_chains_residues(chains))

def Complex.centroid(k: Complex) -> G.Vec3:
  match k:
    case Complex{chains}:
      centroid_of_sum(sum_count_atoms(flatten_residues(flatten_chains_residues(chains))))