## First guard tests for the gun selector + the speed-sensitivity checks for the
## per-tick caching bug class.
##
## Headless: no Java, no server, no battle. Run with plain
##   nim c -r common_libs/tests/test_gun_harness.nim
##
## Selection tests seed the tracker's exported fitness windows directly instead of
## dragging virtual bullets through spawnBullets/tickBullets. That is deliberate:
## it makes exact hit-rates (and therefore tie/rng/floor behaviour) deterministic
## and fast. The spawn/tick pipeline itself is exercised by the droppedBullets
## test below and by the full gauntlet.

import std/[math, random, tables, os]
import gun_harness/gun_interface
import gun_harness/virtual_bullets
import gun_harness/selector
import gun_harness/offline_range
import guns/stop_shot
import guns/displacement
import guns/averaged_lead
import guns/pattern_matcher
import guns/head_on
import guns/linear

var failures = 0

proc check(name: string, ok: bool) =
  if ok:
    echo "PASS: ", name
  else:
    echo "FAIL: ", name
    inc failures

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

proc seedWindow(t: var VirtualTracker, targetId, gunId, binIdx, hits, misses: int) =
  ## Narrowly-scoped test helper: write `hits`/`misses` samples straight into a
  ## gun×bin fitness window (fields are exported by virtual_bullets).
  if targetId notin t.fitness:
    t.fitness[targetId] = newSeq[GunFitness](t.numGuns)
  var fw = addr t.fitness[targetId][gunId].bins[binIdx]
  for _ in 0..<hits:   recordHit(fw[], true)
  for _ in 0..<misses: recordHit(fw[], false)

proc ws(tick: int, ex, ey, espeed, eheading: float): WorldState =
  WorldState(selfX: 100.0, selfY: 100.0, enemyX: ex, enemyY: ey,
             enemySpeed: espeed, enemyHeading: eheading,
             arenaWidth: 1000.0, arenaHeight: 1000.0, tick: tick)

proc pointsDiffer(a, b: GunPrediction): bool =
  abs(a.x - b.x) > 0.5 or abs(a.y - b.y) > 0.5

# ── selector guards ──────────────────────────────────────────────────────────

proc testColdBestGun() =
  var t = initTracker(3)
  check "bestGun on a cold tracker returns 0 (HeadOn)", t.bestGun(-1) == 0

proc testRandomTiebreak() =
  # Two guns with an identical, well-observed hit rate: the tiebreak must expose
  # both ids. Before the random tiebreak landed this always returned index 0.
  var t = initTracker(2)
  seedWindow(t, 7, gunId = 0, binIdx = 0, hits = 50, misses = 0)
  seedWindow(t, 7, gunId = 1, binIdx = 0, hits = 50, misses = 0)
  var seen: array[2, bool]
  for _ in 0..<500:
    let g = t.bestGun(-1)
    if g >= 0 and g < 2: seen[g] = true
  check "random tiebreak returns BOTH tied gun ids (no index-0 determinism)",
        seen[0] and seen[1]

proc testBestGunDeterministicWinner() =
  # gun 2 clearly best and past MinObsBeforeCompete; must win every call.
  var t = initTracker(3)
  seedWindow(t, 7, gunId = 0, binIdx = 0, hits = 25, misses = 25)  # 50 obs, 50%
  seedWindow(t, 7, gunId = 1, binIdx = 0, hits = 0,  misses = 0)   # cold, skipped
  seedWindow(t, 7, gunId = 2, binIdx = 0, hits = 50, misses = 0)   # 50 obs, 100%
  var allTwo = true
  for _ in 0..<100:
    if t.bestGun(-1) != 2: allTwo = false
  check "gun with clearly best rate and >= MinObsBeforeCompete wins deterministically",
        allTwo

# ── sticky selector (hysteresis) ──────────────────────────────────────────────
# `selectGun` is the live path; `bestGun`/`chooseFromFit` stay memoryless. These
# checks pin the dwell, switch-margin, disqualification and floor behaviours.

proc testHysteresisDwell() =
  var t = initTracker(2)
  seedWindow(t, 7, gunId = 0, binIdx = 0, hits = 60,  misses = 40)  # 60%
  seedWindow(t, 7, gunId = 1, binIdx = 0, hits = 100, misses = 0)   # 100%
  t.currentGun = 0
  t.currentSince = 0
  check "hysteresis: a clearly better gun is held off inside the dwell window",
        t.selectGun(7, tick = 5) == 0
  check "hysteresis: dwell expiry lets the better gun take over",
        t.selectGun(7, tick = GunDwellTicks) == 1

