## KNN gun: K-nearest-neighbor statistical targeting inspired by DrussGT's DC gun.
## Builds a feature vector per scan, stores resolved GF outcomes, queries KNN at
## predict time and picks the GF with the highest Gaussian-weighted density.
## ponytail: linear scan O(n*k), cap at 2000 obs — KD-tree if perf matters at scale.

import std/[math]
import gun_harness/gun_interface
import gun_harness/virtual_bullets as vb  # PowerBins
import guns/lead_forecast

const
  MaxObs     = 2000    # ring-buffer cap
  KCap       = 50      # hard ceiling on K
  KernelW    = 0.3     # Gaussian kernel width multiplier
  DensityBins = 60     # scan resolution for peak-GF search
  NFeat      = 8       # feature vector length (see buildFeatures)

type
  Obs = object
    feat: array[NFeat, float]  # normalized feature vector
    gf:   float            # observed GF at wave resolution

  KNNWave = object
    fireX, fireY: float
    fireBearing:  float
    feat:         array[NFeat, float]

  KNNGun* = object
    obs:                seq[Obs]
    obsHead:            int         # ring-buffer write index
    # One wave queue per power bin; matched on bulletSpeed / bulletPower so a
    # resolved bullet only ever learns from a wave fired with the same power.
    waves:              array[len(vb.PowerBins), seq[KNNWave]]
    waveHead:           array[len(vb.PowerBins), int]   # O(1) pop cursor
    waveStoredTick:     array[len(vb.PowerBins), int]   # last tick a wave was queued for this bin
    # per-tick cache
    cachedTick:         int
    vt:                 VelocityTracker  # enemy velocity history (base selection)
    tickWave:           KNNWave     # wave template for the current tick (features computed once)
    # rolling normalization ranges
    featMin:            array[NFeat, float]
    featMax:            array[NFeat, float]
    # state for feature extraction
    lastSpeed:          float
    lastDirection:      float        # +1 or -1
    timeSinceDirChange: int
    wavePushes*:        int
    waveStarved*:       int
    debugGraphics*:     bool

proc initKNNGun*(): KNNGun =
  result.cachedTick    = -1
  result.lastDirection = 1.0
  result.debugGraphics = false
  for b in 0..<len(vb.PowerBins):
    result.waveStoredTick[b] = -1
  for i in 0..<NFeat:
    result.featMin[i] =  1e18
    result.featMax[i] = -1e18
  # Energy (dim NFeat-1) stays on a FIXED [0,1] scale: a threshold behaviour is
  # keyed to the target's ABSOLUTE energy, so min-max normalizing it against the
  # observed range would smear the very boundary the feature exists to expose.
  result.featMin[NFeat - 1] = 0.0
  result.featMax[NFeat - 1] = 1.0

# ── helpers ──────────────────────────────────────────────────────────────────

proc normFeat(g: KNNGun, raw: array[NFeat, float]): array[NFeat, float] =
  for i in 0..<NFeat:
    let span = g.featMax[i] - g.featMin[i]
    result[i] = if span > 1e-9: (raw[i] - g.featMin[i]) / span else: 0.0

proc updateMinMax(g: var KNNGun, raw: array[NFeat, float]) =
  for i in 0..<NFeat - 1:
    if raw[i] < g.featMin[i]: g.featMin[i] = raw[i]
    if raw[i] > g.featMax[i]: g.featMax[i] = raw[i]

proc buildFeatures(state: WorldState, lastSpeed, lastDir: float,
                   tsdc: int): array[NFeat, float] =
  let dx     = state.enemyX - state.selfX
  let dy     = state.enemyY - state.selfY
  let dist   = sqrt(dx*dx + dy*dy)
  let arenaDiag = sqrt(state.arenaWidth*state.arenaWidth + state.arenaHeight*state.arenaHeight)

  # bearing to enemy (0°=East, standard Tank Royale)
  let bearing = arctan2(dy, dx)
  # angle of enemy heading relative to bearing
  let relHead = state.enemyHeading - bearing
  let latVel  = state.enemySpeed * sin(relHead)
  let advVel  = state.enemySpeed * (-cos(relHead))
  let accel   = state.enemySpeed - lastSpeed  # signed delta

  # wall distances: how far enemy can travel fwd/bwd before hitting wall
  # approximate: project enemy heading to nearest wall in each axis
  let ex = state.enemyX
  let ey = state.enemyY
  let eh = state.enemyHeading
  # forward distances to each wall in heading direction
  let fwdX = if cos(eh) > 0: (state.arenaWidth  - ex) / max(abs(cos(eh)), 1e-9)
             else:            ex / max(abs(cos(eh)), 1e-9)
  let fwdY = if sin(eh) > 0: (state.arenaHeight - ey) / max(abs(sin(eh)), 1e-9)
             else:            ey / max(abs(sin(eh)), 1e-9)
  let fwdDist = min(fwdX, fwdY)
  # backward = forward in opposite direction
  let bwdX = if cos(eh) < 0: (state.arenaWidth  - ex) / max(abs(cos(eh)), 1e-9)
             else:            ex / max(abs(cos(eh)), 1e-9)
  let bwdY = if sin(eh) < 0: (state.arenaHeight - ey) / max(abs(sin(eh)), 1e-9)
             else:            ey / max(abs(sin(eh)), 1e-9)
  let bwdDist = min(bwdX, bwdY)

  result[0] = abs(latVel) / 8.0
  result[1] = clamp(advVel / 8.0, -1.0, 1.0) * 0.5 + 0.5  # shift to [0,1]
  result[2] = clamp(dist / arenaDiag, 0.0, 1.0)
  result[3] = clamp(accel / 2.0, -1.0, 1.0) * 0.5 + 0.5
  result[4] = clamp(float(tsdc) / 100.0, 0.0, 1.0)
  result[5] = clamp(fwdDist / arenaDiag, 0.0, 1.0)
  result[6] = clamp(bwdDist / arenaDiag, 0.0, 1.0)
  # Enemy energy. The only feature that can separate behaviours that depend on
  # the target's own remaining energy (e.g. an energy-threshold turner that
  # changes movement below 30). Kept on a fixed [0,1] scale (see initKNNGun).
  result[7] = clamp(state.enemyEnergy / 100.0, 0.0, 1.0)

