## 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]
import gun_harness/gun_interface

# ── 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
  TM_TRACE_SLOTS = 64   # ring buffer of pending traces
  # ponytail: 64 slots >> TRACE_MAX_AGE=40 ticks, safe margin; grow if many guns/bins

type
  TmTrace = object
    predX, predY: float          # key: matches FeedbackEvent.prediction
    input:        TmBinaryVector
    cache:        TmClauseCache
    alive:        bool

  TsetlinGun* = object
    net:         TmNet
    frameBuffer: array[TM_WINDOW_SIZE, TmFrameEncoded]
    bufferCount: int
    traces:      array[TM_TRACE_SLOTS, TmTrace]
    traceHead:   int

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
  randomize()

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
  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 by prediction coords
  let slot = g.traceHead mod TM_TRACE_SLOTS
  g.traces[slot] = TmTrace(predX: predX, predY: predY, input: vec, cache: cache, alive: true)
  g.traceHead = (slot + 1) mod TM_TRACE_SLOTS

  GunPrediction(x: predX, y: predY)

proc onResult*(g: var TsetlinGun, e: FeedbackEvent) =
  # Find matching trace by prediction coords
  for i in 0..<TM_TRACE_SLOTS:
    var t = addr g.traces[i]
    if not t.alive: continue
    if abs(t.predX - e.prediction.x) > 0.01 or abs(t.predY - e.prediction.y) > 0.01:
      continue
    # Shaped reward: residual = (actual enemy pos) - (bullet impact pos)
    # FeedbackEvent carries missDistance but not the direction.
    # We reconstruct: aimed at (predX, predY); miss is distance to current enemy.
    # Use miss distance as magnitude; direction unknown → scale back along aim vector.
    # ponytail: zero-direction residual when miss=0 is fine; TM learns from magnitude via pFeedback
    let missSign = if e.hit: 0.0 else: 1.0
    let residualX = (e.prediction.x - t.predX) * missSign  # trivially 0; real signal is missDistance
    # Better: treat miss distance as residual magnitude along (enemy - pred) direction
    # We don't have enemy pos here directly, but we can scale correction proportionally.
    # Simplest correct signal: pass missDistance as residual magnitude for both dims.
    let rMag = e.missDistance * missSign
    let lits = tmMakeLiterals(t.input)
    # Apply residual equally to both axes (we don't know direction split)
    # ponytail: split 50/50; upgrade to directional when FeedbackEvent carries enemy pos
    let r = rMag / sqrt(2.0)
    g.net.tmLearnOne(0, lits, t.cache, r)
    g.net.tmLearnOne(1, lits, t.cache, r)
    t.alive = false
    break