## Tsetlin-machine pattern-classifier gun — a MIXTURE OF EXPERTS with a learned
## TM gate.
##
## Why this exists (and why it is not the existing `tsetlin.nim` gun): the old
## Tsetlin gun uses the TM as a pixel-correction REGRESSOR on a single baseline
## and ranks 12th of 13. Here the TM does what Granmo's machine is actually
## strong at — supervised multi-class classification — on the frame-stacked
## binary encoding that measures 99.35-99.48% held-out accuracy at 2 frames.
##
## Architecture
## ------------
##   experts : HeadOn, Linear, Circular, WallBounce, Accel (the existing
##             analytic guns, reused unchanged and cheap).
##   gate    : the TM. Every tick it sees the same 2-frame Gray-coded binary
##             vector as the old Tsetlin gun and votes for the expert most
##             likely to be right.
##   label   : EXACTLY observable and supervised — at virtual-bullet
##             resolution the FeedbackEvent carries the enemy's actual
##             position, so the label is simply which expert's stored
##             prediction was CLOSEST. No delayed credit, no eligibility trace.
##   output  : the winning class's expert prediction for the requested
##             bulletSpeed (so the gun implements the ordinary
##             `predict(state, bulletSpeed)` interface and drops into the rack).
##
## TM reuse: encoding (`tmEncodeFrame`/`tmEncodeSelf`/`tmEncodeFullVector`) and
## ALL learning primitives come from `guns/tsetlin.nim` — in particular the
## CORRECTED Granmo feedback (`tmLearnDir`/`tmLearnOne`: Type I conditioned on
## the clause output, reachable Type II, the (T - clip(v,-T,T))/(2T) resource
## allocation, and the Eq. 6 all-Exclude bootstrap). Nothing is re-derived here.
##
## Multi-class formulation: Granmo's standard one-clause-team-per-class. Class c
## is a binary TM (`d = +1` for the winning expert, `-1` for the rest) and the
## predicted class is the argmax of the class votes. `TM_N_OUT = 2` independent
## clause teams already live in one `TmNet`, so 5 classes fit in 3 nets.

import std/[math, random, strformat, algorithm]
import gun_harness/gun_interface
import gun_harness/virtual_bullets as vb
import guns/tsetlin
import guns/head_on
import guns/linear
import guns/circular
import guns/wall_bounce
import guns/accel_predictor

const
  N_EXPERTS*   = 5
  N_NETS       = (N_EXPERTS + 1) div 2   ## 3 nets × 2 outputs = 6 class teams
  SEL_TRACE_SLOTS = 1024                 ## exact (fireTick,powerBin) ring, as tsetlin.nim
  SEL_MIN_OBS  = 50   ## observations before the TM outvotes the bootstrap expert
  DebugSelector* = false

  ExpertNames*: array[N_EXPERTS, string] =
    ["HeadOn", "Linear", "Circular", "WallBounce", "Accel"]

type
  TmSelTrace = object
    fireTick: int
    powerBin: int
    preds:    array[N_EXPERTS, GunPrediction]  ## fire-time expert predictions
    votes:    array[N_NETS * 2, float]      ## fire-time class votes (clamped)
    cache:    array[N_NETS, TmClauseCache]     ## fire-time LEARNING clause outputs
    input:    TmBinaryVector                   ## fire-time encoded input
    alive:    bool

  TmSelectorGun* = object
    nets:       array[N_NETS, TmNet]
    frameBuff:  array[2, TmFrameEncoded]  ## [0]=newest, [1]=previous
    frameCount: int
    lastTick:   int
    input:      TmBinaryVector
    curCaches:  array[N_NETS, TmClauseCache]
    votes:      array[N_NETS * 2, float]
    winCount*:  array[N_EXPERTS, int]  ## cumulative winners (bootstrap only)
    totalObs*:  int
    traces:     array[SEL_TRACE_SLOTS, TmSelTrace]
    # experts
    headOn:     HeadOnGun
    linear:     LinearGun
    circular:   CircularGun
    wallBounce: WallBounceGun
    accel:      AccelGun
    # instrumentation
    predictCalls*: int
    trainCalls*:   int
    traceMisses*:  int
    lastChosen*:   int
    debugGraphics*: bool

