## PPO_Bot — enemy tracker + state vector wired into the game loop.
## Training: trajectory collected per tick, PPO update in background thread.

import std/[os, locks]
import arraymancer
import tankroyale_botapi
import network
import actions
import training
import weights
import ./enemy_tracker
import ./state_vector

const botJsonPath = currentSourcePath().parentDir / "PPO_Bot.json"
const weightsRoot = currentSourcePath().parentDir / "weights"

type PPOBot = ref object of Bot
  tracker:    EnemyTracker
  buffer:     TrajectoryBuffer
  prevEnergy: float32       # own energy last tick
  prevEnemyE: float32       # enemy energy last tick (from tracker)
  lastState:  Tensor[float32]
  lastAction: Tensor[float32]
  lastLogP:   float32
  lastValue:  float32
  hasLastTrans: bool

var ac = initActorCritic()

# ── Background training state ─────────────────────────────────────────────────

type TrainingArgs = object
  ac:          ActorCritic
  buffer:      TrajectoryBuffer
  lastValue:   float32
  roundNum:    int
  weightsRoot: string

var
  trainingThread: Thread[TrainingArgs]
  trainingDone:   bool = true    # true = no training running / last run finished
  trainingLock:   Lock
  resultChan:     Channel[ActorCritic]
  roundCounter:   int = 0

proc trainingThreadProc(args: TrainingArgs) {.thread.} =
  var localAc = args.ac
  ppoUpdate(localAc, args.buffer, lastValue = args.lastValue)
  saveCheckpoint(localAc, args.weightsRoot, args.roundNum)
  resultChan.send(localAc)
  withLock(trainingLock):
    trainingDone = true

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

method onScannedBot*(bot: PPOBot, e: ScannedBotEvent) =
  bot.tracker.update(e.x, e.y, e.direction, e.speed, e.energy)

method onRoundStarted*(bot: PPOBot, e: RoundStartedEvent) =
  bot.tracker       = initEnemyTracker()
  bot.buffer        = initTrajectoryBuffer()
  bot.prevEnergy    = 0.0'f32
  bot.prevEnemyE    = 0.0'f32
  bot.hasLastTrans  = false

method onRoundEnded*(bot: PPOBot, e: RoundEndedEventForBot) =
  inc roundCounter

  # Add round-end score bonus to last transition (if any)
  let roundReward = computeRoundReward(e.results.totalScore.float32)
  if bot.hasLastTrans and bot.buffer.len > 0:
    bot.buffer.transitions[^1].reward += roundReward

  # Pick up result from previous training thread if done
  withLock(trainingLock):
    if trainingDone and roundCounter > 1:
      let (avail, trained) = resultChan.tryRecv()
      if avail:
        ac = trained

  if bot.buffer.len == 0:
    bot.hasLastTrans = false
    return

  # If last thread still running, drop this pass — start fresh with newer data
  # ponytail: simple drop; queue if every round must train
  withLock(trainingLock):
    if not trainingDone:
      bot.buffer.clear()
      bot.hasLastTrans = false
      return
    trainingDone = false

  let args = TrainingArgs(
    ac:          ac,
    buffer:      bot.buffer,
    lastValue:   0.0'f32,
    roundNum:    roundCounter,
    weightsRoot: weightsRoot,
  )
  bot.buffer.clear()
  bot.hasLastTrans = false

  createThread(trainingThread, trainingThreadProc, args)

method run(bot: PPOBot) =
  # Seed energy on first tick
  bot.prevEnergy = getEnergy().float32
  bot.prevEnemyE = if bot.tracker.hasContact: bot.tracker.current.energy.float32 else: 0.0'f32

  while isRunning():
    bot.tracker.deadReckon()

    setRadarTurnRate(bot.tracker.getRadarTurnRate(
      getX(), getY(), getDirection(), getRadarDirection()))

    let botData = BotStateData(
      x:           getX(),
      y:           getY(),
      direction:   getDirection(),
      speed:       getSpeed(),
      energy:      getEnergy(),
      gunDirection: getGunDirection(),
      gunHeat:     getGunHeat(),
      arenaWidth:  float64(getArenaWidth()),
      arenaHeight: float64(getArenaHeight()),
    )

    let state = buildStateVector(botData, bot.tracker)
    let (rawActs, logP) = ac.actorForward(state)
    let value = ac.criticForward(state)
    let acts = mapActions(rawActs, getSpeed().float32, getGunHeat().float32)

    # Compute tick reward from energy deltas
    let curEnergy   = getEnergy().float32
    let curEnemyE   = if bot.tracker.hasContact: bot.tracker.current.energy.float32 else: bot.prevEnemyE
    let myDelta     = curEnergy   - bot.prevEnergy
    let enemyDelta  = curEnemyE   - bot.prevEnemyE
    let tickReward  = computeTickReward(myDelta, enemyDelta)

    # Finalise previous transition with the reward from this tick's state change
    if bot.hasLastTrans:
      let tr = Transition(
        state:   bot.lastState,
        action:  bot.lastAction,
        logProb: bot.lastLogP,
        reward:  tickReward,
        value:   bot.lastValue,
      )
      bot.buffer.add(tr)

    # Store current for next tick
    bot.lastState   = state
    bot.lastAction  = rawActs
    bot.lastLogP    = logP
    bot.lastValue   = value
    bot.prevEnergy  = curEnergy
    bot.prevEnemyE  = curEnemyE
    bot.hasLastTrans = true

    setTargetSpeed(acts.targetSpeed.float)
    setTurnRate(acts.turnRate.float)
    setGunTurnRate(acts.gunTurnRate.float)
    if acts.shouldFire:
      discard setFire(acts.firePower.float)
    go()

when isMainModule:
  initLock(trainingLock)
  resultChan.open()
  createDir(weightsRoot)
  cleanStaleTempDirs(weightsRoot)
  discard loadBestAvailable(ac, weightsRoot)

  var bot = PPOBot(
    tracker: initEnemyTracker(),
    buffer:  initTrajectoryBuffer(),
  )
  start(bot, botJsonPath)