## Tsetlin Machine gun — regression TM for aiming correction on top of linear extrapolation.
## Self-contained: includes binary encoding and TM predictor inline.
## Implements Gun interface: predict(state, bulletSpeed) → GunPrediction, onResult(FeedbackEvent).

import std/[math, random, strformat]
import gun_harness/gun_interface
import gun_harness/virtual_bullets as vb  # PowerBins (power-bin count for trace keys)

# ── Binary encoding (adapted from BNNBot_garage/src/binary_encoding.nim) ─────

const
  TM_FRAME_BITS*  = 83
  TM_SELF_BITS*   = 40
  TM_WINDOW_SIZE* = 10
  TM_TOTAL_BITS* = TM_FRAME_BITS * TM_WINDOW_SIZE + TM_SELF_BITS  # 870
  TM_MAX_DISTANCE = 1414.0  # diagonal of 1000x1000 arena

type
  TmBinaryVector*  = array[TM_TOTAL_BITS, uint8]
  TmFrameEncoded*  = array[TM_FRAME_BITS, uint8]
  TmSelfEncoded*   = array[TM_SELF_BITS, uint8]

proc tmToGray(value: int): int = value xor (value shr 1)

proc tmToBits(value: int, bits: int): seq[uint8] =
  result = newSeq[uint8](bits)
  let gray = tmToGray(value)
  for i in 0..<bits:
    result[bits - 1 - i] = uint8((gray shr i) and 1)

proc tmEncodeFrame*(bearing, distance, velocity, heading,
                    wallN, wallS, wallE, wallW, energy: float): TmFrameEncoded =
  var offset = 0

  # bearing sin: 8 bits
  let bSin = tmToBits(clamp(int((sin(degToRad(bearing)) + 1.0) / 2.0 * 199.0), 0, 199), 8)
  for i in 0..<8: result[offset + i] = bSin[i]
  offset += 8

  # bearing cos: 8 bits
  let bCos = tmToBits(clamp(int((cos(degToRad(bearing)) + 1.0) / 2.0 * 199.0), 0, 199), 8)
  for i in 0..<8: result[offset + i] = bCos[i]
  offset += 8

  # distance %: 7 bits
  let distBits = tmToBits(clamp(int(distance / TM_MAX_DISTANCE * 99.0), 0, 99), 7)
  for i in 0..<7: result[offset + i] = distBits[i]
  offset += 7

  # velocity: 5 bits
  let velBits = tmToBits(clamp(int(velocity + 8.0), 0, 16), 5)
  for i in 0..<5: result[offset + i] = velBits[i]
  offset += 5

  # heading sin: 8 bits
  let hSin = tmToBits(clamp(int((sin(degToRad(heading)) + 1.0) / 2.0 * 199.0), 0, 199), 8)
  for i in 0..<8: result[offset + i] = hSin[i]
  offset += 8

  # heading cos: 8 bits
  let hCos = tmToBits(clamp(int((cos(degToRad(heading)) + 1.0) / 2.0 * 199.0), 0, 199), 8)
  for i in 0..<8: result[offset + i] = hCos[i]
  offset += 8

  # wall distances: 4×7 bits
  for wall in [wallN, wallS, wallE, wallW]:
    let wBits = tmToBits(clamp(int(wall / 1000.0 * 99.0), 0, 99), 7)
    for i in 0..<7: result[offset + i] = wBits[i]
    offset += 7

  # energy: 11 bits
  let eBits = tmToBits(clamp(int(energy * 10.0), 0, 1500), 11)
  for i in 0..<11: result[offset + i] = eBits[i]

proc tmEncodeSelf*(wallN, wallS, wallE, wallW, energy: float, canFire: bool): TmSelfEncoded =
  var offset = 0
  for wall in [wallN, wallS, wallE, wallW]:
    let wBits = tmToBits(clamp(int(wall / 1000.0 * 99.0), 0, 99), 7)
    for i in 0..<7: result[offset + i] = wBits[i]
    offset += 7
  let eBits = tmToBits(clamp(int(energy * 10.0), 0, 1500), 11)
  for i in 0..<11: result[offset + i] = eBits[i]
  offset += 11
  result[offset] = if canFire: 1'u8 else: 0'u8

proc tmEncodeFullVector*(window: array[TM_WINDOW_SIZE, TmFrameEncoded],
                         self: TmSelfEncoded): TmBinaryVector =
  var offset = 0
  for i in 0..<TM_WINDOW_SIZE:
    for j in 0..<TM_FRAME_BITS:
      result[offset] = window[i][j]; inc offset
  for j in 0..<TM_SELF_BITS:
    result[offset] = self[j]; inc offset

