## 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

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]): 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
  return if hasIncluded: 1'u8 else: 0'u8

proc tmForwardWithCache(net: TmNet, input: TmBinaryVector,
                        cache: var TmClauseCache): (float, float) =
  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] = o
    vx += tmPolarity(c) * float(o)
  for c in 0..<TM_N_CLAUSES:
    let o = tmEvalClause(net, 1, c, lits)
    cache[TM_N_CLAUSES + c] = o
    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)

proc tmLearnOne(net: var TmNet, outIdx: int, lits: array[TM_N_LITERALS, uint8],
                cache: TmClauseCache, residual: float) =
  var vote = 0.0
  for c in 0..<TM_N_CLAUSES:
    vote += tmPolarity(c) * float(cache[outIdx * TM_N_CLAUSES + c])
  vote = clamp(vote, -TM_T, TM_T)
  let predicted = vote / TM_T * TM_RESID_MAX
  let error = residual - predicted
  let pFeedback = min(1.0, abs(error) / (2.0 * TM_RESID_MAX))

  for c in 0..<TM_N_CLAUSES:
    if pFeedback <= 0.0: continue
    if rand(1.0) >= pFeedback: continue
    let pol  = tmPolarity(c)
    let cOut = cache[outIdx * TM_N_CLAUSES + c]

    if (error > 0.0 and pol > 0.0) or (error < 0.0 and pol < 0.0):
      # Type I / Ib feedback
      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 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)
        net.states[si] = int16(st)
    else:
      # Type II: shrink false literals in include range
      if cOut == 1'u8:
        for lit in 0..<TM_N_LITERALS:
          if lits[lit] == 0'u8:
            let si = tmStateIdx(outIdx, c, lit)
            var st = int(net.states[si])
            if st > 0:
              net.states[si] = int16(max(st - 1, -TM_N_STATES))

# ── 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 the directional residual
    input:        TmBinaryVector
    cache:        TmClauseCache
    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)
    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 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.selfEnergy,  # use self energy as proxy (enemy energy not in WorldState)
  )
  # 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) = tmForwardWithCache(g.net, vec, cache)

  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,
      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

  # Directional residual: actual enemy pos minus our prediction
  # On hit residual is 0 (we were right); on miss we push toward actual position.
  let rx = if e.hit: 0.0 else: clamp(e.actualX - t.predX, -TM_RESID_MAX, TM_RESID_MAX)
  let ry = if e.hit: 0.0 else: clamp(e.actualY - t.predY, -TM_RESID_MAX, TM_RESID_MAX)
  let lits = tmMakeLiterals(t.input)
  g.net.tmLearnOne(0, lits, t.cache, rx)
  g.net.tmLearnOne(1, lits, t.cache, ry)
  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}) rx={rx:.1f} ry={ry:.1f} hit={e.hit}"
  t.alive = false
  inc g.trainedShots