# TsetlinBot — Regression Tsetlin Machine for aiming.
# Identical to BNNBot except uses tsetlin_predictor instead of hebbian.

import std/[math, os, strutils, random]
import robocode_tankroyale_botapi
import radar_lock/radar_lock as radar_lock
import binary_encoding
import tsetlin_predictor

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

const
  BULLET_SLOTS  = 50
  EPSILON_START = 0.2
  EPSILON_MIN   = 0.05
  EPSILON_DECAY = 0.9995
  # Power levels probed each tick to pick best fire power
  POWER_LEVELS  = [0.1'f64, 0.4, 0.7, 1.0, 1.3, 1.6, 1.9, 2.2, 2.5, 3.0]
  N_POWER       = POWER_LEVELS.len
  HIT_THRESHOLD = 0.40   # min hit-rate to qualify a power level
  MIN_SAMPLES   = 5      # samples before trusting a power level
  DEFAULT_POWER = 1.0    # cold-start fallback

type
  ProbeSlot = object   # lightweight wave probe — no learning, just geometry
    active:      bool
    powerIdx:    int
    fireX:       float
    fireY:       float
    aimAngleDeg: float
    fireDist:    float
    bulletSpeed: float
    age:         int

type
  TsetlinBot = ref object of Bot
    hasContact:   bool
    enemyBearing: float
    lastEnemyX:   float
    lastEnemyY:   float
    hasLastPos:   bool
    velocity:     float
    heading:      float
    distance:     float
    tick:         int
    prevVec:      BinaryVector
    hasPrev:      bool
    frameBuffer:  array[WINDOW_SIZE, array[FRAME_BITS, uint8]]
    bufferCount:  int
    net:          TsetlinNet
    bullets:      array[BULLET_SLOTS, VirtualBullet]
    bulletHead:   int
    virtualHits:  int
    virtualMiss:  int
    epsilon:      float
    aimAngle:     float   # absolute gun aim angle, updated each scan
    firePower:    float   # dynamically selected fire power
    powerStats:   array[N_POWER, tuple[hits, total: int]]
    probes:       array[N_POWER, ProbeSlot]  # one probe per power level per tick

method onScannedBot*(bot: TsetlinBot, e: ScannedBotEvent) =
  let bx = getX(); let by = getY()
  bot.enemyBearing = directionTo(bx, by, e.x, e.y)
  bot.distance     = distanceTo(bx, by, e.x, e.y)
  bot.heading      = e.direction
  bot.velocity     = e.speed
  bot.lastEnemyX   = e.x
  bot.lastEnemyY   = e.y
  bot.hasLastPos   = true
  bot.hasContact   = true

  let arenaW = getArenaWidth().float
  let arenaH = getArenaHeight().float

  let frame = EnemyScanFrame(
    bearing:     bot.enemyBearing,
    distance:    bot.distance,
    velocity:    bot.velocity,
    heading:     bot.heading,
    enemyWallN:  arenaH - e.y,
    enemyWallS:  e.y,
    enemyWallE:  arenaW - e.x,
    enemyWallW:  e.x,
    enemyEnergy: e.energy,
  )
  let encoded = encodeFrame(frame)

  for i in countdown(WINDOW_SIZE - 1, 1):
    bot.frameBuffer[i] = bot.frameBuffer[i - 1]
  bot.frameBuffer[0] = encoded

  if bot.bufferCount < WINDOW_SIZE:
    inc bot.bufferCount
  if bot.bufferCount < WINDOW_SIZE:
    return

  let selfState = SelfState(
    myWallN:  arenaH - by,
    myWallS:  by,
    myWallE:  arenaW - bx,
    myWallW:  bx,
    myEnergy: getEnergy(),
    canFire:  getGunHeat() <= 0.0,
  )
  let selfEncoded = encodeSelf(selfState)
  let vec = encodeFullVector(bot.frameBuffer, selfEncoded)

  # ── pick best fire power ──────────────────────────────────────────────
  bot.firePower = DEFAULT_POWER
  var bestPow   = -1.0
  for pi in 0..<N_POWER:
    let s = bot.powerStats[pi]
    if s.total >= MIN_SAMPLES:
      let rate = s.hits.float / s.total.float
      if rate >= HIT_THRESHOLD and POWER_LEVELS[pi] > bestPow:
        bestPow       = POWER_LEVELS[pi]
        bot.firePower = POWER_LEVELS[pi]

  let selectedSpeed = 20.0 - 3.0 * bot.firePower

  # ── Bug 4 fix: forward pass returns (cx, cy) pixel corrections ────────
  var cache: ClauseCache
  let (cx, cy) = bot.net.forwardWithCache(vec, cache)

  # epsilon-greedy exploration: perturb the correction
  var corrX = cx
  var corrY = cy
  if rand(1.0) < bot.epsilon:
    corrX += rand(20.0) - 10.0
    corrY += rand(20.0) - 10.0
  bot.epsilon = max(EPSILON_MIN, bot.epsilon * EPSILON_DECAY)

  # Convert (cx, cy) correction to aim angle.
  # Linear extrapolation first, then TM residual correction on top.
  let bx2 = getX(); let by2 = getY()
  var aimAngle = bot.enemyBearing
  if bot.hasLastPos:
    let ticksToArrive = bot.distance / selectedSpeed
    let extrapolX = cos(degToRad(bot.heading)) * bot.velocity * ticksToArrive
    let extrapolY = sin(degToRad(bot.heading)) * bot.velocity * ticksToArrive
    let targetX = bot.lastEnemyX + extrapolX + corrX
    let targetY = bot.lastEnemyY + extrapolY + corrY
    aimAngle = directionTo(bx2, by2, targetX, targetY)

  var aimOffset = aimAngle - bot.enemyBearing
  while aimOffset >  180.0: aimOffset -= 360.0
  while aimOffset < -180.0: aimOffset += 360.0

  bot.aimAngle = bot.enemyBearing + aimOffset

  # ── store virtual bullet (with clause cache for eligibility) ──────────
  let slot = bot.bulletHead mod BULLET_SLOTS
  bot.bulletHead = slot + 1
  bot.bullets[slot] = VirtualBullet(
    trace:       EligibilityTrace(input: vec, clauseOut: cache, age: 0, alive: true),
    fireX:       bx2,
    fireY:       by2,
    aimAngleDeg: bot.enemyBearing + aimOffset,
    fireDist:    bot.distance,
    bulletSpeed: selectedSpeed,
    active:      true,
  )

  # ── spawn one probe per power level ──────────────────────────────────
  for pi in 0..<N_POWER:
    let spd = 20.0 - 3.0 * POWER_LEVELS[pi]
    let probeTicks   = bot.distance / spd
    let probeExtraX  = cos(degToRad(bot.heading)) * bot.velocity * probeTicks
    let probeExtraY  = sin(degToRad(bot.heading)) * bot.velocity * probeTicks
    let probeTargetX = bot.lastEnemyX + probeExtraX + corrX
    let probeTargetY = bot.lastEnemyY + probeExtraY + corrY
    let probeAimAngle = directionTo(bx2, by2, probeTargetX, probeTargetY)
    bot.probes[pi] = ProbeSlot(
      active:      true,
      powerIdx:    pi,
      fireX:       bx2,
      fireY:       by2,
      aimAngleDeg: probeAimAngle,
      fireDist:    bot.distance,
      bulletSpeed: spd,
      age:         0,
    )

  # ── stats & echo ─────────────────────────────────────────────────────
  let hamming    = if bot.hasPrev: hammingDistance(bot.prevVec, vec) else: 0
  let similarity = if bot.hasPrev: TOTAL_BITS - hamming else: 0
  let overlap    = if bot.hasPrev: popcount(bitwiseAnd(bot.prevVec, vec)) else: 0
  let totalVirtual = bot.virtualHits + bot.virtualMiss
  let hitPct = if totalVirtual > 0: bot.virtualHits.float / totalVirtual.float * 100.0 else: 0.0
  echo align($bot.tick, 4), " ",
       hamming, " ", similarity, " ", overlap, " ",
       formatFloat(aimOffset, ffDecimal, 2), " ",
       bot.virtualHits, " ", bot.virtualMiss, " ",
       formatFloat(hitPct, ffDecimal, 1), "% ",
       "pwr=", formatFloat(bot.firePower, ffDecimal, 1)

  bot.prevVec = vec
  bot.hasPrev = true

method onRoundStarted*(bot: TsetlinBot, e: RoundStartedEvent) =
  setAdjustGunForBodyTurn(true)
  setAdjustRadarForBodyTurn(true)
  setAdjustRadarForGunTurn(true)
  radar_lock.init()
  bot.hasContact  = false
  bot.hasLastPos  = false
  bot.tick        = 0
  bot.hasPrev     = false
  bot.bufferCount = 0
  for i in 0..<BULLET_SLOTS:
    bot.bullets[i].active = false
  for i in 0..<N_POWER:
    bot.probes[i].active = false
  bot.bulletHead = 0
  # net and epsilon persist across rounds

method onGameStarted*(bot: TsetlinBot, e: GameStartedEventForBot) =
  discard

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

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

    setRadarTurnRate(radar_lock.doRadar(getRadarDirection(), bot.enemyBearing))

    # ── age & settle virtual bullets (game ticks) ─────────────────────
    if bot.hasLastPos:
      for idx in 0..<BULLET_SLOTS:
        var b = addr bot.bullets[idx]
        if not b.active: continue
        inc b.trace.age
        let bulletDist = b.bulletSpeed * float(b.trace.age)
        if bulletDist >= b.fireDist or b.trace.age >= TRACE_MAX_AGE:
          let bulletX   = b.fireX + cos(degToRad(b.aimAngleDeg)) * bulletDist
          let bulletY   = b.fireY + sin(degToRad(b.aimAngleDeg)) * bulletDist
          let residualX = bot.lastEnemyX - bulletX
          let residualY = bot.lastEnemyY - bulletY
          let missDistance = hypot(residualX, residualY)
          bot.net.learn(b.trace, residualX, residualY)
          if missDistance < 36.0: inc bot.virtualHits
          else:                   inc bot.virtualMiss
          b.active = false

      for pi in 0..<N_POWER:
        var p = addr bot.probes[pi]
        if not p.active: continue
        inc p.age
        let pd = p.bulletSpeed * float(p.age)
        if pd >= p.fireDist or p.age >= TRACE_MAX_AGE:
          let px = p.fireX + cos(degToRad(p.aimAngleDeg)) * pd
          let py = p.fireY + sin(degToRad(p.aimAngleDeg)) * pd
          let md = hypot(px - bot.lastEnemyX, py - bot.lastEnemyY)
          inc bot.powerStats[pi].total
          if md < 36.0: inc bot.powerStats[pi].hits
          p.active = false

    # Aim gun and fire when ready
    let gunDir    = getGunDirection()
    let gunDelta  = normalizeRelativeAngle(bot.aimAngle - gunDir)
    setGunTurnRate(gunDelta.clamp(-20.0, 20.0))
    if getGunHeat() <= 0.0 and abs(gunDelta) < 2.0:
      discard setFire(bot.firePower)

    go()

when isMainModule:
  randomize()
  var bot = TsetlinBot(
    net:       initTsetlinNet(),
    epsilon:   EPSILON_START,
    firePower: DEFAULT_POWER,
  )
  start(bot, botJsonPath)