## Virtual bullet tracker.
## Spawns virtual bullets per gun×power bin every tick (no real firing).
## Resolves by travel distance. Rolling window fitness per gun×power.
## Calls onResult() on the owning gun when a bullet resolves.

import std/math
import std/tables
import gun_interface

const
  PowerBins*    = [1.0, 1.5, 2.0, 3.0]  ## 4 bins; ponytail: fixed array, add runtime config if needed
  WindowSize*   = 100                     ## rolling window ticks for fitness
  MaxBullets*   = 2048                    ## hard cap; ponytail: ring buffer, resize if more guns added
  MinHitRate*          = 0.40   ## 40% threshold for acceptable power selection
  MinObsBeforeCompete* = 15    ## min observations before a gun×bin enters competition

type
  GunId*   = int  ## index into the guns seq

  VirtualBullet* = object
    gunId*:      GunId
    powerBin*:   int       ## index into PowerBins
    targetId*:   int       ## enemy bot ID this bullet was aimed at
    fireX*, fireY*: float
    aimX*, aimY*:   float  ## predicted target (absolute)
    bulletSpeed*:  float
    travelDist*:   float   ## accumulated px so far
    fireDist*:     float   ## distance to target at fire time
    active*:       bool

  FitnessWindow* = object
    ## Ring buffer of hit booleans.
    hits*:  array[WindowSize, bool]
    count*: int  ## total samples so far (capped at WindowSize for rate)
    head*:  int

  GunFitness* = object
    bins*: array[len(PowerBins), FitnessWindow]

  VirtualTracker* = object
    bullets*:   array[MaxBullets, VirtualBullet]
    head*:      int  ## ring buffer head
    numGuns*:   int
    fitness*:   Table[int, seq[GunFitness]]  ## keyed by enemy bot ID, indexed by GunId

proc initTracker*(numGuns: int): VirtualTracker =
  result.numGuns = numGuns

proc hitRate*(fw: FitnessWindow): float =
  ## Returns fraction of hits in the rolling window. 0.0 when no data.
  if fw.count == 0: return 0.0
  let n = min(fw.count, WindowSize)
  var h = 0
  for i in 0..<n: h += (if fw.hits[i]: 1 else: 0)
  result = h.float / n.float

proc record(fw: var FitnessWindow, hit: bool) =
  fw.hits[fw.head] = hit
  fw.head = (fw.head + 1) mod WindowSize
  inc fw.count

proc spawnBullets*(t: var VirtualTracker, gunId: GunId,
                   predictions: array[len(PowerBins), GunPrediction],
                   state: WorldState, targetId: int) =
  ## Call once per gun per tick with predictions for all power bins.
  ## Lazily creates fitness entry for targetId on first spawn.
  if targetId notin t.fitness:
    t.fitness[targetId] = newSeq[GunFitness](t.numGuns)
  for binIdx in 0..<len(PowerBins):
    let power   = PowerBins[binIdx]
    let speed   = bulletSpeed(power)
    let pred    = predictions[binIdx]
    let fireDist = hypot(pred.x - state.selfX, pred.y - state.selfY)
    let slot    = t.head mod MaxBullets
    t.bullets[slot] = VirtualBullet(
      gunId:      gunId,
      powerBin:   binIdx,
      targetId:   targetId,
      fireX:      state.selfX,
      fireY:      state.selfY,
      aimX:       pred.x,
      aimY:       pred.y,
      bulletSpeed: speed,
      travelDist: 0.0,
      fireDist:   fireDist,
      active:     true,
    )
    t.head = (t.head + 1) mod MaxBullets

const StaleTicks* = 20  ## discard bullet if target not seen within this many ticks