# ── Tsetlin Machine (adapted from BNNBot_garage/src/tsetlin_predictor.nim) ───

const
  TM_N_IN*        = TM_TOTAL_BITS         # 870
  TM_N_OUT*       = 2                     # cx, cy pixel corrections
  TM_N_LITERALS*  = TM_N_IN * 2          # 1740
  TM_N_CLAUSES*   = 50                    # per output; issue #184 default
  TM_HALF*        = TM_N_CLAUSES div 2
  TM_N_STATES*    = 32                    # automaton range [-32..32]
  TM_T*           = float(TM_HALF)        # vote clamped to [-T, T]
  TM_S*           = 1.5                   # specificity
  TM_RESID_MAX   = 80.0                  # pixel correction range
  # ponytail: TM_N_STATES=32 needs int16 (int8 only fits ≤127, fine here); raise N_CLAUSES if underfitting

type
  TmClauseCache* = array[TM_N_OUT * TM_N_CLAUSES, uint8]

  TmNet* = object
    states*: array[TM_N_OUT * TM_N_CLAUSES * TM_N_LITERALS, int16]
    # ponytail: int16 to safely hold [-32..32]; TM_N_STATES=32 fits int8 too but int16 is safer

proc tmStateIdx*(outIdx, clause, lit: int): int {.inline.} =
  (outIdx * TM_N_CLAUSES + clause) * TM_N_LITERALS + lit

proc tmPolarity*(clause: int): float {.inline.} =
  if clause < TM_HALF: 1.0 else: -1.0

type
  TmClauseStats* = object
    ## Include-count statistics over both clause teams (TM_N_OUT * TM_N_CLAUSES).
    ## "Active" = a clause with at least one included literal at the sampled
    ## moment. A healthy TM settles into a SPARSE regime (single digits to low
    ## tens of literals per clause). Saturation at hundreds of literals/clause
    ## (measured mean ≈714 on the energy-threshold fixture before the fix) is the
    ## failure signature of the broken Type I feedback rule.
    nClauses*:        int
    nActive*:         int
    minIncluded*:     int
    maxIncluded*:     int
    meanIncluded*:    float  ## mean include count over ACTIVE clauses
    meanIncludedAll*: float  ## mean include count over ALL clauses (incl. empty)

proc tmMakeLiterals*(input: TmBinaryVector): array[TM_N_LITERALS, uint8] =
  for i in 0..<TM_N_IN:
    result[i]            = input[i]
    result[i + TM_N_IN]  = 1'u8 - input[i]

proc tmEvalClause*(net: TmNet, outIdx, clause: int,
                  lits: array[TM_N_LITERALS, uint8],
                  learning = false): uint8 =
  var hasIncluded = false
  for lit in 0..<TM_N_LITERALS:
    let s = net.states[tmStateIdx(outIdx, clause, lit)]
    if s > 0:
      hasIncluded = true
      if lits[lit] == 0: return 0'u8
  if hasIncluded: return 1'u8
  # Granmo Eq. 6: an all-Exclude clause evaluates to 1 during LEARNING (the
  # empty conjunction is vacuously true) and 0 during classification. The
  # learning value is what bootstraps the clauses: at initialisation every
  # automaton sits at the Exclude boundary, and with the corrected Type I rule
  # (c=0, lk=1 -> toward Exclude) a classification-only 0 would deadlock every
  # clause at empty forever.
  return if learning: 1'u8 else: 0'u8

proc tmForwardWithCache*(net: TmNet, input: TmBinaryVector,
                        cache: var TmClauseCache): (float, float, float, float) =
  ## Returns (correctionX, correctionY, voteX, voteY). `cache` receives the
  ## clause outputs under LEARNING semantics (empty clause = 1) for tmLearnOne;
  ## the returned correction and votes use classification semantics (empty = 0).
  let lits = tmMakeLiterals(input)
  var vx = 0.0; var vy = 0.0
  for c in 0..<TM_N_CLAUSES:
    let o = tmEvalClause(net, 0, c, lits)
    cache[c] = tmEvalClause(net, 0, c, lits, learning = true)
    vx += tmPolarity(c) * float(o)
  for c in 0..<TM_N_CLAUSES:
    let o = tmEvalClause(net, 1, c, lits)
    cache[TM_N_CLAUSES + c] = tmEvalClause(net, 1, c, lits, learning = true)
    vy += tmPolarity(c) * float(o)
  vx = clamp(vx, -TM_T, TM_T)
  vy = clamp(vy, -TM_T, TM_T)
  (vx / TM_T * TM_RESID_MAX, vy / TM_T * TM_RESID_MAX, vx, vy)

