## 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

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

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

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

  KNNGun* = object
    obs:                seq[Obs]
    obsHead:            int         # ring-buffer write index
    waves:              seq[KNNWave]
    # per-tick cache
    cachedTick:         int
    cachedWaveStored:   bool
    # rolling normalization ranges
    featMin:            array[7, float]
    featMax:            array[7, float]
    # state for feature extraction
    lastSpeed:          float
    lastDirection:      float        # +1 or -1
    timeSinceDirChange: int
    debugGraphics*:     bool

proc initKNNGun*(): KNNGun =
  result.cachedTick    = -1
  result.lastDirection = 1.0
  result.debugGraphics = false
  for i in 0..6:
    result.featMin[i] =  1e18
    result.featMax[i] = -1e18

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

proc normFeat(g: KNNGun, raw: array[7, float]): array[7, float] =
  for i in 0..6:
    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[7, float]) =
  for i in 0..6:
    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[7, 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)

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

# ── 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 dist    = sqrt(dx*dx + dy*dy)
  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.cachedWaveStored = false

    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

  # Store wave once per tick
  if not g.cachedWaveStored:
    let feat = buildFeatures(state, g.lastSpeed, g.lastDirection, g.timeSinceDirChange)
    g.updateMinMax(feat)
    g.waves.add KNNWave(
      fireX:       state.selfX,
      fireY:       state.selfY,
      fireBearing: bearing,
      feat:        feat,
    )
    g.lastSpeed = state.enemySpeed
    g.cachedWaveStored = true

  # Cold start — no data yet
  if g.obs.len == 0:
    return GunPrediction(
      x: clamp(state.selfX + cos(bearing) * dist, BotRadius, state.arenaWidth  - BotRadius),
      y: clamp(state.selfY + sin(bearing) * dist, 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(state.selfX + cos(bearing) * dist, BotRadius, state.arenaWidth  - BotRadius),
      y: clamp(state.selfY + sin(bearing) * dist, 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(state.selfX + cos(bearing) * dist, BotRadius, state.arenaWidth  - BotRadius),
      y: clamp(state.selfY + sin(bearing) * dist, 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 = bearing + clamp(bestGF, -1.0, 1.0) * mea
  let px = state.selfX + cos(aimAngle) * dist
  let py = state.selfY + sin(aimAngle) * dist

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

proc onResult*(g: var KNNGun, e: FeedbackEvent) =
  if g.waves.len == 0: return

  let w = g.waves[0]
  g.waves.delete(0)

  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