# BNNBot — Hebbian weight matrix with virtual bullet learning.
# 870-bit input → 7-bit aim angle output via forward pass.
# Learns from virtual bullets (no real firing) via three-factor Hebbian rule.
# CSV data collection: BNNBOT_CSV=1 writes per-round CSV to data/battle_{round}.csv

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

let csvEnabled = getEnv("BNNBOT_CSV", "0") == "1"

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

const
  BULLET_SLOTS  = 50
  BULLET_SPEED  = 14.0  # power 2
  EPSILON_START = 0.2
  EPSILON_MIN   = 0.05
  EPSILON_DECAY = 0.9995

# CSV: 9 raw decimal fields per frame (matches analysis/backtest.py column names)
# bearing_sin(0-199), bearing_cos(0-199), distance(0-99), velocity(0-15),
# heading_sin(0-199), heading_cos(0-199), enemy_x(px), enemy_y(px), enemy_energy(float)
type
  FrameRaw = object
    bSin, bCos:    int     # 0-199
    dist:          int     # 0-99
    vel:           int     # 0-15
    hSin, hCos:   int     # 0-199
    ex, ey:        float   # absolute pixel coords
    energy:        float

  BNNBot = 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]]
    frameRawBuf:  array[WINDOW_SIZE, FrameRaw]   # decimal mirror of frameBuffer
    bufferCount:  int
    net:          HebbianNet
    bullets:      array[BULLET_SLOTS, VirtualBullet]
    bulletHead:   int
    virtualHits:  int
    virtualMiss:  int
    epsilon:      float
    # CSV state
    csvFile:      File
    csvOpen:      bool
    roundNum:     int

# ── CSV helpers ──────────────────────────────────────────────────────────────

proc csvPath(roundNum: int): string =
  getAppDir() / "data" / fmt"battle_{roundNum}.csv"

proc buildHeader(): string =
  result = "tick"
  for fr in 0..<WINDOW_SIZE:
    for fn in ["bearing_sin", "bearing_cos", "distance", "velocity",
               "heading_sin", "heading_cos", "enemy_x", "enemy_y", "enemy_energy"]:
      result &= fmt",f{fr}_{fn}"

proc openCsv(bot: BNNBot) =
  if not csvEnabled: return
  createDir(csvPath(bot.roundNum).parentDir)
  bot.csvFile = open(csvPath(bot.roundNum), fmWrite)
  bot.csvOpen = true
  bot.csvFile.writeLine(buildHeader())

proc closeCsv(bot: BNNBot) =
  if bot.csvOpen:
    bot.csvFile.close()
    bot.csvOpen = false

proc writeRow(bot: BNNBot) =
  if not bot.csvOpen: return
  var line = $bot.tick
  for i in 0..<WINDOW_SIZE:
    let f = bot.frameRawBuf[i]
    line &= fmt",{f.bSin},{f.bCos},{f.dist},{f.vel},{f.hSin},{f.hCos},{f.ex:.2f},{f.ey:.2f},{f.energy:.2f}"
  bot.csvFile.writeLine(line)

# ── Bot methods ──────────────────────────────────────────────────────────────

method onScannedBot*(bot: BNNBot, 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)

  # Shift window — index 0 = newest
  for i in countdown(WINDOW_SIZE - 1, 1):
    bot.frameBuffer[i] = bot.frameBuffer[i - 1]
    bot.frameRawBuf[i] = bot.frameRawBuf[i - 1]
  bot.frameBuffer[0] = encoded

  # Build raw decimal record for this frame
  let bSin = int((sin(degToRad(bot.enemyBearing)) + 1.0) / 2.0 * 199.0)
  let bCos = int((cos(degToRad(bot.enemyBearing)) + 1.0) / 2.0 * 199.0)
  let dist = int(clamp(bot.distance / 1414.0 * 99.0, 0.0, 99.0))
  let vel  = int(clamp(bot.velocity + 8.0, 0.0, 16.0))
  let hSin = int((sin(degToRad(bot.heading)) + 1.0) / 2.0 * 199.0)
  let hCos = int((cos(degToRad(bot.heading)) + 1.0) / 2.0 * 199.0)
  bot.frameRawBuf[0] = FrameRaw(bSin: bSin, bCos: bCos, dist: dist, vel: vel,
                                  hSin: hSin, hCos: hCos,
                                  ex: e.x, ey: e.y, energy: e.energy)

  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)

  # ── age & settle existing virtual bullets ────────────────────────────
  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 missDistance = hypot(bulletX - bot.lastEnemyX, bulletY - bot.lastEnemyY)

      # Shaped reward: +1.0 for perfect hit, decays toward -1.0 as miss distance grows
      let reward = 2.0 * exp(-missDistance / 36.0) - 1.0
      bot.net.learn(b.trace, reward)

      if missDistance < 36.0:
        inc bot.virtualHits
      else:
        inc bot.virtualMiss
      b.active = false

  # ── forward pass → aim angle ─────────────────────────────────────────
  var outVec = bot.net.forward(vec)

  # epsilon-greedy exploration
  if rand(1.0) < bot.epsilon:
    let flipIdx = rand(OUTPUT_BITS - 1)
    outVec[flipIdx] = 1'u8 - outVec[flipIdx]
  bot.epsilon = max(EPSILON_MIN, bot.epsilon * EPSILON_DECAY)

  let aimOffset = decodeOutput(outVec)

  # ── store virtual bullet ──────────────────────────────────────────────
  let slot = bot.bulletHead mod BULLET_SLOTS
  bot.bulletHead = slot + 1
  bot.bullets[slot] = VirtualBullet(
    trace:       EligibilityTrace(input: vec, output: outVec, age: 0, alive: true),
    fireX:       bx,
    fireY:       by,
    aimAngleDeg: bot.enemyBearing + aimOffset,
    fireDist:    bot.distance,
    bulletSpeed: BULLET_SPEED,
    active:      true,
  )

  # ── CSV row ───────────────────────────────────────────────────────────
  bot.writeRow()

  # ── 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), "%"

  bot.prevVec = vec
  bot.hasPrev = true

method onRoundStarted*(bot: BNNBot, 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
  inc bot.roundNum
  bot.openCsv()
  # net and epsilon persist across rounds (learning carries over)

method onRoundEnded*(bot: BNNBot, e: RoundEndedEventForBot) =
  bot.closeCsv()

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

method run*(bot: BNNBot) =
  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))
    go()

when isMainModule:
  randomize()
  var bot = BNNBot(
    net:     initHebbianNet(),
    epsilon: EPSILON_START,
  )
  start(bot, botJsonPath)