## State vector builder — produces 57-float normalized tensor for PPO policy.
## No bot API imports; takes plain BotState + EnemyTracker structs.

import std/math
import arraymancer
import ./enemy_tracker

type
  BotStateData* = object
    ## Plain data mirror of the bot's observable state.
    x*, y*: float64
    direction*: float64
    speed*: float64
    energy*: float64
    gunDirection*: float64
    gunHeat*: float64
    arenaWidth*, arenaHeight*: float64

  BulletData* = object
    ## Enemy bullet in flight (absolute arena coords + fire power).
    x*, y*: float64
    power*: float64  # fire power in [0.1, 3.0]; speed = 20 - 3*power

proc buildStateVector*(bot: BotStateData; enemy: EnemyTracker;
                       remainingGotoDistance: float64 = 0.0;
                       remainingGunAngle: float64 = 0.0;
                       bullets: openArray[BulletData] = [];
                       bulletCount: int = 0): Tensor[float32] =
  ## Build the 57-float normalized state tensor.
  ## Indices 0-43: existing features. Indices 44-55: up to 3 bullet slots (4 floats each).
  ## Index 56: scan staleness (ticksSinceLastScan / 30, clamped to 1).
  ## All values clipped to roughly [-1, 1] via division by physical maxima.
  result = zeros[float32](57)

  let aW = bot.arenaWidth
  let aH = bot.arenaHeight
  let diag = sqrt(aW * aW + aH * aH)  # ≈ 1700 for 1200×800
  let wallMax = max(aW, aH)

  # --- Current tick: own bot (indices 0-6) ---
  result[0] = float32(bot.x / aW)
  result[1] = float32(bot.y / aH)
  result[2] = float32(bot.direction / 360.0)
  result[3] = float32(bot.speed / 8.0)
  result[4] = float32(bot.energy / 100.0)
  result[5] = float32(bot.gunDirection / 360.0)
  result[6] = float32(bot.gunHeat / 1.8)

  # --- Current tick: enemy (indices 7-13) ---
  if enemy.hasContact:
    result[7]  = float32(enemy.current.x / aW)
    result[8]  = float32(enemy.current.y / aH)
    result[9]  = float32(enemy.current.direction / 360.0)
    result[10] = float32(enemy.current.speed / 8.0)
    result[11] = float32(enemy.current.energy / 100.0)
    result[12] = float32(if enemy.current.hasFired: 1.0 else: 0.0)
    result[13] = float32(enemy.current.lastFirePower / 3.0)
  # else: remain 0.0

  # --- Derived features (indices 14-21) ---
  if enemy.hasContact:
    # Enemy acceleration (speed delta from last history entry)
    # Max delta is ±8 (stopped ↔ full speed); divide by 8 to normalize
    if enemy.historyCount >= 1:
      result[14] = float32((enemy.current.speed - enemy.history[0].speed) / 8.0)
    # Enemy turn rate (direction delta from last history entry, normalized to [-1,1])
    # Divide by 180 (max possible relative rotation) rather than 10 (max body turn rate)
    # ponytail: /180 covers all cases; use /10 if you want sensitivity to small turns
    if enemy.historyCount >= 1:
      let dirDelta = ((enemy.current.direction - enemy.history[0].direction) + 540.0) mod 360.0 - 180.0
      result[15] = float32(dirDelta / 180.0)
    # Relative bearing to enemy (signed, from bot perspective)
    let dx = enemy.current.x - bot.x
    let dy = enemy.current.y - bot.y
    let absDir = (180.0 * arctan2(dx, dy) / PI + 360.0) mod 360.0
    let relBearing = ((absDir - bot.direction) + 540.0) mod 360.0 - 180.0
    result[16] = float32(relBearing / 180.0)
    # Distance to enemy
    let dist = sqrt(dx * dx + dy * dy)
    result[17] = float32(dist / diag)

  # Wall distances (indices 18-21): top, bottom, left, right
  # top = distance from bot to top wall (y=aH), bottom = distance to bottom (y=0)
  # left = distance to left (x=0), right = distance to right (x=aW)
  result[18] = float32((aH - bot.y) / wallMax)  # top
  result[19] = float32(bot.y / wallMax)           # bottom
  result[20] = float32(bot.x / wallMax)           # left
  result[21] = float32((aW - bot.x) / wallMax)   # right

  # --- History: 5 ticks × 4 floats = 20 floats (indices 22-41) ---
  for i in 0 ..< 5:
    let base = 22 + i * 4
    if i < enemy.historyCount:
      result[base + 0] = float32(enemy.history[i].x / aW)
      result[base + 1] = float32(enemy.history[i].y / aH)
      result[base + 2] = float32(enemy.history[i].direction / 360.0)
      result[base + 3] = float32(enemy.history[i].speed / 8.0)
    # else: remain 0.0 (pad)

  # --- Goto controller inputs (indices 42-43) ---
  result[42] = float32(remainingGotoDistance / diag)
  result[43] = float32(remainingGunAngle / 180.0)

  # --- Bullet tracking (indices 44-55): up to 3 enemy bullets, 4 floats each ---
  # Per bullet: relX/aW, relY/aH, speed/20, ticksToImpact/diag
  # Positions are relative to bot (useful for dodging). Slots beyond bulletCount stay 0.
  for i in 0 ..< min(bulletCount, 3):
    let b = bullets[i]
    let bSpeed = 20.0 - 3.0 * b.power  # Tank Royale bullet speed formula
    let bdx = b.x - bot.x
    let bdy = b.y - bot.y
    let bdist = sqrt(bdx * bdx + bdy * bdy)
    let ticks = if bSpeed > 0.0: bdist / bSpeed else: 0.0
    let base = 44 + i * 4
    result[base + 0] = float32(bdx / bot.arenaWidth)
    result[base + 1] = float32(bdy / bot.arenaHeight)
    result[base + 2] = float32(bSpeed / 20.0)
    result[base + 3] = float32(ticks / diag)

  # --- Scan staleness (index 56) ---
  result[56] = float32(min(enemy.current.ticksSinceLastScan.float64 / 30.0, 1.0))