# ── helpers ──────────────────────────────────────────────────────────────────

proc selBinForSpeed(spd: float): int {.inline.} =
  for i in 0..<len(vb.PowerBins):
    if abs(spd - bulletSpeed(vb.PowerBins[i])) < 1e-6:
      return i
  -1

proc selTraceSlot(fireTick, binIdx: int): int {.inline.} =
  ((fireTick * len(vb.PowerBins)) + binIdx) mod SEL_TRACE_SLOTS

proc expertPred*(g: var TmSelectorGun, idx: int, state: WorldState,
                 bulletSpeed: float): GunPrediction =
  case idx
  of 0: g.headOn.predict(state, bulletSpeed)
  of 1: g.linear.predict(state, bulletSpeed)
  of 2: g.circular.predict(state, bulletSpeed)
  of 3: g.wallBounce.predict(state, bulletSpeed)
  of 4: g.accel.predict(state, bulletSpeed)
  else: g.headOn.predict(state, bulletSpeed)

proc initTmSelectorGun*(): TmSelectorGun =
  # states start at the Exclude boundary (0); one Type I step crosses into Include.
  for k in 0..<N_NETS:
    for s in result.nets[k].states.mitems: s = 0'i16
  result.lastTick = -1
  result.lastChosen = 2  # Circular
  randomize()
  result.debugGraphics = false

proc isWarmedUp*(g: TmSelectorGun): bool {.inline.} =
  ## The gun needs the 2-frame window to encode; experts handle colder states.
  g.frameCount >= 2

# ── class selection ──────────────────────────────────────────────────────────

proc chooseClass(g: TmSelectorGun): int =
  ## Bootstrap to the empirically best expert until the TM has SEL_MIN_OBS
  ## labels; afterwards take the argmax class vote (ties broken at random, so no
  ## index-0 bias toward HeadOn).
  if g.totalObs < SEL_MIN_OBS:
    var bestCount = -1
    for c in 0..<N_EXPERTS:
      if g.winCount[c] > bestCount:
        bestCount = g.winCount[c]
        result = c
    if bestCount <= 0: return 2  # Circular — sensible cold default
    return
  var bestV = -Inf
  for c in 0..<N_EXPERTS:
    if g.votes[c] > bestV: bestV = g.votes[c]
  var tied: seq[int]
  for c in 0..<N_EXPERTS:
    if g.votes[c] >= bestV - 1e-9: tied.add c
  result = tied[rand(tied.len - 1)]

# ── Gun interface ────────────────────────────────────────────────────────────