proc testHysteresisSwitchMargin() =
  # 80% incumbent vs 83% challenger: inside a 5% relative margin -> hold.
  var t = initTracker(2)
  seedWindow(t, 7, gunId = 0, binIdx = 0, hits = 80, misses = 20)
  seedWindow(t, 7, gunId = 1, binIdx = 0, hits = 83, misses = 17)
  t.currentGun = 0
  t.currentSince = 0
  check "hysteresis: a challenger inside the switch margin does not displace the incumbent",
        t.selectGun(7, tick = 100) == 0

  # 80% incumbent vs 90% challenger: beyond the margin -> switch.
  var t2 = initTracker(2)
  seedWindow(t2, 7, gunId = 0, binIdx = 0, hits = 80, misses = 20)
  seedWindow(t2, 7, gunId = 1, binIdx = 0, hits = 90, misses = 10)
  t2.currentGun = 0
  t2.currentSince = 0
  check "hysteresis: a challenger beyond the switch margin displaces the incumbent",
        t2.selectGun(7, tick = 100) == 1

proc testHysteresisDisqualified() =
  var t = initTracker(2)
  seedWindow(t, 7, gunId = 0, binIdx = 0, hits = 10, misses = 0)  # < MinObsBeforeCompete
  seedWindow(t, 7, gunId = 1, binIdx = 0, hits = 60, misses = 0)
  t.currentGun = 0
  t.currentSince = 0
  check "hysteresis: a disqualified incumbent is replaced even inside dwell",
        t.selectGun(7, tick = 3) == 1

proc testHysteresisFloor() =
  var t = initTracker(2)
  seedWindow(t, 7, gunId = 0, binIdx = 0, hits = 0, misses = 60)
  seedWindow(t, 7, gunId = 1, binIdx = 0, hits = 0, misses = 60)
  t.currentGun = 1
  t.currentSince = 0
  t.peakRateRef = 1.0
  check "hysteresis: floor collapse returns HeadOn even mid-dwell",
        t.selectGun(7, tick = 3) == 0

proc testBestPowerCold() =
  var t = initTracker(3)
  let (bin, power) = t.bestPower(0, -1)
  check "bestPower on a zero-observation gun returns bin 0 / power 1.0",
        bin == 0 and power == 1.0

proc testBestPowerWarmBin3() =
  var t = initTracker(3)
  seedWindow(t, 7, gunId = 0, binIdx = 3, hits = 50, misses = 0)  # 100% >= MinHitRate
  let (bin, power) = t.bestPower(0, -1)
  check "bestPower on a warm gun whose bin 3 rate >= MinHitRate returns bin 3",
        bin == 3 and power == 3.0

proc testFitnessForDeterministic() =
  # Same per-enemy data inserted in opposite orders must aggregate identically.
  # Before fitnessFor sorted enemy ids, std/tables hash order leaked in.
  var t1 = initTracker(2)
  seedWindow(t1, 5, gunId = 0, binIdx = 0, hits = 10, misses = 5)
  seedWindow(t1, 3, gunId = 0, binIdx = 0, hits = 5,  misses = 10)

  var t2 = initTracker(2)
  seedWindow(t2, 3, gunId = 0, binIdx = 0, hits = 5,  misses = 10)
  seedWindow(t2, 5, gunId = 0, binIdx = 0, hits = 10, misses = 5)

  var same = true
  for _ in 0..<20:
    let r1 = t1.fitnessFor(-1)[0].bins[0].hitRate()
    let r2 = t2.fitnessFor(-1)[0].bins[0].hitRate()
    if r1 != r2: same = false
  let expected = 15.0 / 30.0
  check "fitnessFor is deterministic across insertion orders",
        same and abs(t1.fitnessFor(-1)[0].bins[0].hitRate() - expected) < 1e-12

