# SNNBot — SNN core + aiming loop prototype (issue #152).
# 80-input population-coded → 12 LIF hidden → polar-coded (sin/cos) output decoded via atan2.
# Inputs: bearing [0-35], velocity direction [36-71], speed [72-79].
# No movement, no firing. SuperSpike three-factor rule on all weights (issue #158).
# Binary reservoir alternative toggled via USE_RESERVOIR const (issue #159).
#
# State machine:
#   DECIDE   → feed bearing+velocity into SNN/reservoir, store targetAngle, → WAITING
#   WAITING  → aimTo() each tick; when error < 2° → EVALUATE
#   EVALUATE → measure error, compute SuperSpike/reservoir update, log, → DECIDE

import std/[math, random, os, strutils, bitops]
import robocode_tankroyale_botapi
import radar_lock/radar_lock as radar_lock
import reservoir

const USE_RESERVOIR* = true  # ponytail: kept for SNN path fallback; remove when binary aimer is validated

# ── Constants ──────────────────────────────────────────────────────────────────

const botJsonPath = currentSourcePath().parentDir / "SNNBot.json"

const
  N_IN          = 80          # input neurons: bearing(36) + vel_dir(36) + speed(8)
  N_HID         = 12          # hidden LIF neurons
  BAND_DEG      = 10.0        # degrees per input band
  BEARING_OFFSET  = 0         # neurons 0-35: relative bearing
  VEL_DIR_OFFSET  = 36        # neurons 36-71: velocity direction
  SPEED_OFFSET    = 72        # neurons 72-79: speed bands
  N_SPEED_BANDS   = 8
  SPEED_BAND_WIDTH = 1.0      # units/tick per band
  MAX_SPEED       = 8.0
  LEAK     = 0.9         # LIF membrane leak factor
  THRESH   = 0.08        # ponytail: THRESH=0.08 tuned for 80-input layer; raise if firing saturates
  MAX_GUN_TURN = 20.0    # max gun turn per tick (degrees)
  ETA      = 0.1         # SuperSpike learning rate for hidden→output weights (r_0 from paper, bumped 2x)
  ETA_IH   = 0.02        # ponytail: ETA_IH=0.02 bumped 4x for input→hidden (was 0.005, too conservative)
  # ponytail: separate input→hidden rate; add RMaxProp optimizer if convergence still unstable
  N_INFER  = 10          # inference window ticks per DECIDE
  # ponytail: N_INFER=10, increase if output still noisy; decrease if too slow per tick
  TRACE_DECAY = 0.9      # pre-synaptic trace decay (exponential low-pass)
  W_CLAMP  = 1.0         # ponytail: W_CLAMP=1.0, paper says ±0.1 but that's for mV-scale voltages; our unitless THRESH=1.0 needs larger weights
  BETA     = 1.0         # surrogate steepness (unitless; potentials are unitless)
  NOISE_AMP = 0.3        # exploration noise amplitude
  # ponytail: uniform exploration noise; upgrade to annealed Gaussian if convergence needs tuning
  FIRE_POWER  = 1.0      # fixed firing power; bullet speed = 20 - 3*FIRE_POWER
  BULLET_SPEED = 20.0 - 3.0 * FIRE_POWER

# ── SNN types ─────────────────────────────────────────────────────────────────

type
  SNN = object
    wih:  array[N_IN * N_HID, float]   # 80×12 input→hidden weights
    wSin: array[N_HID, float]          # hidden→sin-channel weights
    wCos: array[N_HID, float]          # hidden→cos-channel weights
    vHid: array[N_HID, float]          # hidden membrane potentials
    preTrace: array[N_IN, float]       # pre-synaptic traces (low-pass of input spikes)
    # Random feedback weights B[h, o] for hidden error projection; fixed, never updated.
    # Indexed as bFb[h * 2 + o], o=0 → sin, o=1 → cos.
    bFb: array[N_HID * 2, float]
    tick: int
    lastSinOut: float                  # last forward pass sin output (weighted sum)
    lastCosOut: float                  # last forward pass cos output (weighted sum)
    lastSnnAngle: float                # last atan2 result (degrees)