proc predict*(g: var TmSelectorGun, state: WorldState, bulletSpeed: float): GunPrediction =
  inc g.predictCalls

  # Encode the current frame and refresh the TM votes at most once per tick.
  # The harness calls predict() once per power bin (4×/tick); the TM input does
  # not depend on bulletSpeed, so the forward pass is tick-guarded exactly like
  # the old Tsetlin gun's window shift.
  if state.tick != g.lastTick:
    g.lastTick = state.tick
    g.frameBuff[1] = g.frameBuff[0]
    let dist    = hypot(state.enemyX - state.selfX, state.enemyY - state.selfY)
    let bearing = radToDeg(arctan2(state.enemyY - state.selfY, state.enemyX - state.selfX))
    g.frameBuff[0] = tmEncodeFrame(
      bearing, dist, state.enemySpeed, state.enemyHeading,
      state.arenaHeight - state.enemyY, state.enemyY,
      state.arenaWidth - state.enemyX, state.enemyX,
      state.enemyEnergy)
    if g.frameCount < 2: inc g.frameCount

    # 2-frame window (the measured sweet spot: 39.6 effective literals/clause vs
    # 152.8 at 10 frames). All TM_WINDOW_SIZE slots are filled with a REAL frame
    # (alternating newest/previous) rather than zero-padding: constant-zero
    # literals have an always-true negation, which Type I then includes en masse
    # and saturates every clause (measured mean 256-755 included literals).
    # Duplicating the two real frames keeps every literal variable. Reuses the
    # shared 10-frame encoder/vector so the corrected TM primitives apply
    # unchanged.
    var window: array[TM_WINDOW_SIZE, TmFrameEncoded]
    let prev = if g.frameCount >= 2: g.frameBuff[1] else: g.frameBuff[0]
    for i in 0..<TM_WINDOW_SIZE:
      window[i] = if (i and 1) == 0: g.frameBuff[0] else: prev
    let selfState = tmEncodeSelf(
      state.arenaHeight - state.selfY, state.selfY,
      state.arenaWidth - state.selfX, state.selfX,
      state.selfEnergy, true)
    g.input = tmEncodeFullVector(window, selfState)
    for k in 0..<N_NETS:
      let (_, _, vx, vy) = tmForwardWithCache(g.nets[k], g.input, g.curCaches[k])
      g.votes[2 * k]     = vx
      g.votes[2 * k + 1] = vy

  var preds: array[N_EXPERTS, GunPrediction]
  for c in 0..<N_EXPERTS:
    preds[c] = g.expertPred(c, state, bulletSpeed)

  let chosen = g.chooseClass()
  g.lastChosen = chosen
  when DebugSelector:
    echo fmt"[sel-dbg] tick={state.tick} bin={selBinForSpeed(bulletSpeed)} chosen={ExpertNames[chosen]} " &
         fmt"votes=[{g.votes[0]:.1f},{g.votes[1]:.1f},{g.votes[2]:.1f},{g.votes[3]:.1f},{g.votes[4]:.1f}] obs={g.totalObs}"

  let binIdx = selBinForSpeed(bulletSpeed)
  if binIdx >= 0:
    let slot = selTraceSlot(state.tick, binIdx)
    g.traces[slot] = TmSelTrace(
      fireTick: state.tick, powerBin: binIdx,
      preds: preds, votes: g.votes, cache: g.curCaches, input: g.input, alive: true)

  preds[chosen]

proc onResult*(g: var TmSelectorGun, e: FeedbackEvent) =
  let binIdx =
    if e.powerBin >= 0 and e.powerBin < len(vb.PowerBins): e.powerBin
    else: selBinForSpeed(bulletSpeed(e.bulletPower))
  if binIdx < 0:
    inc g.traceMisses
    return
  let slot = selTraceSlot(e.fireTick, binIdx)
  var t = addr g.traces[slot]
  if not t.alive or t.fireTick != e.fireTick or t.powerBin != binIdx:
    inc g.traceMisses
    return

  # The label: which expert's fire-time prediction was closest to the actual
  # enemy position the virtual bullet resolved against. Exact and supervised.
  var winner = 0
  var bestD = Inf
  for c in 0..<N_EXPERTS:
    let d = hypot(t.preds[c].x - e.actualX, t.preds[c].y - e.actualY)
    if d < bestD:
      bestD = d
      winner = c
  inc g.winCount[winner]
  inc g.totalObs
  inc g.trainCalls

  let lits = tmMakeLiterals(t.input)
  for c in 0..<N_EXPERTS:
    let k = c div 2
    let o = c mod 2
    let d = if c == winner: 1.0 else: -1.0
    g.nets[k].tmLearnDir(o, lits, t.cache[k], t.votes[c], d)
  t.alive = false

# ── interpretability ─────────────────────────────────────────────────────────