proc euclidean(a, b: array[NFeat, float]): float {.inline.} =
  for i in 0..<NFeat:
    let d = a[i] - b[i]
    result += d * d
  result = sqrt(result)

proc binForSpeed(spd: float): int {.inline.} =
  ## Map a virtual-bullet speed back to its power-bin index. All four bin speeds
  ## are exactly representable floats; the epsilon is belt-and-braces only.
  for i in 0..<len(vb.PowerBins):
    if abs(spd - bulletSpeed(vb.PowerBins[i])) < 1e-6:
      return i
  -1

proc binForPower(power: float): int {.inline.} =
  ## Map a FeedbackEvent.bulletPower back to its power-bin index.
  for i in 0..<len(vb.PowerBins):
    if abs(power - vb.PowerBins[i]) < 1e-6:
      return i
  -1

proc takeOldestWave(g: var KNNGun, binIdx: int): (bool, KNNWave) =
  ## Pop the oldest unresolved wave for this bin (O(1) amortized via waveHead).
  if binIdx < 0 or g.waveHead[binIdx] >= g.waves[binIdx].len:
    return (false, KNNWave())
  result = (true, g.waves[binIdx][g.waveHead[binIdx]])
  inc g.waveHead[binIdx]
  if g.waveHead[binIdx] >= 64 and
     g.waveHead[binIdx] * 2 >= g.waves[binIdx].len:
    g.waves[binIdx] = g.waves[binIdx][g.waveHead[binIdx] .. g.waves[binIdx].high]
    g.waveHead[binIdx] = 0

# ── Gun interface ─────────────────────────────────────────────────────────────