proc initSNN(snn: var SNN) =
  randomize()
  for w in snn.wih.mitems:  w = rand(0.2) - 0.1   # init within clamp ±0.1, room to grow to ±1.0
  for w in snn.wSin.mitems: w = rand(0.2) - 0.1
  for w in snn.wCos.mitems: w = rand(0.2) - 0.1
  for b in snn.bFb.mitems:  b = rand(2.0) - 1.0   # N(0,1)-ish; fixed forever
  for t in snn.preTrace.mitems: t = 0.0
  snn.tick = 0

proc encodeBearing(inputs: var array[N_IN, float], bearing: float, offset: int) =
  ## Population-code a -180..+180 angle into 36 neurons starting at offset.
  ## Triangular interpolation between lo and hi band.
  let norm = ((bearing + 180.0) / BAND_DEG)  # 0..36
  let lo   = int(norm) mod 36
  let hi   = (lo + 1) mod 36
  let frac = norm - float(int(norm))
  inputs[offset + lo] = 1.0 - frac
  inputs[offset + hi] = frac

proc encodeInput(bearing: float): array[N_IN, float] =
  ## Bearing-only encode for overlay (velocity channels stay 0).
  encodeBearing(result, bearing, BEARING_OFFSET)

proc encodeInputFull(bearing: float, velDirDeg: float, speed: float,
                     hasVel: bool): array[N_IN, float] =
  ## Full 80-neuron encode: bearing + velocity direction + speed.
  encodeBearing(result, bearing, BEARING_OFFSET)
  if hasVel:
    encodeBearing(result, velDirDeg, VEL_DIR_OFFSET)
    # Speed: triangular interpolation over N_SPEED_BANDS bands, clamp to [0, MAX_SPEED]
    let s    = speed.clamp(0.0, MAX_SPEED)
    let norm = s / SPEED_BAND_WIDTH
    let lo   = min(int(norm), N_SPEED_BANDS - 1)
    let hi   = min(lo + 1, N_SPEED_BANDS - 1)
    let frac = norm - float(int(norm))
    result[SPEED_OFFSET + lo] += 1.0 - frac
    result[SPEED_OFFSET + hi] += frac

proc surrogateDerivative(v: float): float {.inline.} =
  ## σ'(U) = (1 + |β(U − ϑ)|)^{-2}  — peaks at threshold, gives gradient direction.
  let x = BETA * (v - THRESH)
  result = 1.0 / ((1.0 + abs(x)) * (1.0 + abs(x)))

proc forward(snn: var SNN, inputs: array[N_IN, float],
             spikesOut: var array[N_HID, float],
             vHidSnap: var array[N_HID, float],
             sinOut: var float, cosOut: var float) =
  ## One SNN tick. Writes hidden spikes, pre-spike voltages, and raw sin/cos outputs.
  ## Caller accumulates sinOut/cosOut across N_INFER ticks, then calls atan2.
  inc snn.tick

  # Pre-synaptic trace: low-pass of input spikes
  for i in 0 ..< N_IN:
    snn.preTrace[i] = TRACE_DECAY * snn.preTrace[i] + inputs[i]

  # Hidden layer: LIF update
  for h in 0 ..< N_HID:
    var wsum = 0.0
    for i in 0 ..< N_IN:
      wsum += inputs[i] * snn.wih[i * N_HID + h]
    # ponytail: noise removed — SuperSpike surrogate derivative provides gradient direction, no exploration needed
    snn.vHid[h] = LEAK * snn.vHid[h] + wsum
    vHidSnap[h] = snn.vHid[h]   # snapshot voltage before reset (for surrogate)
    if snn.vHid[h] >= THRESH:
      spikesOut[h] = 1.0
      snn.vHid[h] = 0.0
    else:
      spikesOut[h] = 0.0

  # Output: polar-coded via sin/cos channels (raw; caller does atan2)
  sinOut = 0.0
  cosOut = 0.0
  for h in 0 ..< N_HID:
    sinOut += spikesOut[h] * snn.wSin[h]
    cosOut += spikesOut[h] * snn.wCos[h]
  snn.lastSinOut = sinOut
  snn.lastCosOut = cosOut
  snn.lastSnnAngle = arctan2(sinOut, cosOut) * 180.0 / PI

