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

import std/[os, strformat, strutils, math, times, algorithm]
import arraymancer
import tankroyale_botapi
import PPO_Bot/network
import PPO_Bot/actions
import PPO_Bot/training
import PPO_Bot/weights
import PPO_Bot/enemy_tracker
import PPO_Bot/state_vector

# ── Hyperparameters from env vars (PPOB_ prefix) ─────────────────────────────
# All optional; defaults match ppoUpdate signature in training.nim.

proc getEnvFloat(name: string, default: float32): float32 =
  let v = getEnv(name)
  if v.len == 0: default else: parseFloat(v).float32

proc getEnvInt(name: string, default: int): int =
  let v = getEnv(name)
  if v.len == 0: default else: parseInt(v)

var
  hpLr:             float32 = getEnvFloat("PPOB_LR", 3e-4'f32)
  hpClipEpsilon:    float32 = getEnvFloat("PPOB_CLIP_EPSILON", 0.2'f32)
  hpEntropyCoeff:   float32 = getEnvFloat("PPOB_ENTROPY_COEFF", 0.01'f32)
  hpValueLossCoeff: float32 = getEnvFloat("PPOB_VALUE_LOSS_COEFF", 0.5'f32)
  hpMaxGradNorm:    float32 = getEnvFloat("PPOB_MAX_GRAD_NORM", 0.5'f32)
  hpGamma:          float32 = getEnvFloat("PPOB_GAMMA", 0.99'f32)
  hpLam:            float32 = getEnvFloat("PPOB_LAM", 0.95'f32)
  hpEpochs:         int     = getEnvInt("PPOB_EPOCHS", 4)
  hpMiniBatchSize:  int     = getEnvInt("PPOB_MINI_BATCH_SIZE", 64)
  hpUpdateInterval: int     = getEnvInt("PPOB_UPDATE_INTERVAL", 10)

# Wire logStd tunable params into network module vars (read before initActorCritic)
logStdFloor   = getEnvFloat("PPOB_LOG_STD_FLOOR", -3.0'f32)
logStdCeiling = getEnvFloat("PPOB_LOG_STD_CEILING", 0.5'f32)
initialLogStd = getEnvFloat("PPOB_INITIAL_LOG_STD", 0.0'f32)

# ── Structured log output ─────────────────────────────────────────────────────

let logFile = getEnv("PPOB_LOG_FILE")  # empty → no JSON logging
let evalOnly = getEnv("PPOB_EVAL_ONLY") == "1"  # freeze training (pure evaluation)

proc appendJsonLine(path, line: string) =
  ## Append a JSON line to path; no-op if path is empty.
  if path.len == 0: return
  let f = open(path, fmAppend)
  f.writeLine(line)
  f.close()

proc jsonFloat(v: float32): string =
  ## Serialize a float for JSON; non-finite → null (keeps JSONL parseable).
  if v == v and v > -1e30'f32 and v < 1e30'f32: $v else: "null"

proc hyperparmSnapshot(): string =
  ## Compact JSON object of current hyperparams (no outer braces).
  &"\"lr\":{hpLr},\"clipEpsilon\":{hpClipEpsilon}," &
  &"\"entropyCoeff\":{hpEntropyCoeff},\"valueLossCoeff\":{hpValueLossCoeff}," &
  &"\"maxGradNorm\":{hpMaxGradNorm},\"gamma\":{hpGamma},\"lam\":{hpLam}," &
  &"\"epochs\":{hpEpochs},\"miniBatchSize\":{hpMiniBatchSize}," &
  &"\"logStdFloor\":{logStdFloor},\"logStdCeiling\":{logStdCeiling},\"initialLogStd\":{initialLogStd}"

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

type
  InFlightBullet = object
    x, y:   float64  # current position
    vx, vy: float64  # velocity (pixels/tick)
    power:  float64

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:     array[STATE_DIM, float32]    # plain arrays — tensors NEVER cross threads
  lastAction:    array[ACTION_DIM, float32]
  lastLogP:      float32
  lastValue:     float32
  hasLastTrans:  bool
  lastActions:   BotActions    # previous tick's decoded actions (for state vector)
  roundRewardSum: float32      # cumulative reward this round (for live display)
  roundTicks:     int          # ticks this round
  # Fixed-size bullet buffer — NO heap on the shared bot object. A seq here is
  # allocated by the per-round bot thread and freed by the next round's thread
  # (bot.bullets = @[] on round start) → foreign-heap free under --threads:on +
  # ORC → SIGSEGV. N=4: the state vector only consumes the closest 3 slots.
  bullets:       array[4, InFlightBullet]
  bulletCount:   int

var ac = initActorCritic()
var gAdamStates: ACAdamStates  # persists across rounds
var roundCounter = 0
var roundsSinceUpdate = 0

# ── 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) =
  setAdjustGunForBodyTurn(true)
  setAdjustRadarForBodyTurn(true)
  setAdjustRadarForGunTurn(true)
  bot.tracker        = initEnemyTracker()
  # Do NOT clear bot.buffer here — transitions accumulate across rounds
  # until hpUpdateInterval rounds have passed (cleared in onRoundEnded).
  bot.prevEnergy     = 0.0'f32
  bot.prevEnemyE     = 0.0'f32
  bot.hasLastTrans   = false
  bot.roundRewardSum = 0.0'f32
  bot.roundTicks     = 0
  bot.bulletCount    = 0

method onRoundEnded*(bot: PPOBot, e: RoundEndedEventForBot) =
  inc roundCounter
  inc roundsSinceUpdate
  debugLog("[PO-ENTER] round=" & $roundCounter & " tid=" & $getThreadId())

  # Add round-end score bonus to last transition and mark it as episode boundary
  let roundReward = computeRoundReward(e.results.totalScore.float32)
  if bot.hasLastTrans and bot.buffer.len > 0:
    let lastIdx = bot.buffer.len - 1
    bot.buffer.transitions[lastIdx].reward += roundReward
    bot.buffer.transitions[lastIdx].done = true

  # Training progress display — one line per round in the UI console
  let ticks = bot.buffer.len
  var avgR = 0.0'f32
  if ticks > 0:
    var rewardSum = 0.0'f32
    for i in 0 ..< bot.buffer.len: rewardSum += bot.buffer.transitions[i].reward
    avgR = rewardSum / ticks.float32
    let avgRStr = formatFloat(avgR.float, ffDecimal, 3)
    printToStdOut(&"R:{roundCounter} ticks:{ticks} buf:{bot.buffer.len} sinceUpd:{roundsSinceUpdate}/{hpUpdateInterval} avgR:{avgRStr} score:{e.results.totalScore}\n")
    echo &"R:{roundCounter} ticks:{ticks} buf:{bot.buffer.len} sinceUpd:{roundsSinceUpdate}/{hpUpdateInterval} avgR:{avgRStr} score:{e.results.totalScore}"

  writeFile(weightsRoot / "round_counter.txt", $roundCounter)

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

  # Emit per-round game-stats JSON line
  let ts = int(epochTime())
  let jline = &"""{{"type":"round","round":{roundCounter},"ticks":{ticks},"bufLen":{bot.buffer.len},"avgReward":{avgR},"score":{e.results.totalScore},"ts":{ts}}}"""
  appendJsonLine(logFile, jline)

  # PPOB_EVAL_ONLY=1 → freeze training (pure evaluation): skip ppoUpdate and
  # checkpoint save, but keep advancing/writing round_counter.txt so run.sh's
  # remaining-rounds bookkeeping still works, and keep the game line above.
  if evalOnly:
    bot.buffer.clear()
    bot.hasLastTrans = false
    roundsSinceUpdate = 0
    return

  # Always save weights every round so round_counter.txt stays current.
  saveCheckpoint(ac, gAdamStates, weightsRoot, roundCounter)

  # Only update policy every hpUpdateInterval rounds (~3000 transitions).
  if roundsSinceUpdate >= hpUpdateInterval:
    # ponytail: synchronous update. Arraymancer tensors can't cross threads under
    # ORC — training-thread ppoUpdate frees/rebinds tensors owned by the bot
    # thread's heap (SIGSEGV, reproduced with a lone trainer thread on a fixed
    # buffer; save/channel/forward exonerated). The bot API runs events on one
    # bot thread, so inline is single-threaded. Revert to a background thread
    # only if tensors are rebuilt from plain data on that thread.
    printToStdOut(&"  train→ R:{roundCounter} bufLen:{bot.buffer.len}\n")
    echo &"  train→ R:{roundCounter} bufLen:{bot.buffer.len}"
    let m = ppoUpdate(ac, bot.buffer,
                      lastValue      = 0.0'f32,
                      adamStates     = gAdamStates,
                      epochs         = hpEpochs,
                      miniBatchSize  = hpMiniBatchSize,
                      clipEpsilon    = hpClipEpsilon,
                      entropyCoeff   = hpEntropyCoeff,
                      valueLossCoeff = hpValueLossCoeff,
                      lr             = hpLr,
                      maxGradNorm    = hpMaxGradNorm,
                      gamma          = hpGamma,
                      lam            = hpLam)
    printToStdOut(&"  trained R:{roundCounter} aLoss:{formatFloat(m.actorLoss.float, ffDecimal, 4)} vLoss:{formatFloat(m.valueLoss.float, ffDecimal, 4)} gNorm:{formatFloat(m.gradNorm.float, ffDecimal, 3)}\n")
    echo &"  trained R:{roundCounter} aLoss:{formatFloat(m.actorLoss.float, ffDecimal, 4)} vLoss:{formatFloat(m.valueLoss.float, ffDecimal, 4)} gNorm:{formatFloat(m.gradNorm.float, ffDecimal, 3)}"
    # Emit training-health JSON line
    let hp = hyperparmSnapshot()
    let ts2 = int(epochTime())
    let jline2 = &"""{{"type":"train","round":{roundCounter},"actorLoss":{jsonFloat(m.actorLoss)},"valueLoss":{jsonFloat(m.valueLoss)},"gradNorm":{jsonFloat(m.gradNorm)},"ts":{ts2},{hp}}}"""
    appendJsonLine(logFile, jline2)
    bot.buffer.clear()
    roundsSinceUpdate = 0

  bot.hasLastTrans = false
  debugLog("[PO-EXIT] round=" & $roundCounter & " tid=" & $getThreadId())

method run(bot: PPOBot) =
  debugLog("[RUN-ENTER] tid=" & $getThreadId())
  # Seed energy and goto/aimTo targets on first tick (remainingDistance = 0 initially)
  bot.prevEnergy = getEnergy().float32
  bot.prevEnemyE = if bot.tracker.hasContact: bot.tracker.current.energy.float32 else: 0.0'f32
  bot.lastActions.gotoX  = getX()
  bot.lastActions.gotoY  = getY()
  bot.lastActions.aimToX = getX()
  bot.lastActions.aimToY = getY()

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

    # Spawn bullet when enemy fired last scan tick. Edge-triggered: the flag is
    # consumed later (after the state build) so a radar gap (deadReckon ticks)
    # can't spawn k phantom bullets from one shot — exactly one bullet, and
    # state index 12 still pulses 1 on the detection tick.
    if bot.tracker.hasContact and bot.tracker.current.hasFired:
      if bot.bulletCount < 4:
        let power = bot.tracker.current.lastFirePower
        let speed = 20.0 - 3.0 * power
        # Approximate gun direction: bearing from enemy toward our position
        let myX = getX(); let myY = getY()
        let ang = arctan2(myY - bot.tracker.current.y, myX - bot.tracker.current.x)
        bot.bullets[bot.bulletCount] = InFlightBullet(
          x:     bot.tracker.current.x,
          y:     bot.tracker.current.y,
          vx:    speed * cos(ang),
          vy:    speed * sin(ang),
          power: power,
        )
        inc bot.bulletCount

    # Advance in-flight bullets and prune those off-arena (in-place compaction
    # into the fixed buffer — no per-tick heap churn).
    let aW = float64(getArenaWidth()); let aH = float64(getArenaHeight())
    var n = 0
    for i in 0 ..< bot.bulletCount:
      let b = bot.bullets[i]
      let nx = b.x + b.vx; let ny = b.y + b.vy
      if nx >= 0.0 and nx <= aW and ny >= 0.0 and ny <= aH:
        bot.bullets[n] = InFlightBullet(x: nx, y: ny, vx: b.vx, vy: b.vy, power: b.power)
        inc n
    bot.bulletCount = n

    # Convert to BulletData for state vector (closest 3 by distance to us)
    let myX2 = getX(); let myY2 = getY()
    var bulletData: array[4, BulletData]
    var bulletDataCount = 0
    for i in 0 ..< bot.bulletCount:
      let b = bot.bullets[i]
      bulletData[bulletDataCount] = BulletData(x: b.x, y: b.y, power: b.power)
      inc bulletDataCount
    # sort ascending by distance so the nearest threats fill slots 0-2
    if bulletDataCount > 1:
      bulletData.toOpenArray(0, bulletDataCount - 1).sort(proc(a, b: BulletData): int =
        let da = hypot(a.x - myX2, a.y - myY2)
        let db = hypot(b.x - myX2, b.y - myY2)
        cmp(da, db))

    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 remainingGotoDistance = hypot(bot.lastActions.gotoX - botData.x,
                                      bot.lastActions.gotoY - botData.y)
    let remainingGunAngle = abs(normalizeRelativeAngle(
      directionTo(botData.x, botData.y, bot.lastActions.aimToX, bot.lastActions.aimToY) -
      botData.gunDirection))
    let state = buildStateVector(botData, bot.tracker, remainingGotoDistance, remainingGunAngle, bulletData, bulletDataCount)
    # Consume the fired flag AFTER the state build: state index 12 saw the
    # detection-tick pulse, and the next iteration's spawn check sees false —
    # one shot → exactly one bullet, even across deadReckon gaps.
    bot.tracker.current.hasFired = false
    let (rawActs, logP) = ac.actorForward(state, deterministic = evalOnly)
    let value = ac.criticForward(state)
    let ex = if bot.tracker.hasContact: bot.tracker.current.x else: botData.arenaWidth  / 2.0
    let ey = if bot.tracker.hasContact: bot.tracker.current.y else: botData.arenaHeight / 2.0
    let acts = mapActions(rawActs,
                          getGunHeat().float,
                          botData.arenaWidth, botData.arenaHeight,
                          botData.x, botData.y,
                          botData.direction, botData.speed, botData.gunDirection,
                          ex, ey)

    # Compute tick reward from energy deltas + dense shaping
    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 arenaDiag   = float32(sqrt(botData.arenaWidth * botData.arenaWidth +
                                   botData.arenaHeight * botData.arenaHeight))
    let distEnemy   = if bot.tracker.hasContact:
                        float32(hypot(bot.tracker.current.x - botData.x,
                                      bot.tracker.current.y - botData.y))
                      else: arenaDiag
    let gunToEnemy  = if bot.tracker.hasContact:
                        abs(normalizeRelativeAngle(
                          arctan2(bot.tracker.current.y - botData.y,
                                  bot.tracker.current.x - botData.x) * 180.0 / PI -
                          botData.gunDirection)).float32
                      else: 180.0'f32
    let tickReward  = computeTickReward(myDelta, enemyDelta,
                                        distToEnemy    = distEnemy,
                                        maxDist        = arenaDiag,
                                        gunBearingAbs  = gunToEnemy)

    # Track running reward for in-game display
    bot.roundRewardSum += tickReward
    inc bot.roundTicks

    # 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,
        done:    false,  # episode boundary set in onRoundEnded
      )
      bot.buffer.add(tr)

    # Store current for next tick — plain arrays only. `state`/`rawActs` tensors
    # live and die on this thread; a NEW bot thread runs each round, so storing
    # tensors in the shared bot object would free round-N's heap memory from
    # round N+1's thread (SIGSEGV; confirmed empirically).
    bot.lastState   = stateToArr(state)
    bot.lastAction  = actionToArr(rawActs)
    bot.lastLogP    = logP
    bot.lastValue   = value
    bot.prevEnergy  = curEnergy
    bot.prevEnemyE  = curEnemyE
    bot.hasLastTrans = true
    bot.lastActions  = acts

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

    # In-game training progress overlay
    let avgR = if bot.roundTicks > 0: bot.roundRewardSum / bot.roundTicks.float32
               else: 0.0'f32
    let avgRStr = formatFloat(avgR.float, ffDecimal, 2)
    drawText(&"R:{roundCounter} avg:{avgRStr}", getX(), getY() - 40.0)

    go()

when isMainModule:
  createDir(weightsRoot)
  cleanStaleTempDirs(weightsRoot)
  let loadResult = loadBestAvailable(ac, gAdamStates, weightsRoot)
  if loadResult.loaded:
    roundCounter = loadResult.roundNum

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