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

# --- SASA lite v0: Shrake-Rupley with 6 octahedral sample points.
# Uniform radii: rad scales the sample sphere, rc2 is the squared occlusion
# cutoff, area_pt is the surface area each exposed point stands for (4*pi*r^2/6).
# The occluder set `ys` must contain every atom INCLUDING the source atom;
# the source atom is skipped by serial (which callers must keep unique),
# because its own center sits exactly `rad` from each of its sample points
# and would otherwise self-occlude whenever rc2 > rad^2.

def sasa_offsets() -> List<&2, G.Vec3>:
  [G.V3{1.0, 0.0, 0.0}, G.V3{(0.0 - 1.0 : F32), 0.0, 0.0}, G.V3{0.0, 1.0, 0.0}, G.V3{0.0, (0.0 - 1.0 : F32), 0.0}, G.V3{0.0, 0.0, 1.0}, G.V3{0.0, 0.0, (0.0 - 1.0 : F32)}]

def occ_go(same: Bool, +pos: G.Vec3, +pt: G.Vec3, +rc2: F32) -> Bool:
  match same:
    case True{}:
      False{}
    case False{}:
      F32.is_lt(G.Vec3.dist2(pos, pt), rc2)

def occ_head(h: P.Atom, +pt: G.Vec3, +rc2: F32, skip: U32) -> Bool:
  match h:
    case P.Atom{serial, elem, +pos}:
      occ_go(U32.is_eq(serial, skip), pos, pt, rc2)

# Any atom but `skip` within rc2 of the sample point occludes it.
def has_other_close(
  xs: List<&2, P.Atom>, +pt: G.Vec3, +rc2: F32, +skip: U32
) -> Bool:
  match xs:
    case Nil{}:
      False{}
    case h <> t:
      Bool.or(has_other_close(t, pt, rc2, skip), occ_head(h, pt, rc2, skip))

def sasa_pt(
  +c: G.Vec3, +o: G.Vec3, +rad: F32, +ys: List<&2, P.Atom>, +rc2: F32,
  +skip: U32
) -> Nat:
  T.hit_nat_go(Bool.not(has_other_close(ys, G.Vec3.add(c, G.Vec3.scale(rad, o)), rc2, skip)))

def sasa_cover(
  +c: G.Vec3, +offs: List<&2, G.Vec3>, +rad: F32, +ys: List<&2, P.Atom>,
  +rc2: F32, +skip: U32
) -> Nat:
  match offs:
    case Nil{}:
      0n
    case o <> t:
      Nat.add(sasa_cover(c, t, rad, ys, rc2, skip), sasa_pt(c, o, rad, ys, rc2, skip))

def sasa_of_atom(
  h: P.Atom, +offs: List<&2, G.Vec3>, +rad: F32, +ys: List<&2, P.Atom>,
  +rc2: F32, +area_pt: F32
) -> F32:
  match h:
    case P.Atom{+serial, elem, +pos}:
      (area_pt * F32.from_nat(sasa_cover(pos, offs, rad, ys, rc2, serial)) : F32)

def sasa_all(
  xs: List<&2, P.Atom>, +offs: List<&2, G.Vec3>, +rad: F32,
  +ys: List<&2, P.Atom>, +rc2: F32, +area_pt: F32
) -> List<&2, F32>:
  match xs:
    case Nil{}:
      Nil{}
    case h <> t:
      sasa_of_atom(h, offs, rad, ys, rc2, area_pt) <> sasa_all(t, offs, rad, ys, rc2, area_pt)

def sasa_of(
  +xs: List<&2, P.Atom>, +rad: F32, +rc2: F32, +area_pt: F32
) -> List<&2, F32>:
  sasa_all(xs, sasa_offsets(), rad, xs, rc2, area_pt)

# Total area: sum of per-atom areas.
def sasa_total(areas: List<&2, F32>) -> F32:
  match areas:
    case Nil{}:
      0.0
    case h <> t:
      (h + sasa_total(t) : F32)