proc superSpikeUpdate(snn: var SNN,
                      spikes: array[N_HID, float],
                      vSnap: array[N_HID, float],
                      preTrace: array[N_IN, float],
                      targetAngle: float) =
  ## SuperSpike three-factor weight update.
  ## Δw = η × pre_trace × σ'(U) × error
  ## spikes/vSnap/preTrace: captured at DECIDE time for this inference window.
  ## Output error: target_rate − actual_rate (rate-coded target).
  ## Hidden error: projected via fixed random feedback weights B.

  # Target rates for sin/cos channels: map [-1,1] → [0,1]
  let tSin = (sin(degToRad(targetAngle)) + 1.0) / 2.0
  let tCos = (cos(degToRad(targetAngle)) + 1.0) / 2.0

  # Normalize accumulated spike counts to rates in [0,1]
  # Actual output activity using spike rates
  var sinAct = 0.0; var cosAct = 0.0
  for h in 0 ..< N_HID:
    let rate = spikes[h] / float(N_INFER)
    sinAct += rate * snn.wSin[h]
    cosAct += rate * snn.wCos[h]
  # Map actual output to [0,1] for rate comparison
  let sinActNorm = (sinAct.clamp(-1.0, 1.0) + 1.0) / 2.0
  let cosActNorm = (cosAct.clamp(-1.0, 1.0) + 1.0) / 2.0

  # Output error (target_rate − actual_rate)
  let errSin = tSin - sinActNorm
  let errCos = tCos - cosActNorm

  # Update hidden→output weights: Δw = η × rate_h × σ'(U_h) × error_o
  for h in 0 ..< N_HID:
    let sg   = surrogateDerivative(vSnap[h])
    let rate = spikes[h] / float(N_INFER)
    snn.wSin[h] += ETA * rate * sg * errSin
    snn.wSin[h] = snn.wSin[h].clamp(-W_CLAMP, W_CLAMP)
    snn.wCos[h] += ETA * rate * sg * errCos
    snn.wCos[h] = snn.wCos[h].clamp(-W_CLAMP, W_CLAMP)

  # Update input→hidden weights: Δw = η × preTrace_j × σ'(U_h) × error_h
  # Hidden error projected via fixed random feedback: error_h = Σ_o B[h,o] × error_o
  for h in 0 ..< N_HID:
    let sg     = surrogateDerivative(vSnap[h])
    let errHid = snn.bFb[h * 2 + 0] * errSin + snn.bFb[h * 2 + 1] * errCos
    for i in 0 ..< N_IN:
      snn.wih[i * N_HID + h] += ETA_IH * preTrace[i] * sg * errHid
      snn.wih[i * N_HID + h] = snn.wih[i * N_HID + h].clamp(-W_CLAMP, W_CLAMP)

# ── Bot state machine ─────────────────────────────────────────────────────────

type
  Phase = enum DECIDE, WAITING, EVALUATE

  SNNBot = ref object of Bot
    snn:             SNN
    res:             BinaryAimer
    phase:           Phase
    targetAngle:     float    # SNN output (absolute bearing)
    enemyBearing:    float    # last known enemy bearing
    enemyDist:       float    # last known distance to enemy
    hasContact:      bool
    tick:            int
    lastSpikes:      array[N_HID, float]  # accumulated spike counts over N_INFER ticks
    lastVSnap:       array[N_HID, float]  # hidden voltages before reset from last tick (for SuperSpike)
    lastRelBearing:  float              # relative bearing at DECIDE time (fixed for learning)
    lastEnemyX:      float64            # previous tick enemy position (for velocity)
    lastEnemyY:      float64
    hasLastPos:      bool
    velDirDeg:       float64            # velocity direction (degrees) from last scan delta
    velSpeed:        float64            # speed (units/tick) from last scan delta
    lastDecideGunDir: float             # gun heading captured at DECIDE time for EVALUATE
    lastBinInput:    BitVec80           # DECIDE-time binary input, reused in EVALUATE
    lastInput:       BitVec80           # previous EVALUATE-time input (for change detection)

# ── aimTo helper ──────────────────────────────────────────────────────────────

proc aimTo(targetAngle, gunDir: float) {.inline.} =
  let delta = normalizeRelativeAngle(targetAngle - gunDir)
  setGunTurnRate(delta.clamp(-MAX_GUN_TURN, MAX_GUN_TURN))