proc testDroppedBullets() =
  var t = initTracker(1)
  let state = ws(0, 500.0, 500.0, 0.0, 0.0)
  let preds = [GunPrediction(x: 500.0, y: 500.0),
               GunPrediction(x: 500.0, y: 500.0),
               GunPrediction(x: 500.0, y: 500.0),
               GunPrediction(x: 500.0, y: 500.0)]
  # Fill the ring exactly (4 bullets per spawn, no tickBullets -> never resolve).
  for _ in 0..<(MaxBullets div len(PowerBins)):
    t.spawnBullets(0, preds, state, 5)
  check "droppedBullets stays 0 until the ring wraps", t.droppedBullets == 0
  t.spawnBullets(0, preds, state, 5)
  check "droppedBullets counts unresolved bullets clobbered by the ring",
        t.droppedBullets == 4

# ── caching-bug speed sensitivity (Task 5) ───────────────────────────────────

proc testStopShotSpeedSensitivity() =
  let spd0 = bulletSpeed(PowerBins[0])
  let spd3 = bulletSpeed(PowerBins[3])

  var ss = initStopShotGun()
  # Constant speed 4: warm two frames, then compare on the same tick. Before the
  # fix the tick-only cache returned bin 0's lead for every bin.
  discard ss.predict(ws(1, 400.0, 100.0, 4.0, 0.0), spd0)
  discard ss.predict(ws(2, 400.0, 100.0, 4.0, 0.0), spd0)
  let s3 = ws(3, 400.0, 100.0, 4.0, 0.0)
  let p0 = ss.predict(s3, spd0)
  let p3 = ss.predict(s3, spd3)
  check "stop_shot: same tick, different bulletSpeed -> different point",
        pointsDiffer(p0, p3)

  # Task 1a: deceleration is actually detected (8 -> 4 px/tick). The old ordering
  # made prev == speed, so this branch was unreachable and the gun was Linear.
  var ss2 = initStopShotGun()
  discard ss2.predict(ws(1, 400.0, 100.0, 8.0, 0.0), spd0)
  let pd = ss2.predict(ws(2, 400.0, 100.0, 4.0, 0.0), spd0)
  # Stop point is 400 + 4 + 2 = 406 px (BrakeDecel=2); linear lead would be ~470.
  check "stop_shot: deceleration branch reaches the simulated stop point",
        abs(pd.x - 406.0) < 1.0

proc testDisplacementSpeedSensitivity() =
  let spd0 = bulletSpeed(PowerBins[0])
  let spd3 = bulletSpeed(PowerBins[3])

  var dg = initDisplacementGun()
  # Warm 16 ticks emulating the real harness: 4 predict() calls (one per power
  # bin) on every tick. Feed 16 ticks of constant +5 px/tick motion so the
  # 15-tick window is ready.
  for tick in 1..16:
    for bin in 0..<len(PowerBins):
      discard dg.predict(ws(tick, 300.0 + 5.0 * tick.float, 200.0, 5.0, 0.0),
                         bulletSpeed(PowerBins[bin]))
  let s17 = ws(17, 300.0 + 5.0 * 17.0, 200.0, 5.0, 0.0)
  let d0 = dg.predict(s17, spd0)
  let d3 = dg.predict(s17, spd3)
  check "displacement: same tick, different bulletSpeed -> different point",
        pointsDiffer(d0, d3)

  # The real displacement bug: the speed-in-key cache advanced the ring ~4x per
  # tick, so the nominal 15-tick window spanned ~4 ticks. Sampling once per tick
  # means 4 calls/tick must be identical to 1 call/tick.
  var dgMulti = initDisplacementGun()
  var dgOnce  = initDisplacementGun()
  for tick in 1..16:
    let s = ws(tick, 300.0 + 5.0 * tick.float, 200.0, 5.0, 0.0)
    for bin in 0..<len(PowerBins):
      discard dgMulti.predict(s, bulletSpeed(PowerBins[bin]))
    discard dgOnce.predict(s, spd0)
  let s18 = ws(18, 300.0 + 5.0 * 18.0, 200.0, 5.0, 0.0)
  let a = dgMulti.predict(s18, spd0)
  let b = dgOnce.predict(s18, spd0)
  check "displacement: ring advances exactly once per tick (4 calls == 1 call)",
        not pointsDiffer(a, b)

proc testAveragedLeadSpeedSensitivity() =
  let spd0 = bulletSpeed(PowerBins[0])
  let spd3 = bulletSpeed(PowerBins[3])

  var al = initAveragedLeadGun()
  discard al.predict(ws(1, 400.0, 100.0, 3.0, 0.0), spd0)  # warm circular's omega
  let s2 = ws(2, 400.0, 100.0, 3.0, 0.0)
  let a0 = al.predict(s2, spd0)
  let a3 = al.predict(s2, spd3)
  check "averaged_lead: same tick, different bulletSpeed -> different point",
        pointsDiffer(a0, a3)