proc tmLearnDir*(net: var TmNet, outIdx: int, lits: array[TM_N_LITERALS, uint8],
                 cache: TmClauseCache, vote: float, d: float) =
  ## One faithful Granmo Table 2/3 clause update with an EXPLICIT desired vote
  ## direction `d` in {-1, +1}. This is the exact corrected core the Tsetlin gun
  ## uses; `tmLearnOne` is the regression wrapper that derives `d` from a
  ## continuous residual, and the multi-class selector passes the class label
  ## directly.
  ##
  ## `cache` holds clause outputs under LEARNING semantics (empty = 1, Eq. 6).
  ## `vote` is the classification-semantics clause sum at prediction time,
  ## already clamped to [-TM_T, TM_T]. The Granmo resource allocation collapses
  ## to a single probability p = (T - d*clip(v,-T,T)) / (2T): it is high when the
  ## vote opposes `d` and falls to 0 once the class is already won.
  let pFeedback = (TM_T - d * vote) / (2.0 * TM_T)
  if pFeedback <= 0.0: return

  for c in 0..<TM_N_CLAUSES:
    if rand(1.0) >= pFeedback: continue
    let pol  = tmPolarity(c)
    let cOut = cache[outIdx * TM_N_CLAUSES + c]
    if pol * d > 0.0:
      # Type I (Table 2), collapsed to the resulting state move:
      #   c=1, lk=1 -> +1 (toward Include) w.p. (s-1)/s
      #   c=0, lk=1 -> -1 (toward Exclude) w.p.  1/s   <- the missing counter-force
      #   lk=0      -> -1 (toward Exclude) w.p.  1/s
      for lit in 0..<TM_N_LITERALS:
        let si = tmStateIdx(outIdx, c, lit)
        var st = int(net.states[si])
        if lits[lit] == 1'u8:
          if cOut == 1'u8:
            if rand(1.0) < (TM_S - 1.0) / TM_S: st = min(st + 1, TM_N_STATES)
          else:
            if rand(1.0) < 1.0 / TM_S:          st = max(st - 1, -TM_N_STATES)
        else:
          if rand(1.0) < 1.0 / TM_S:            st = max(st - 1, -TM_N_STATES)
        net.states[si] = int16(st)
    else:
      # Type II (Table 3): the only non-Inaction cell is
      #   c=1, lk=0, Exclude action -> Penalty -> toward Include.
      # (The old code penalised INCLUDED false literals, which is unreachable
      # when c=1 and the wrong direction.)
      if cOut == 1'u8:
        for lit in 0..<TM_N_LITERALS:
          if lits[lit] == 0'u8:
            let si = tmStateIdx(outIdx, c, lit)
            if net.states[si] <= 0:
              net.states[si] = int16(min(int(net.states[si]) + 1, TM_N_STATES))

proc tmLearnOne*(net: var TmNet, outIdx: int, lits: array[TM_N_LITERALS, uint8],
                 cache: TmClauseCache, vote: float, residual: float) =
  ## One faithful Granmo Table 2/3 update against a continuous residual target.
  ##
  ## `cache` holds clause outputs under LEARNING semantics (empty = 1, Eq. 6).
  ## `vote` is the classification-semantics clause sum at prediction time,
  ## already clamped to [-TM_T, TM_T]. `residual` is the correction target
  ## delta = actual - linear baseline (see onResult), NOT actual - prediction.
  ##
  ## Regression adaptation: the "label" direction is d = sign(residual -
  ## predicted), i.e. which way the correction must move. The resource
  ## allocation of Granmo Eq. 8-11 collapses to a single probability
  ## p = (T - d*clip(v,-T,T)) / (2T), applied both to the aligned clauses
  ## (Type I) and the opposed ones (Type II), exactly as Algorithm 1 lines 11-22.
  let predicted = vote / TM_T * TM_RESID_MAX
  let error = residual - predicted
  let d = if error > 0.0: 1.0 elif error < 0.0: -1.0 else: return
  net.tmLearnDir(outIdx, lits, cache, vote, d)

# ── TsetlinGun public type ────────────────────────────────────────────────────