# ── Debug overlay ─────────────────────────────────────────────────────────────

proc drawOverlay(bot: SNNBot, myX, myY, gunDir, enemyBearing, targetAngle: float) =
  let LINE_LEN = bot.enemyDist + 50.0
  # Convert degree bearing to SVG direction (SVG y-axis is inverted; 0=east, CCW+)
  template toXY(bearing, len: float): (float, float) =
    (myX + cos(degToRad(bearing)) * len,
     myY + sin(degToRad(bearing)) * len)

  # Green: enemy direction (truth)
  setStrokeColor(GREEN)
  setStrokeWidth(2.0)
  let (ex, ey) = toXY(enemyBearing, LINE_LEN)
  drawLine(myX, myY, ex, ey)

  # Red: gun direction
  setStrokeColor(RED)
  let (gx, gy) = toXY(gunDir, LINE_LEN)
  drawLine(myX, myY, gx, gy)

  # Yellow: SNN target angle
  setStrokeColor(YELLOW)
  let (tx, ty) = toXY(targetAngle, LINE_LEN)
  drawLine(myX, myY, tx, ty)

# ── Event handlers ────────────────────────────────────────────────────────────

method onScannedBot*(bot: SNNBot, e: ScannedBotEvent) =
  let bx = getX(); let by = getY()
  bot.enemyBearing = directionTo(bx, by, e.x, e.y)
  bot.enemyDist    = distanceTo(bx, by, e.x, e.y)
  bot.hasContact   = true
  if bot.hasLastPos:
    let dx = e.x - bot.lastEnemyX
    let dy = e.y - bot.lastEnemyY
    bot.velDirDeg = arctan2(dy, dx) * 180.0 / PI
    bot.velSpeed  = sqrt(dx * dx + dy * dy)
  bot.lastEnemyX = e.x; bot.lastEnemyY = e.y
  bot.hasLastPos = true

method onRoundStarted*(bot: SNNBot, e: RoundStartedEvent) =
  setAdjustGunForBodyTurn(true)
  setAdjustRadarForBodyTurn(true)
  setAdjustRadarForGunTurn(true)
  radar_lock.init()
  bot.hasContact   = false
  bot.hasLastPos   = false
  bot.velDirDeg    = 0.0
  bot.velSpeed     = 0.0
  bot.phase        = DECIDE
  bot.tick         = 0
  setTargetSpeed(0.0)
  setTurnRate(0.0)

method onGameStarted*(bot: SNNBot, e: GameStartedEventForBot) =
  initSNN(bot.snn)
  bot.res = initBinaryAimer()

# ── Reservoir helpers ─────────────────────────────────────────────────────────