proc tickBullets*(t: var VirtualTracker, state: WorldState,
                  enemies: Table[int, tuple[x, y: float, lastSeenTick: int, alive: bool]],
                  onResolved: proc(gunId: GunId, binIdx: int, e: FeedbackEvent)) =
  ## Advance all active bullets one tick. Resolve when bullet reaches target distance.
  ## Uses per-target enemy position from enemies table for accurate miss distance.
  ## Discards bullet without scoring if target is dead or stale (> StaleTicks).
  for i in 0..<MaxBullets:
    var b = addr t.bullets[i]
    if not b.active: continue
    b.travelDist += b.bulletSpeed
    if b.travelDist < b.fireDist: continue

    # Resolved: look up the correct enemy position
    var ex, ey: float
    if b.targetId in enemies:
      let e = enemies[b.targetId]
      if not e.alive or (state.tick - e.lastSeenTick) > StaleTicks:
        b.active = false
        continue
      ex = e.x; ey = e.y
    else:
      # No data for this target — fall back to selected enemy in state
      ex = state.enemyX; ey = state.enemyY

    let dx    = b.aimX - b.fireX
    let dy    = b.aimY - b.fireY
    let dist  = hypot(dx, dy)
    let (bx, by) =
      if dist < 1e-6: (b.aimX, b.aimY)
      else: (b.fireX + dx / dist * b.travelDist,
             b.fireY + dy / dist * b.travelDist)
    let missDist = hypot(bx - ex, by - ey)
    let hit = missDist < BotRadius

    if b.targetId in t.fitness:
      t.fitness[b.targetId][b.gunId].bins[b.powerBin].record(hit)

    let fe = FeedbackEvent(
      prediction:  GunPrediction(x: b.aimX, y: b.aimY),
      actualX:     ex,
      actualY:     ey,
      bulletPower: PowerBins[b.powerBin],
      missDistance: missDist,
      hit:         hit,
    )
    onResolved(b.gunId, b.powerBin, fe)
    b.active = false

proc fitnessFor(t: VirtualTracker, targetId: int): seq[GunFitness] =
  ## Returns fitness seq for targetId, or merges all enemies as fallback.
  ## ponytail: merge is O(enemies*guns*bins), fine for small counts
  if targetId >= 0 and targetId in t.fitness:
    return t.fitness[targetId]
  # Aggregate across all enemies
  result = newSeq[GunFitness](t.numGuns)
  for perEnemy in t.fitness.values:
    for gunId in 0..<t.numGuns:
      for binIdx in 0..<len(PowerBins):
        let src = perEnemy[gunId].bins[binIdx]
        var dst = addr result[gunId].bins[binIdx]
        for k in 0..<min(src.count, WindowSize):
          dst[].record(src.hits[k])

proc bestPower*(t: VirtualTracker, gunId: GunId, targetId: int = -1): (int, float) =
  ## Returns (binIdx, power) with highest power that has >= MinHitRate.
  ## Falls back to lowest power bin if nothing qualifies yet.
  ## Uses per-enemy fitness when targetId >= 0 and data exists; else aggregate.
  let fit = t.fitnessFor(targetId)
  result = (0, PowerBins[0])
  for binIdx in countdown(len(PowerBins) - 1, 0):
    let rate = fit[gunId].bins[binIdx].hitRate()
    if rate >= MinHitRate or fit[gunId].bins[binIdx].count == 0:
      return (binIdx, PowerBins[binIdx])

proc bestGun*(t: VirtualTracker, targetId: int = -1): GunId =
  ## Pick gun with highest hit rate across all power bins.
  ## Guns with fewer than MinObsBeforeCompete observations are skipped
  ## unless every gun is below threshold (then fall back to best of all).
  ## Uses per-enemy fitness when targetId >= 0 and data exists; else aggregate.
  ## ponytail: O(n*bins), fine for small gun counts
  let fit = t.fitnessFor(targetId)
  proc bestAmong(fit: seq[GunFitness], requireMin: bool): GunId =
    var bestRate = -1.0
    result = 0
    for gunId in 0..<fit.len:
      var maxCount = 0
      for binIdx in 0..<len(PowerBins):
        maxCount = max(maxCount, fit[gunId].bins[binIdx].count)
      if requireMin and maxCount < MinObsBeforeCompete: continue
      for binIdx in 0..<len(PowerBins):
        let r = fit[gunId].bins[binIdx].hitRate()
        if r > bestRate:
          bestRate = r
          result = gunId

  var anyQualifies = false
  for gunId in 0..<fit.len:
    for binIdx in 0..<len(PowerBins):
      if fit[gunId].bins[binIdx].count >= MinObsBeforeCompete:
        anyQualifies = true
        break
    if anyQualifies: break

  result = if anyQualifies: fit.bestAmong(true) else: fit.bestAmong(false)