const
  # Ring of pending traces keyed EXACTLY by (fireTick, powerBin). A power-3 shot
  # can take ~fireDist/speed ~ 128 ticks to resolve, and the rack stores 4 traces
  # per tick, so 1024 slots (> 128*4) guarantee a live trace is never overwritten
  # by a newer one. The old 64-slot ring held only ~13 ticks of traces.
  TM_TRACE_SLOTS = 1024
  DebugTM* = false      # set true to print [tm-dbg] lines per onResult call

type
  TmTrace = object
    fireTick:     int            # key part: tick the bullet was fired
    powerBin:     int            # key part: power bin the bullet belonged to
    predX, predY: float          # stored prediction (for debug/provenance)
    linearX, linearY: float      # baseline the correction was added to (fix 4)
    voteX, voteY: float          # classification clause sum at fire time
    input:        TmBinaryVector
    cache:        TmClauseCache  # LEARNING clause outputs (empty = 1)
    alive:        bool

  TsetlinGun* = object
    net:            TmNet
    frameBuffer:    array[TM_WINDOW_SIZE, TmFrameEncoded]
    bufferCount:    int
    frameTick:      int          # last tick the window was shifted (once per tick)
    traces:         array[TM_TRACE_SLOTS, TmTrace]
    shotCount:      int  ## total onResult calls received
    trainedShots*:  int  ## onResult calls that found and trained their exact trace
    traceMisses*:   int  ## onResult calls whose trace was gone (integrity counter)
    # Instrumentation: is the TM actually producing a nonzero correction? Before
    # the feedback-rule fix vHits was effectively identical to Linear because
    # the clauses over-specified to a mean of ~714 included literals and almost
    # never fired (correction was nonzero on only 8 of 764 predict calls).
    predictCalls*:        int    ## predict() invocations (4/tick in the harness)
    correctionsNonzero*:  int    ## predict() calls whose |cx|+|cy| > 1e-9
    lastCx*, lastCy*:     float  ## last correction emitted (debugging/assertions)
    debugGraphics*: bool

proc tmBinForSpeed(spd: float): int {.inline.} =
  ## Map a virtual-bullet speed back to its power-bin index.
  for i in 0..<len(vb.PowerBins):
    if abs(spd - bulletSpeed(vb.PowerBins[i])) < 1e-6:
      return i
  -1

proc tmTraceSlot(fireTick, binIdx: int): int {.inline.} =
  ## Exact (fireTick, powerBin) key -> ring slot. TM_TRACE_SLOTS is a multiple of
  ## the bin count and larger than maxResolveTicks*bins, so live traces never
  ## collide with newer ones; unresolved traces are evicted after ~256 ticks.
  ((fireTick * len(vb.PowerBins)) + binIdx) mod TM_TRACE_SLOTS

proc initTsetlinGun*(): TsetlinGun =
  # states init at 0 (boundary); one Type I step crosses into Include
  for s in result.net.states.mitems: s = 0'i16
  result.frameTick = -1
  randomize()
  result.debugGraphics = false

proc isWarmedUp*(g: TsetlinGun): bool {.inline.} =
  g.bufferCount >= TM_WINDOW_SIZE

proc tmClauseStats*(g: TsetlinGun): TmClauseStats =
  ## Sample the current include-count distribution of the automata. O(outputs ×
  ## clauses × literals); call on demand, not per-tick.
  result.nClauses = TM_N_OUT * TM_N_CLAUSES
  result.minIncluded = high(int)
  var total = 0
  for o in 0..<TM_N_OUT:
    for c in 0..<TM_N_CLAUSES:
      var inc = 0
      for lit in 0..<TM_N_LITERALS:
        if g.net.states[tmStateIdx(o, c, lit)] > 0: inc += 1
      total += inc
      if inc > 0:
        inc result.nActive
        result.minIncluded = min(result.minIncluded, inc)
        result.maxIncluded = max(result.maxIncluded, inc)
  if result.nActive > 0:
    result.meanIncluded = total.float / result.nActive.float
  else:
    result.minIncluded = 0
  result.meanIncludedAll = total.float / result.nClauses.float