const
  FrameBitNames: array[TM_FRAME_BITS, string] = block:
    var a: array[TM_FRAME_BITS, string]
    for i in 0..<8:  a[i]      = "bearSin"
    for i in 0..<8:  a[8+i]    = "bearCos"
    for i in 0..<7:  a[16+i]   = "dist"
    for i in 0..<5:  a[23+i]   = "vel"
    for i in 0..<8:  a[28+i]   = "headSin"
    for i in 0..<8:  a[36+i]   = "headCos"
    for i in 0..<7:  a[44+i]   = "wallN"
    for i in 0..<7:  a[51+i]   = "wallS"
    for i in 0..<7:  a[58+i]   = "wallE"
    for i in 0..<7:  a[65+i]   = "wallW"
    for i in 0..<11: a[72+i]   = "energy"
    a

proc bitName(bit: int): string =
  ## Human name for a base vector bit index (0..TM_N_IN-1).
  if bit < TM_FRAME_BITS * TM_WINDOW_SIZE:
    let frame = bit div TM_FRAME_BITS
    let off   = bit mod TM_FRAME_BITS
    result = fmt"f{frame}.{FrameBitNames[off]}"
  else:
    let off = bit - TM_FRAME_BITS * TM_WINDOW_SIZE
    if off < 7: result = "self.wallN"
    elif off < 14: result = "self.wallS"
    elif off < 21: result = "self.wallE"
    elif off < 28: result = "self.wallW"
    elif off < 39: result = "self.energy"
    else: result = "self.canFire"

proc litName(lit: int): string =
  ## Literal `lit` is positive when `lit < TM_N_IN`, negated otherwise.
  if lit < TM_N_IN: bitName(lit) & "=1"
  else: bitName(lit - TM_N_IN) & "=0"

proc selectorClauseStats*(g: TmSelectorGun): tuple[nClauses, nActive: int,
                                                   meanIncluded: float] =
  ## Include-count over all class clause teams (5 × 50 clauses).
  var total = 0
  result.nClauses = N_EXPERTS * TM_N_CLAUSES
  for c in 0..<N_EXPERTS:
    let k = c div 2
    let o = c mod 2
    for cl in 0..<TM_N_CLAUSES:
      var inc = 0
      for lit in 0..<TM_N_LITERALS:
        if g.nets[k].states[tmStateIdx(o, cl, lit)] > 0: inc += 1
      total += inc
      if inc > 0: inc result.nActive
  if result.nActive > 0:
    result.meanIncluded = total.float / result.nActive.float

proc describeClauses*(g: TmSelectorGun, topN = 10): string =
  ## Per expert class, the feature literals included most often (and only in
  ## slots f0/f1 — slots 2..9 are the zero padding). One line per top feature:
  ## `feature=value × count` where count is how many of the class's 50 clauses
  ## include it. This is the interpretability payoff: it shows WHAT the gate
  ## switches on.
  for c in 0..<N_EXPERTS:
    let k = c div 2
    let o = c mod 2
    # count clauses per literal
    var counts: array[TM_N_LITERALS, int]
    for cl in 0..<TM_N_CLAUSES:
      for lit in 0..<TM_N_LITERALS:
        if g.nets[k].states[tmStateIdx(o, cl, lit)] > 0: inc counts[lit]
    var idx: seq[int]
    for lit in 0..<TM_N_LITERALS:
      # Frames alternate f0 (newest) / f1 (previous); slots 2..9 are duplicates,
      # so the unique selection signal lives in the first two frames. Report
      # only those to keep the dump readable.
      let base = if lit < TM_N_IN: lit else: lit - TM_N_IN
      if base >= TM_FRAME_BITS * 2 and base < TM_FRAME_BITS * TM_WINDOW_SIZE: continue
      if counts[lit] > 0: idx.add lit
    idx.sort(proc(a, b: int): int = counts[b] - counts[a])
    result.add fmt"class {c} ({ExpertNames[c]}):"
    if idx.len == 0:
      result.add " <no active literals>\n"
      continue
    result.add " "
    for i in 0..<min(topN, idx.len):
      result.add fmt"{litName(idx[i])}×{counts[idx[i]]}"
      if i < min(topN, idx.len) - 1: result.add ", "
    result.add "\n"