proc toBinaryInput(bearing: float, velDir: float, velSpeed: float, hasVel: bool): BitVec80 =
  ## Binary-native encoding: ~5 bits per channel, ~13 total active bits
  result = [0'u64, 0'u64]

  # Bearing: bits 0-35 (36 bits, 10° bands)
  # Activate center band + 2 neighbors on each side = 5 bits
  let bNorm = (bearing + 180.0) / 10.0  # 0..36
  let bCenter = int(bNorm) mod 36
  for offset in -2 .. 2:
    let idx = (bCenter + offset + 36) mod 36
    let word = idx div 64
    let bit = idx mod 64
    result[word] = result[word] or (1'u64 shl bit)

  if hasVel:
    # Velocity direction: bits 36-71 (36 bits, 10° bands)
    # Same scheme: 5 bits active
    let vNorm = (velDir + 180.0) / 10.0
    let vCenter = int(vNorm) mod 36
    for offset in -2 .. 2:
      let idx = 36 + (vCenter + offset + 36) mod 36
      let word = idx div 64
      let bit = idx mod 64
      result[word] = result[word] or (1'u64 shl bit)

    # Speed: bits 72-79 (8 bits, 1 unit/tick bands)
    # Activate center + 1 neighbor on each side = 3 bits
    let sCenter = clamp(int(velSpeed), 0, 7)
    for offset in -1 .. 1:
      let idx = 72 + clamp(sCenter + offset, 0, 7)
      let word = idx div 64
      let bit = idx mod 64
      result[word] = result[word] or (1'u64 shl bit)

# ── Main loop ─────────────────────────────────────────────────────────────────

method run*(bot: SNNBot) =
  while isRunning():
    inc bot.tick
    setTargetSpeed(0.0)
    setTurnRate(0.0)

    if not bot.hasContact:
      setRadarTurnRate(45.0)
      go()
      continue

    let myX    = getX()
    let myY    = getY()
    let gunDir = getGunDirection()

    case bot.phase
    of DECIDE:
      # absBearing: world-frame bearing [0,360) from directionTo — independent of bot heading
      let absBearing = bot.enemyBearing
      bot.lastRelBearing = normalizeRelativeAngle(absBearing - gunDir)
      when USE_RESERVOIR:
        let binInput = toBinaryInput(absBearing, bot.velDirDeg, bot.velSpeed, bot.hasLastPos)
        let aimRaw   = bot.res.forward(binInput)   # [0,360) or -1.0 sentinel
        bot.lastBinInput     = binInput
        bot.lastDecideGunDir = gunDir
        if aimRaw < 0.0:
          # cold start: aim directly at enemy
          bot.targetAngle = absBearing
          echo "RES tick=" & $bot.tick & " cold-start aim=" & formatFloat(absBearing, ffDecimal, 1)
        else:
          # aimRaw is absolute [0,360): use directly as target
          bot.targetAngle = aimRaw
          echo "RES tick=" & $bot.tick & " aim=" & formatFloat(aimRaw, ffDecimal, 1)
      else:
        let inputs = encodeInputFull(absBearing, bot.velDirDeg, bot.velSpeed, bot.hasLastPos)
        # Multi-tick inference: accumulate sin/cos and spike counts over N_INFER ticks
        var totalSin = 0.0; var totalCos = 0.0
        var tickSpikes: array[N_HID, float]
        var tickVSnap:  array[N_HID, float]
        for h in 0 ..< N_HID: bot.lastSpikes[h] = 0.0
        for _ in 0 ..< N_INFER:
          var sinT, cosT: float
          bot.snn.forward(inputs, tickSpikes, tickVSnap, sinT, cosT)
          totalSin += sinT; totalCos += cosT
          for h in 0 ..< N_HID:
            bot.lastSpikes[h] += tickSpikes[h]   # accumulate counts
        # Store last-tick voltages for learning
        bot.lastVSnap = tickVSnap
        bot.snn.lastSinOut = totalSin; bot.snn.lastCosOut = totalCos
        bot.snn.lastSnnAngle = arctan2(totalSin, totalCos) * 180.0 / PI
        bot.targetAngle = gunDir + bot.snn.lastSnnAngle
        bot.lastDecideGunDir = gunDir
        # Log total spike count across inference window
        var spikeCount = 0
        var maxV = 0.0
        for h in 0 ..< N_HID:
          spikeCount += int(bot.lastSpikes[h])
          maxV = max(maxV, bot.snn.vHid[h])
        echo "tick=" & $bot.tick & " infer_spk=" & $spikeCount & "/" & $(N_INFER * N_HID) & " maxV=" & formatFloat(maxV, ffDecimal, 3)
      bot.phase = WAITING

    of WAITING:
      aimTo(bot.targetAngle, gunDir)
      let err = abs(normalizeRelativeAngle(bot.targetAngle - gunDir))

      # Adaptive dead zone: scale learning threshold by enemy velocity
      # Stationary → 0.5° (stable), full speed (8 units/tick) → 0.3° (fast tracking)
      let t = (bot.velSpeed / 8.0).clamp(0.0, 1.0)
      let adaptiveDeadZone = 0.5 + (0.3 - 0.5) * t  # lerp(0.5, 0.3, t)

      # Fire gate: only fire when error within adaptive dead zone
      if err < adaptiveDeadZone:
        discard setFire(FIRE_POWER)
        bot.phase = EVALUATE
      else:
        discard setFire(0.0)

    of EVALUATE:
      # Predictive error signal: extrapolate enemy position at bullet impact time.
      if bot.hasLastPos:
        let travelTime = bot.enemyDist / BULLET_SPEED
        let velRad  = degToRad(bot.velDirDeg)
        let futureX = bot.lastEnemyX + cos(velRad) * bot.velSpeed * travelTime
        let futureY = bot.lastEnemyY + sin(velRad) * bot.velSpeed * travelTime
        # correctAngle: absolute world-frame bearing [0,360) to predicted enemy position
        let correctAngle = directionTo(myX, myY, futureX, futureY)

        # Adaptive learning dead zone: scale by enemy velocity (same as fire gate)
        let t = (bot.velSpeed / 8.0).clamp(0.0, 1.0)
        let adaptiveDeadZone = 0.5 + (0.3 - 0.5) * t  # lerp(0.5, 0.3, t)

        when USE_RESERVOIR:
          let aimRaw = bot.res.forward(bot.lastBinInput)
          # Detect significant input change via Hamming distance
          let hammingDist = popcount(bot.lastBinInput[0] xor bot.lastInput[0]) + popcount(bot.lastBinInput[1] xor bot.lastInput[1])
          let inputChanged = hammingDist > 2  # 3+ bits flipped = situation changed

          # Learn if: error > dead zone OR input pattern changed significantly
          if aimRaw >= 0.0:
            let aimErr = abs(normalizeRelativeAngle(aimRaw - correctAngle))
            if aimErr > adaptiveDeadZone or inputChanged:
              bot.res.learn(bot.lastBinInput, correctAngle)
          else:
            # Cold start: always learn
            bot.res.learn(bot.lastBinInput, correctAngle)
          let aimErr = if aimRaw >= 0.0: abs(normalizeRelativeAngle(aimRaw - correctAngle)) else: -1.0
          echo "RES tick=" & $bot.tick &
               " aim=" & (if aimRaw >= 0.0: formatFloat(aimRaw, ffDecimal, 1) else: "cold") &
               " correct=" & formatFloat(correctAngle, ffDecimal, 1) &
               " err=" & (if aimRaw >= 0.0: formatFloat(aimErr, ffDecimal, 1) else: "n/a") &
               " hammingDist=" & $hammingDist &
               " inputChanged=" & $inputChanged &
               " exemplars=" & $bot.res.count
        else:
          let targetRel = normalizeRelativeAngle(correctAngle - bot.lastDecideGunDir)
          bot.snn.superSpikeUpdate(bot.lastSpikes, bot.lastVSnap, bot.snn.preTrace, targetRel)
          let err = abs(normalizeRelativeAngle(gunDir - bot.enemyBearing))
          # Compute verbose logging metrics
          var spikeCount = 0
          var maxV = 0.0
          for h in 0 ..< N_HID:
            spikeCount += int(bot.lastSpikes[h])
            maxV = max(maxV, bot.snn.vHid[h])
          var meanWih = 0.0
          for w in bot.snn.wih:
            meanWih += abs(w)
          meanWih /= float(N_IN * N_HID)
          var meanWout = 0.0
          for w in bot.snn.wSin:
            meanWout += abs(w)
          for w in bot.snn.wCos:
            meanWout += abs(w)
          meanWout /= float(2 * N_HID)
          echo "tick=" & $bot.tick & " err=" & formatFloat(err, ffDecimal, 1) & "° infer_spk=" & $spikeCount & "/" & $(N_INFER * N_HID) & " maxV=" & formatFloat(maxV, ffDecimal, 3) & " |wih|=" & formatFloat(meanWih, ffDecimal, 4) & " |wOut|=" & formatFloat(meanWout, ffDecimal, 4) & " sin=" & formatFloat(bot.snn.lastSinOut, ffDecimal, 3) & " cos=" & formatFloat(bot.snn.lastCosOut, ffDecimal, 3) & " snnAngle=" & formatFloat(bot.snn.lastSnnAngle, ffDecimal, 1)

      # Update lastInput for next EVALUATE cycle's change detection
      bot.lastInput = bot.lastBinInput
      bot.phase = DECIDE

    # Radar lock
    setRadarTurnRate(radar_lock.doRadar(getRadarDirection(), bot.enemyBearing))

    # Debug overlay
    drawOverlay(bot, myX, myY, gunDir, bot.enemyBearing, bot.targetAngle)

    go()

# ── Entry point ───────────────────────────────────────────────────────────────

when isMainModule:
  var bot = SNNBot(phase: DECIDE)
  start(bot, botJsonPath)