proc testPatternMatcherSpeedSensitivity() =
  let spd0 = bulletSpeed(PowerBins[0])
  let spd3 = bulletSpeed(PowerBins[3])

  var pm = PatternMatcherGun()
  for tick in 1..25:
    discard pm.predict(ws(tick, 300.0 + 5.0 * tick.float, 200.0, 5.0, 0.0), spd0)
  let s26 = ws(26, 300.0 + 5.0 * 26.0, 200.0, 5.0, 0.0)
  let m0 = pm.predict(s26, spd0)
  let m3 = pm.predict(s26, spd3)
  check "pattern_matcher: same tick, different bulletSpeed -> different point",
        pointsDiffer(m0, m3)

# ── range-aware firing gate (Task B) ────────────────────────────────────────

proc testToleranceStrictlyDecreases() =
  # In the unclamped band the tolerance must fall monotonically with distance.
  # 100..10000 px sits strictly inside [floor, ceiling] for the shipped
  # SafetyFactor, so clamping cannot mask a flat or rising curve.
  let dists = [100.0, 200.0, 400.0, 800.0, 1600.0, 3200.0, 10000.0]
  var ok = true
  for i in 0..<dists.high:
    if aimToleranceDeg(dists[i]) <= aimToleranceDeg(dists[i+1]): ok = false
  check "range gate: tolerance strictly decreases as distance grows (unclamped)",
        ok

proc testToleranceClamps() =
  let near = aimToleranceDeg(1.0)     # point blank -> ceiling
  let far  = aimToleranceDeg(1.0e9)   # effectively infinite -> floor
  let mid  = aimToleranceDeg(500.0)
  check "range gate: point-blank clamps to MaxAimThresholdDeg",
        near == MaxAimThresholdDeg
  check "range gate: extreme range clamps to MinAimThresholdDeg",
        far == MinAimThresholdDeg
  check "range gate: mid-range tolerance is strictly inside the clamps",
        mid > MinAimThresholdDeg and mid < MaxAimThresholdDeg

proc testToleranceFormula() =
  let d = 500.0
  let expected = radToDeg(arctan(BotRadius * SafetyFactor / d))
  check "range gate: tolerance matches radToDeg(arctan(BotRadius*SF/dist))",
        abs(aimToleranceDeg(d) - expected) < 1e-9

proc testPerfectAlignmentAlwaysFires() =
  var allPass = true
  for d in [0.0, 1.0, 100.0, 1000.0, 1.0e9]:
    if not shouldFire(100.0, 100.0, 0.0, d): allPass = false
  check "range gate: a perfectly aligned cool gun passes at every distance",
        allPass

proc testGrossMisalignmentFailsLongRange() =
  # 5 deg is far outside the ~0.8 deg cone at 800 px.
  check "range gate: gross misalignment fails at long range",
        not shouldFire(100.0, 105.0, 0.0, 800.0)

proc testDegenerateDistance() =
  let t = aimToleranceDeg(0.0)
  check "range gate: distPx = 0 falls back to the ceiling",
        t == MaxAimThresholdDeg
  check "range gate: distPx = 0 yields a finite, non-NaN tolerance",
        t == t and t < Inf and t > -Inf
  check "range gate: NaN distance also falls back to the ceiling",
        aimToleranceDeg(NaN) == MaxAimThresholdDeg

# ── offline range (Task 5) ────────────────────────────────────────────────────

proc testFixtureRoundTrip() =
  let fx = synthesizeConstantVelocity(ticks = 40)
  let path = getTempDir() / "gun_range_roundtrip.jsonl"
  saveFixture(path, fx)
  let back = loadFixture(path)
  check "range: fixture round-trip preserves tick count",
        back.states.len == fx.states.len
  var same = back.states.len == fx.states.len
  for i in 0..<min(back.states.len, fx.states.len):
    let a = fx.states[i]
    let b = back.states[i]
    if a.tick != b.tick or abs(a.enemyX - b.enemyX) > 1e-9 or
       abs(a.enemyY - b.enemyY) > 1e-9 or
       abs(a.enemyHeading - b.enemyHeading) > 1e-9 or
       abs(a.enemySpeed - b.enemySpeed) > 1e-9 or
       abs(a.selfX - b.selfX) > 1e-9:
      same = false
  check "range: fixture round-trip preserves per-tick fields", same
  check "range: fixture round-trip defaults arena to 800x600",
        back.meta.arenaW == 800.0 and back.meta.arenaH == 600.0
  check "range: fixture round-trip preserves enemy id", back.enemyId == fx.enemyId
  removeFile(path)

