## WiSARD regression predictor — K=14, ~50 neurons, 16384 entries each.
## Input: TOTAL_BITS-bit BinaryVector. Output: (dx, dy) correction.
## Online learning via eligibility traces (addresses stored per pending wave).

import std/[math, random]
import binary_encoding

const
  K*         = 14
  N_BITS_PAD = ((TOTAL_BITS + K - 1) div K) * K   # pad to multiple of K
  N_NEURONS* = N_BITS_PAD div K                    # ceil(690/14) = 50
  LAST_BITS  = if TOTAL_BITS mod K == 0: K else: TOTAL_BITS mod K  # bits in last neuron
  LUT_SIZE   = 1 shl K                              # 16384
  LAST_LUT_SIZE = 1 shl LAST_BITS                   # smaller LUT for last neuron
  BLEACH_THRESHOLD* = 1  # only use entries with count > this

type
  LutEntry = object
    sumX:  float
    sumY:  float
    count: int

  WiSARDNet* = object
    perm*: array[N_BITS_PAD, int]
    luts*: array[N_NEURONS, array[LUT_SIZE, LutEntry]]

  # Eligibility trace: LUT addresses computed at fire time, replayed on resolution
  WaveTrace* = object
    addrs*: array[N_NEURONS, int]
    age*:   int
    valid*: bool

  # Virtual bullet slot (mirrors tsetlin_predictor's VirtualBullet)
  VirtualBullet* = object
    trace*:       WaveTrace
    fireX*:       float
    fireY*:       float
    aimAngleDeg*: float
    fireDist*:    float
    bulletSpeed*: float
    active*:      bool

const TRACE_MAX_AGE* = 40

proc initWiSARD*(seed: int64 = 42): WiSARDNet =
  ## Build fixed random permutation over TOTAL_BITS; padding slots duplicate
  ## valid indices (random) to avoid zero-bias.
  var rng = initRand(seed)
  for i in 0..<TOTAL_BITS: result.perm[i] = i
  # Fisher-Yates shuffle over real bits only
  for i in countdown(TOTAL_BITS - 1, 1):
    let j = rng.rand(i)
    swap(result.perm[i], result.perm[j])
  # Padding slots get random valid indices (no bias toward bit 0)
  for i in TOTAL_BITS..<N_BITS_PAD:
    result.perm[i] = rng.rand(TOTAL_BITS - 1)

proc computeAddresses*(ws: WiSARDNet, vec: BinaryVector): array[N_NEURONS, int] =
  for n in 0..<N_NEURONS:
    var laddr = 0
    let bits = if n == N_NEURONS - 1: LAST_BITS else: K
    for b in 0..<bits:
      laddr = (laddr shl 1) or int(vec[ws.perm[n * K + b]])
    result[n] = laddr

proc predictCorrection*(ws: WiSARDNet, addrs: array[N_NEURONS, int]): tuple[cx, cy: float] =
  ## Average (sumX/count, sumY/count) across neurons with count > BLEACH_THRESHOLD.
  var sx = 0.0; var sy = 0.0; var active = 0
  for n in 0..<N_NEURONS:
    let e = ws.luts[n][addrs[n]]
    if e.count > BLEACH_THRESHOLD:
      sx += e.sumX / float(e.count)
      sy += e.sumY / float(e.count)
      inc active
  if active == 0: return (0.0, 0.0)
  (sx / float(active), sy / float(active))

proc learnCorrection*(ws: var WiSARDNet, addrs: array[N_NEURONS, int],
                      dx, dy: float) =
  ## Accumulate (dx, dy) residuals at the addressed LUT entries.
  if dx.isNaN or dy.isNaN or dx.classify == fcInf or dx.classify == fcNegInf or
     dy.classify == fcInf or dy.classify == fcNegInf: return
  for n in 0..<N_NEURONS:
    let a = addrs[n]
    ws.luts[n][a].sumX  += dx
    ws.luts[n][a].sumY  += dy
    ws.luts[n][a].count += 1