# Regression Tsetlin Machine for aiming correction.
# Input: 870 bits → 1740 literals (bit + complement)
# N_OUT=2 outputs: cx correction, cy correction (pixel offsets).
# 64 clauses per output (32 positive polarity, 32 negative).
# Clause output = AND of all included literals.
# Regression: vote ∈ [-N_CLAUSES/2, N_CLAUSES/2] → scaled to [-RESID_MAX, RESID_MAX].
# Online learning via Type I / Ib / Type II stochastic feedback.

import binary_encoding
import std/[math, random]

const
  N_IN*        = TOTAL_BITS        # 870
  N_OUT*       = 2                 # cx, cy pixel corrections
  N_LITERALS   = N_IN * 2         # 1740 (bit + complement)
  N_CLAUSES    = 64               # per output; first 32 = pos polarity
  HALF_CLAUSES = N_CLAUSES div 2
  N_STATES     = 15               # automaton range [-N_STATES..N_STATES] (int8 ok)
  T*           = float(HALF_CLAUSES)  # = 32.0; vote clamped to [-T, T]
  S            = 4.0              # specificity (higher = sparser clauses)
  RESID_MAX*   = 80.0             # residual correction range (on top of linear extrapolation)
  # ponytail: N_STATES=15 fits int8, keeps array small; raise if underfitting

  TRACE_MAX_AGE* = 40

type
  # int8 automaton: positive = include literal, negative = exclude
  TsetlinNet* = object
    # states[out][clause][literal] — flattened
    states: array[N_OUT * N_CLAUSES * N_LITERALS, int8]

  # Eligibility: store clause outputs (one bit per clause per output)
  ClauseCache* = array[N_OUT * N_CLAUSES, uint8]

  EligibilityTrace* = object
    input*:      BinaryVector
    clauseOut*:  ClauseCache
    age*:        int
    alive*:      bool

  VirtualBullet* = object
    trace*:       EligibilityTrace
    fireX*:       float
    fireY*:       float
    aimAngleDeg*: float
    fireDist*:    float
    bulletSpeed*: float
    active*:      bool

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

proc stateIdx(outIdx, clause, lit: int): int {.inline.} =
  (outIdx * N_CLAUSES + clause) * N_LITERALS + lit

proc clausePolarity(clause: int): float {.inline.} =
  # first HALF_CLAUSES = +1 polarity, rest = -1
  if clause < HALF_CLAUSES: 1.0 else: -1.0

proc evalClause(net: TsetlinNet, outIdx, clause: int,
                literals: array[N_LITERALS, uint8]): uint8 =
  ## Returns 1 if clause fires (AND of all included literals satisfied).
  ## Empty clause (no included literals) returns 0 — silent, not vacuously true.
  var hasIncluded = false
  for lit in 0..<N_LITERALS:
    let s = net.states[stateIdx(outIdx, clause, lit)]
    if s > 0:
      hasIncluded = true
      if literals[lit] == 0:
        return 0'u8
  return if hasIncluded: 1'u8 else: 0'u8

proc makeLiterals(input: BinaryVector): array[N_LITERALS, uint8] =
  for i in 0..<N_IN:
    result[i]        = input[i]
    result[i + N_IN] = 1'u8 - input[i]

proc computeVote(net: TsetlinNet, outIdx: int,
                 literals: array[N_LITERALS, uint8]): float =
  for c in 0..<N_CLAUSES:
    result += clausePolarity(c) * float(evalClause(net, outIdx, c, literals))
  result = clamp(result, -T, T)

# ── public API ─────────────────────────────────────────────────────────

proc initTsetlinNet*(): TsetlinNet =
  # Bug 2 fix: init at 0 (boundary). One Type I step crosses into Include.
  for s in result.states.mitems:
    s = 0'i8

proc forward*(net: TsetlinNet, input: BinaryVector): (float, float) =
  ## Returns (cx, cy) pixel corrections.
  let literals = makeLiterals(input)
  let cx = computeVote(net, 0, literals) / T * RESID_MAX
  let cy = computeVote(net, 1, literals) / T * RESID_MAX
  return (cx, cy)

proc forwardWithCache*(net: TsetlinNet, input: BinaryVector,
                       cache: var ClauseCache): (float, float) =
  ## Like forward but also populates cache for eligibility tracing.
  let literals = makeLiterals(input)
  var vx = 0.0
  var vy = 0.0
  for c in 0..<N_CLAUSES:
    let o = evalClause(net, 0, c, literals)
    cache[c] = o
    vx += clausePolarity(c) * float(o)
  for c in 0..<N_CLAUSES:
    let o = evalClause(net, 1, c, literals)
    cache[N_CLAUSES + c] = o
    vy += clausePolarity(c) * float(o)
  vx = clamp(vx, -T, T)
  vy = clamp(vy, -T, T)
  return (vx / T * RESID_MAX, vy / T * RESID_MAX)

proc learnOne(net: var TsetlinNet, outIdx: int, literals: array[N_LITERALS, uint8],
              clauseOut: ClauseCache, residual: float) =
  ## Regression TM update for one output dimension.
  # Recompute predicted from cached clause outputs
  var vote = 0.0
  for c in 0..<N_CLAUSES:
    vote += clausePolarity(c) * float(clauseOut[outIdx * N_CLAUSES + c])
  vote = clamp(vote, -T, T)
  let predicted = vote / T * RESID_MAX
  let error = residual - predicted
  # Bug 1 fix: proper feedback probability gated on normalized error
  let pFeedback = min(1.0, abs(error) / (2.0 * RESID_MAX))

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

    if (error > 0.0 and pol > 0.0) or (error < 0.0 and pol < 0.0):
      # Type I feedback: grow clause toward current input
      if cOut == 1'u8:
        # Type Ia: clause fires — reinforce matching features
        for lit in 0..<N_LITERALS:
          let si  = stateIdx(outIdx, c, lit)
          var st  = int(net.states[si])
          if literals[lit] == 1'u8:
            if rand(1.0) < (S - 1.0) / S:
              st = min(st + 1, N_STATES)
          else:
            if rand(1.0) < 1.0 / S:
              st = max(st - 1, -N_STATES)
          net.states[si] = int8(st)
      else:
        # Type Ib: clause silent, should fire — grow toward current input
        for lit in 0..<N_LITERALS:
          let si = stateIdx(outIdx, c, lit)
          var st = int(net.states[si])
          if literals[lit] == 1'u8:
            if rand(1.0) < (S - 1.0) / S:
              st = min(st + 1, N_STATES)
          else:
            if rand(1.0) < 1.0 / S:
              st = max(st - 1, -N_STATES)
          net.states[si] = int8(st)
    else:
      # Bug 3 fix: Type II — decrement (toward exclude) false literals in Include range
      if cOut == 1'u8:
        for lit in 0..<N_LITERALS:
          if literals[lit] == 0'u8:
            let si = stateIdx(outIdx, c, lit)
            var st = int(net.states[si])
            if st > 0:  # only if currently in Include range
              st = max(st - 1, -N_STATES)
              net.states[si] = int8(st)

proc learn*(net: var TsetlinNet, trace: EligibilityTrace,
            residualX: float, residualY: float) =
  ## Online TM update from a resolved virtual bullet.
  ## residualX/Y: pixel correction needed (actual_target - aimed_point).
  if not trace.alive: return
  let literals = makeLiterals(trace.input)
  net.learnOne(0, literals, trace.clauseOut, residualX)
  net.learnOne(1, literals, trace.clauseOut, residualY)