proc predict*(g: var TsetlinGun, state: WorldState, bulletSpeed: float): GunPrediction =
  # Encode current frame and push into window
  let dist = hypot(state.enemyX - state.selfX, state.enemyY - state.selfY)
  let bearing = radToDeg(arctan2(state.enemyY - state.selfY, state.enemyX - state.selfX))

  let frame = tmEncodeFrame(
    bearing, dist, state.enemySpeed, state.enemyHeading,
    state.arenaHeight - state.enemyY, state.enemyY,
    state.arenaWidth - state.enemyX, state.enemyX,
    state.enemyEnergy,  # fix 6: enemy energy, not a duplicate of self energy
  )
  # Shift window: index 0 = newest. Do this at most once per tick — the harness
  # calls predict() 4-5x/tick (once per power bin), which used to shift the
  # 10-frame window ~4-5x/tick (representing ~2 real ticks and tripping
  # isWarmedUp after 2-3 ticks instead of 10).
  if state.tick != g.frameTick:
    g.frameTick = state.tick
    for i in countdown(TM_WINDOW_SIZE - 1, 1):
      g.frameBuffer[i] = g.frameBuffer[i - 1]
    g.frameBuffer[0] = frame
    if g.bufferCount < TM_WINDOW_SIZE: inc g.bufferCount

  # Warm-up: until window is full, fall back to linear extrapolation
  let ticksToArrive = if bulletSpeed > 0.0: dist / bulletSpeed else: 1.0
  let headingRad = degToRad(state.enemyHeading)
  let linearX = state.enemyX + cos(headingRad) * state.enemySpeed * ticksToArrive
  let linearY = state.enemyY + sin(headingRad) * state.enemySpeed * ticksToArrive

  if g.bufferCount < TM_WINDOW_SIZE:
    return GunPrediction(x: clamp(linearX, 0.0, state.arenaWidth),
                         y: clamp(linearY, 0.0, state.arenaHeight))

  let selfState = tmEncodeSelf(
    state.arenaHeight - state.selfY, state.selfY,
    state.arenaWidth - state.selfX, state.selfX,
    state.selfEnergy,
    true,  # canFire not in WorldState; assume true
  )
  let vec = tmEncodeFullVector(g.frameBuffer, selfState)

  var cache: TmClauseCache
  let (cx, cy, vx, vy) = tmForwardWithCache(g.net, vec, cache)
  inc g.predictCalls
  g.lastCx = cx
  g.lastCy = cy
  if abs(cx) + abs(cy) > 1e-9: inc g.correctionsNonzero

  let predX = clamp(linearX + cx, 0.0, state.arenaWidth)
  let predY = clamp(linearY + cy, 0.0, state.arenaHeight)

  # Store trace keyed exactly by (fireTick, powerBin) so the resolution event
  # can find it no matter how many other guns/bins fired in between.
  let binIdx = tmBinForSpeed(bulletSpeed)
  if binIdx >= 0:
    let slot = tmTraceSlot(state.tick, binIdx)
    g.traces[slot] = TmTrace(
      fireTick: state.tick,
      powerBin: binIdx,
      predX:    predX,
      predY:    predY,
      linearX:  linearX,
      linearY:  linearY,
      voteX:    vx,
      voteY:    vy,
      input:    vec,
      cache:    cache,
      alive:    true,
    )

  GunPrediction(x: predX, y: predY)

proc onResult*(g: var TsetlinGun, e: FeedbackEvent) =
  inc g.shotCount
  # Exact pairing: index the trace by the tick the bullet was fired and the power
  # bin it belonged to. The old coordinate-matched 64-slot ring lost the trace
  # long before a long shot resolved, so the TM never trained and its output was
  # pure linear extrapolation.
  let binIdx = if e.powerBin >= 0 and e.powerBin < len(vb.PowerBins): e.powerBin
               else: tmBinForSpeed(bulletSpeed(e.bulletPower))
  if binIdx < 0:
    inc g.traceMisses
    return
  let slot = tmTraceSlot(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

  # Fixes 4+5: train on delta = actual - LINEAR baseline. The old code trained
  # on actual - predX = delta - cx, so tmLearnOne's error became delta - 2*cx and
  # the fixed point was cx = delta/2. Hits are NOT zeroed: a hit means
  # |miss| < BotRadius, not residual == 0.
  let dx = clamp(e.actualX - t.linearX, -TM_RESID_MAX, TM_RESID_MAX)
  let dy = clamp(e.actualY - t.linearY, -TM_RESID_MAX, TM_RESID_MAX)
  let lits = tmMakeLiterals(t.input)
  g.net.tmLearnOne(0, lits, t.cache, t.voteX, dx)
  g.net.tmLearnOne(1, lits, t.cache, t.voteY, dy)
  when DebugTM:
    echo fmt"[tm-dbg] shot={g.shotCount} tick={e.fireTick} bin={binIdx} miss={e.missDistance:.1f}px predicted=({t.predX:.0f},{t.predY:.0f}) actual=({e.actualX:.0f},{e.actualY:.0f}) dx={dx:.1f} dy={dy:.1f} hit={e.hit}"
  t.alive = false
  inc g.trainedShots