proc testFixtureEndMarkerRoundTrip() =
  let fx = synthesizeStationary(ticks = 10)
  var withEnd = fx
  withEnd.enemyDied = true
  let path = getTempDir() / "gun_range_end.jsonl"
  saveFixture(path, withEnd)
  let back = loadFixture(path)
  check "range: fixture round-trip preserves the enemyDied end marker", back.enemyDied
  removeFile(path)

proc testReplayDeterminism() =
  let fx = synthesizeCircular(ticks = 80)
  let r1 = replayFixture(fx, @[makeDriver("HeadOn", HeadOnGun()),
                               makeDriver("Linear", LinearGun())])
  let r2 = replayFixture(fx, @[makeDriver("HeadOn", HeadOnGun()),
                               makeDriver("Linear", LinearGun())])
  var same = r1.len == r2.len
  for i in 0..<r1.len:
    if r1[i].name != r2[i].name or r1[i].shots != r2[i].shots or
       r1[i].hits != r2[i].hits:
      same = false
    for b in 0..<len(PowerBins):
      if r1[i].bins[b].shots != r2[i].bins[b].shots or
         r1[i].bins[b].hits != r2[i].bins[b].hits:
        same = false
  check "range: replaying the same fixture twice is byte-identical (deterministic guns)",
        same

proc testRangeGroundTruthStationary() =
  let fx = synthesizeStationary(ticks = 120)
  let r = replayFixture(fx, @[makeDriver("HeadOn", HeadOnGun())])
  # The last bullets are still in flight when the fixture ends, so shots < 480;
  # every resolved shot must still be a hit.
  check "range: stationary enemy -> HeadOn scores 100%",
        r[0].shots > 300 and r[0].hits == r[0].shots

proc testRangeConstantVelocityLinearWins() =
  # Model-specific property: under the point model HeadOn aims at the current
  # position and misses a moving target, while Linear leads it. (Under the
  # shipped path model every gun scores 100% here, so the check pins the
  # metric it was calibrated against.)
  let fx = synthesizeConstantVelocity(ticks = 120)
  let r = replayFixture(fx, @[makeDriver("HeadOn", HeadOnGun()),
                               makeDriver("Linear", LinearGun())],
                        metric = bmPoint)
  check "range: constant velocity -> Linear beats HeadOn (point model)",
        r[1].hitRate() > r[0].hitRate()

proc testEnergyThresholdFixtureRule() =
  # RULE e(t) = max(5, 50 - 0.5*t); e drops below 30 at t = 41.
  let fx = synthesizeEnergyThresholdTurner(ticks = 100, e0 = 50.0,
                                           decay = 0.5, threshold = 30.0)
  let idx = 41
  check "range: energy-threshold fixture crosses the rule threshold at t=41",
        fx.states[idx].enemyEnergy < 30.0 and
        fx.states[idx-1].enemyEnergy >= 30.0

# ── driver ───────────────────────────────────────────────────────────────────

randomize()

testColdBestGun()
testRandomTiebreak()
testBestGunDeterministicWinner()
testHysteresisDwell()
testHysteresisSwitchMargin()
testHysteresisDisqualified()
testHysteresisFloor()
testBestPowerCold()
testBestPowerWarmBin3()
testFitnessForDeterministic()
testDroppedBullets()
testStopShotSpeedSensitivity()
testDisplacementSpeedSensitivity()
testAveragedLeadSpeedSensitivity()
testPatternMatcherSpeedSensitivity()
testToleranceStrictlyDecreases()
testToleranceClamps()
testToleranceFormula()
testPerfectAlignmentAlwaysFires()
testGrossMisalignmentFailsLongRange()
testDegenerateDistance()
testFixtureRoundTrip()
testFixtureEndMarkerRoundTrip()
testReplayDeterminism()
testRangeGroundTruthStationary()
testRangeConstantVelocityLinearWins()
testEnergyThresholdFixtureRule()

if failures > 0:
  echo "\n", failures, " check(s) FAILED"
  quit(1)
echo "\nAll gun-harness checks passed."