proc predict*(g: var KNNGun, state: WorldState, bulletSpd: float): GunPrediction =
  if bulletSpd <= 0.0:
    return GunPrediction(x: state.enemyX, y: state.enemyY)

  let dx      = state.enemyX - state.selfX
  let dy      = state.enemyY - state.selfY
  let bearing = arctan2(dy, dx)
  let mea     = arcsin(clamp(8.0 / bulletSpd, -1.0, 1.0))

  # Track direction change — update state once per tick
  if state.tick != g.cachedTick:
    g.cachedTick = state.tick
    g.vt.observe(state)

    let relHead  = state.enemyHeading - bearing
    let latVel   = state.enemySpeed * sin(relHead)
    let newDir   = if latVel >= 0: 1.0 else: -1.0
    if newDir != g.lastDirection and abs(latVel) > 0.01:
      g.timeSinceDirChange = 0
      g.lastDirection = newDir
    else:
      inc g.timeSinceDirChange

    # Compute the tick's feature vector ONCE, before lastSpeed is advanced, so
    # every power bin fired this tick shares identical features. lastSpeed is
    # only advanced here (once/tick), not once per bin.
    let feat = buildFeatures(state, g.lastSpeed, g.lastDirection, g.timeSinceDirChange)
    g.updateMinMax(feat)
    g.tickWave = KNNWave(
      fireX:       state.selfX,
      fireY:       state.selfY,
      fireBearing: bearing,
      feat:        feat,
    )
    g.lastSpeed = state.enemySpeed

  # Base forecast: the KNN learns the GF residual against a self-consistent base
  # prediction (see lead_forecast.nim).
  let f = forecastRadialBlend(state, bulletSpd, g.vt)

  # Queue at most one wave per (tick, power bin). The fire site's extra predict()
  # call for the selected bin lands on the same tick and reuses the queued wave.
  let binIdx = binForSpeed(bulletSpd)
  if binIdx >= 0 and g.waveStoredTick[binIdx] != state.tick:
    # fireBearing is the base forecast bearing, which is per-power (flight time
    # differs per bin); override the shared per-tick template here.
    var w = g.tickWave
    w.fireBearing = f.bearing
    g.waves[binIdx].add w
    g.waveStoredTick[binIdx] = state.tick
    inc g.wavePushes

  # Cold start — no data yet: fall back to the self-consistent linear forecast.
  if g.obs.len == 0:
    return GunPrediction(
      x: clamp(f.x, BotRadius, state.arenaWidth  - BotRadius),
      y: clamp(f.y, BotRadius, state.arenaHeight - BotRadius),
    )

  # Build query feature vector (use current state)
  let queryRaw = buildFeatures(state, g.lastSpeed, g.lastDirection, g.timeSinceDirChange)
  let query    = g.normFeat(queryRaw)

  # KNN: linear scan, pick k = max(5, min(sqrt(n), KCap))
  # Fall back to head-on when not enough neighbors to be meaningful
  let n = g.obs.len
  if n < 5:
    return GunPrediction(
      x: clamp(f.x, BotRadius, state.arenaWidth  - BotRadius),
      y: clamp(f.y, BotRadius, state.arenaHeight - BotRadius),
    )
  let k = max(5, min(int(sqrt(float(n))), KCap))

  # Partial selection: maintain k-best by tracking max distance in result set
  # ponytail: O(n*k) insertion; fine for n<=2000, k<=50
  var bestDists = newSeq[float](k)
  var bestGFs   = newSeq[float](k)
  var worstIdx  = 0
  var filled    = 0

  for i in 0..<n:
    let normFeat = g.normFeat(g.obs[i].feat)
    let d = euclidean(query, normFeat)
    if filled < k:
      bestDists[filled] = d
      bestGFs[filled]   = g.obs[i].gf
      inc filled
      if filled == k:
        # find worst
        worstIdx = 0
        for j in 1..<k:
          if bestDists[j] > bestDists[worstIdx]: worstIdx = j
    elif d < bestDists[worstIdx]:
      bestDists[worstIdx] = d
      bestGFs[worstIdx]   = g.obs[i].gf
      worstIdx = 0
      for j in 1..<k:
        if bestDists[j] > bestDists[worstIdx]: worstIdx = j

  if filled == 0:
    return GunPrediction(
      x: clamp(f.x, BotRadius, state.arenaWidth  - BotRadius),
      y: clamp(f.y, BotRadius, state.arenaHeight - BotRadius),
    )

  # Inverse-distance weights, Gaussian (same as DrussGT getBearingGaussian)
  var sumDist = 1e-30
  for i in 0..<filled: sumDist += bestDists[i]
  let invAvg = float(filled) / sumDist

  # Find GF range of neighbors
  var gfMin = bestGFs[0]
  var gfMax = bestGFs[0]
  for i in 1..<filled:
    if bestGFs[i] < gfMin: gfMin = bestGFs[i]
    if bestGFs[i] > gfMax: gfMax = bestGFs[i]

  # Scan DensityBins points in [gfMin, gfMax] for peak density
  let span  = max(gfMax - gfMin, 1e-9)
  let step  = span / float(DensityBins - 1)
  var bestGF    = gfMin
  var bestScore = -1.0

  for b in 0..<DensityBins:
    let testGF = gfMin + float(b) * step
    var score  = 0.0
    for i in 0..<filled:
      let w = exp(-0.5 * (bestDists[i] * invAvg) * (bestDists[i] * invAvg))
      let dg = (testGF - bestGFs[i]) / max(span * KernelW, 1e-9)
      score += w * exp(-0.5 * dg * dg)
    if score > bestScore:
      bestScore = score
      bestGF    = testGF

  let aimAngle = f.bearing + clamp(bestGF, -1.0, 1.0) * mea
  let px = state.selfX + cos(aimAngle) * f.dist
  let py = state.selfY + sin(aimAngle) * f.dist

  GunPrediction(
    x: clamp(px, BotRadius, state.arenaWidth  - BotRadius),
    y: clamp(py, BotRadius, state.arenaHeight - BotRadius),
  )

proc onResult*(g: var KNNGun, e: FeedbackEvent) =
  let binIdx = binForPower(e.bulletPower)
  if binIdx < 0: return
  let (found, w) = g.takeOldestWave(binIdx)
  if not found:
    inc g.waveStarved
    return

  let speed  = bulletSpeed(e.bulletPower)
  let mea    = arcsin(clamp(8.0 / speed, -1.0, 1.0))
  let ax     = e.actualX - w.fireX
  let ay     = e.actualY - w.fireY
  var delta  = arctan2(ay, ax) - w.fireBearing
  while delta >  PI: delta -= 2.0 * PI
  while delta < -PI: delta += 2.0 * PI

  let gf = if mea > 1e-10: clamp(delta / mea, -1.0, 1.0) else: 0.0

  g.updateMinMax(w.feat)

  if g.obs.len < MaxObs:
    g.obs.add Obs(feat: w.feat, gf: gf)
  else:
    # ring buffer
    g.obs[g.obsHead] = Obs(feat: w.feat, gf: gf)
    g.obsHead = (g.obsHead + 1) mod MaxObs