## AUDIT TOOL (new file, read-only w.r.t. shipped source).
##
## Measures, over the committed fixture set:
##   1. speed usage      — does each gun emit 4 DIFFERENT predictions for the 4
##                         PowerBins, or does it ignore bulletSpeed?
##   2. lead / accuracy  — bmPoint miss distance (arrival-time aim error) and the
##                         bmPath "virtual" rate, so the optimism ratio is explicit.
##   3. pairwise agreement of predictions (redundancy).
##   4. conditional point hit rate by distance / speed / wall / reversal.
##   5. feedback-integrity counters (waveStarved / traceMisses / trained).
##
## Usage:
##   nim c -r --path:common_libs common_libs/tests/audit_virtual_guns.nim

import std/[os, math, strformat, tables, algorithm, strutils, json, random]
import gun_harness/offline_range
import gun_harness/gun_interface
import gun_harness/virtual_bullets as vb
import guns/head_on
import guns/linear
import guns/tsetlin
import guns/circular
import guns/guess_factor
import guns/pattern_matcher
import guns/wall_bounce
import guns/accel_predictor
import guns/stop_shot
import guns/displacement
import guns/averaged_lead
import guns/decay_gf
import guns/knn_gun
import guns/tm_selector

const
  NGun* = 14
  GunNames*: array[NGun, string] = ["HeadOn","Linear","Tsetlin","Circular",
    "GuessFactor","Pattern","WallBounce","Accel","StopShot","Displace","AvgLead",
    "DecayGF","KNN","TMSelect"]
  NBins = len(vb.PowerBins)

type
  Box[G] = ref object
    g: G

  Drivers* = object
    d*: seq[GunDriver]
    headOn*:   Box[HeadOnGun]
    linear*:   Box[LinearGun]
    tsetlin*:  Box[TsetlinGun]
    circular*: Box[CircularGun]
    gf*:       Box[GFGun]
    pattern*:  Box[PatternMatcherGun]
    wall*:     Box[WallBounceGun]
    accel*:    Box[AccelGun]
    stop*:     Box[StopShotGun]
    disp*:     Box[DisplacementGun]
    avg*:      Box[AveragedLeadGun]
    decay*:    Box[DecayGFGun]
    knn*:      Box[KNNGun]
    tmsel*:    Box[TmSelectorGun]

proc mk[G](name: string, g: Box[G]): GunDriver =
  result.name = name
  result.predictCb = proc(s: WorldState, sp: float): GunPrediction = g.g.predict(s, sp)
  result.resultCb  = proc(e: FeedbackEvent) = g.g.onResult(e)
  when compiles(g.g.isWarmedUp()):
    result.readyCb = proc(): bool = g.g.isWarmedUp()
  else:
    result.readyCb = nil

proc newBox[G](v: G): Box[G] =
  new(result); result.g = v

proc initDrivers*(seed = 1): Drivers =
  result.headOn   = newBox(HeadOnGun())
  result.linear   = newBox(LinearGun())
  result.tsetlin  = newBox(initTsetlinGun())
  result.circular = newBox(CircularGun())
  result.gf       = newBox(initGFGun())
  result.pattern  = newBox(PatternMatcherGun())
  result.wall     = newBox(initWallBounceGun())
  result.accel    = newBox(initAccelGun())
  result.stop     = newBox(initStopShotGun())
  result.disp     = newBox(initDisplacementGun())
  result.avg      = newBox(initAveragedLeadGun())
  result.decay    = newBox(initDecayGFGun())
  result.knn      = newBox(initKNNGun())
  result.tmsel    = newBox(initTmSelectorGun())
  if seed >= 0: randomize(seed)
  result.d = @[
    mk("HeadOn", result.headOn), mk("Linear", result.linear),
    mk("Tsetlin", result.tsetlin), mk("Circular", result.circular),
    mk("GuessFactor", result.gf), mk("Pattern", result.pattern),
    mk("WallBounce", result.wall), mk("Accel", result.accel),
    mk("StopShot", result.stop), mk("Displace", result.disp),
    mk("AvgLead", result.avg), mk("DecayGF", result.decay),
    mk("KNN", result.knn), mk("TMSelect", result.tmsel),
  ]

# ── accumulators ──────────────────────────────────────────────────────────────

type
  BinAcc = object
    count*, hits*: int
    sumMiss*: float
    within20*, within50*: int
    sumPairDist*: float

  ContextAcc = object
    count*, hits*: int

  GunAcc = object
    bins: array[NBins, BinAcc]
    # point-metric results (arrival-time aim accuracy)
    ptCount*, ptHits*: int
    ptSumMiss*: float
    # prediction-diversity
    identicalTicks*, predTicks*: int
    sumSpread*: float            ## mean over ticks of max pairwise pred distance
    # conditional point hit rate (pooled bins), marginal buckets
    dist*: array[4, ContextAcc]
    speed*: array[4, ContextAcc]
    wall*: array[2, ContextAcc]
    rev*: array[2, ContextAcc]
    # virtual-path rate
    vpCount*, vpHits*: int

  Audit* = object
    guns: array[NGun, GunAcc]
    pairCount: array[NGun, array[NGun, int]]
    pairWithin20: array[NGun, array[NGun, int]]
    pairWithin50: array[NGun, array[NGun, int]]
    pairSumDist: array[NGun, array[NGun, float]]
    ticks*: int
    fixtures*: int

proc ctxDist(s: WorldState): int =
  let d = hypot(s.enemyX - s.selfX, s.enemyY - s.selfY)
  if d < 200.0: 0 elif d < 400.0: 1 elif d < 600.0: 2 else: 3

proc ctxSpeed(s: WorldState): int =
  let sp = abs(s.enemySpeed)
  if sp < 1.0: 0 elif sp < 4.0: 1 elif sp < 7.5: 2 else: 3

proc ctxWall(s: WorldState): int =
  let m = min(min(s.enemyX, s.arenaWidth - s.enemyX),
              min(s.enemyY, s.arenaHeight - s.enemyY))
  if m < 60.0: 1 else: 0

proc feedPredictions(a: var Audit, gi: int, preds: array[NBins, GunPrediction]) =
  ## Speed usage: are the 4 bin predictions distinct?
  var maxD = 0.0
  for i in 0..<NBins:
    for j in i+1..<NBins:
      let d = hypot(preds[i].x - preds[j].x, preds[i].y - preds[j].y)
      if d > maxD: maxD = d
  inc a.guns[gi].predTicks
  a.guns[gi].sumSpread += maxD
  if maxD < 0.01: inc a.guns[gi].identicalTicks

proc feedPairs(a: var Audit, preds: array[NGun, array[NBins, GunPrediction]]) =
  for i in 0..<NGun:
    for j in i+1..<NGun:
      for b in 0..<NBins:
        let d = hypot(preds[i][b].x - preds[j][b].x, preds[i][b].y - preds[j][b].y)
        inc a.pairCount[i][j]
        a.pairSumDist[i][j] += d
        if d < 20.0: inc a.pairWithin20[i][j]
        if d < 50.0: inc a.pairWithin50[i][j]

proc auditReplay*(fx: Fixture, drivers: Drivers, metric: BulletMetric,
                  a: ref Audit, capPreds: bool) =
  ## One replay pass. `metric` selects the virtual scoring model. Captures
  ## bmPoint FeedbackEvents into the point/context accumulators; bmPath into the
  ## virtual-rate accumulators; prediction diversity/pairwise always.
  let n = drivers.d.len
  let tid = fx.enemyId
  a.ticks += fx.states.len
  a.fixtures += 1
  var tracker = vb.initTracker(n, metric)
  var byTick = initTable[int, int]()
  for si in 0..<fx.states.len: byTick[fx.states[si].tick] = si
  var prevSpeed = 0.0

  for si in 0..<fx.states.len:
    let state = fx.states[si]
    var allPreds: array[NGun, array[NBins, GunPrediction]]
    for gi in 0..<n:
      var preds: array[NBins, GunPrediction]
      for i in 0..<NBins:
        preds[i] = drivers.d[gi].predictCb(state, bulletSpeed(vb.PowerBins[i]))
      allPreds[gi] = preds
      if capPreds: feedPredictions(a[], gi, preds)
      let ready = if drivers.d[gi].readyCb == nil: true else: drivers.d[gi].readyCb()
      if ready:
        tracker.spawnBullets(gi, preds, state, tid)
    if capPreds: feedPairs(a[], allPreds)

    let act = state
    var enemyPositions: Table[int, tuple[x, y: float, lastSeenTick: int, alive: bool]]
    for e in act.enemies:
      enemyPositions[e.id] = (x: e.x, y: e.y, lastSeenTick: e.lastSeenTick, alive: true)

    prevSpeed = state.enemySpeed

    tracker.tickBullets(state, enemyPositions,
      proc(gunId: GunId, binIdx: int, e: FeedbackEvent) =
        if gunId < 0 or gunId >= NGun: return
        drivers.d[gunId].resultCb(e)   # feed the owning gun its own feedback
        if metric == bmPoint:
          let ga = addr a.guns[gunId]
          inc ga.ptCount
          ga.ptSumMiss += e.missDistance
          if e.hit: inc ga.ptHits
          # context buckets (from fire-time state)
          let fsIdx = byTick.getOrDefault(e.fireTick, si)
          let fs = fx.states[fsIdx]
          inc ga.dist[ctxDist(fs)].count
          inc ga.speed[ctxSpeed(fs)].count
          inc ga.wall[ctxWall(fs)].count
          let r = if fsIdx > 0:
                    let ps = fx.states[fsIdx-1].enemySpeed
                    ps != 0.0 and fs.enemySpeed != 0.0 and (ps > 0.0) != (fs.enemySpeed > 0.0)
                  else: false
          inc ga.rev[ord(r)].count
          if e.hit:
            inc ga.dist[ctxDist(fs)].hits
            inc ga.speed[ctxSpeed(fs)].hits
            inc ga.wall[ctxWall(fs)].hits
            inc ga.rev[ord(r)].hits
        else:
          inc a.guns[gunId].vpCount
          if e.hit: inc a.guns[gunId].vpHits
        discard binIdx)

proc addGunAcc(dst: var GunAcc, src: GunAcc) =
  for b in 0..<NBins:
    dst.bins[b].count += src.bins[b].count
    dst.bins[b].hits += src.bins[b].hits
    dst.bins[b].sumMiss += src.bins[b].sumMiss
    dst.bins[b].within20 += src.bins[b].within20
    dst.bins[b].within50 += src.bins[b].within50
    dst.bins[b].sumPairDist += src.bins[b].sumPairDist
  dst.ptCount += src.ptCount; dst.ptHits += src.ptHits; dst.ptSumMiss += src.ptSumMiss
  dst.identicalTicks += src.identicalTicks; dst.predTicks += src.predTicks
  dst.sumSpread += src.sumSpread
  for k in 0..<4:
    dst.dist[k].count += src.dist[k].count; dst.dist[k].hits += src.dist[k].hits
    dst.speed[k].count += src.speed[k].count; dst.speed[k].hits += src.speed[k].hits
  for k in 0..<2:
    dst.wall[k].count += src.wall[k].count; dst.wall[k].hits += src.wall[k].hits
    dst.rev[k].count += src.rev[k].count; dst.rev[k].hits += src.rev[k].hits
  dst.vpCount += src.vpCount; dst.vpHits += src.vpHits

proc rate(h, n: int): float = (if n == 0: 0.0 else: h.float / n.float * 100.0)

# ── feedback-integrity counters ───────────────────────────────────────────────
type
  CounterAcc* = object
    gfPush*, gfStarve*, decPush*, decStarve*, knnPush*, knnStarve*: int
    tmTrained*, tmMisses*, tmPredict*, tmCorr*: int
    selTrain*, selMiss*, selObs*, selPredict*: int

var ctrPath: CounterAcc

proc collectCounters(d: Drivers, c: var CounterAcc) =
  c.gfPush     += d.gf.g.wavePushes;     c.gfStarve  += d.gf.g.waveStarved
  c.decPush    += d.decay.g.wavePushes;  c.decStarve += d.decay.g.waveStarved
  c.knnPush    += d.knn.g.wavePushes;    c.knnStarve += d.knn.g.waveStarved
  c.tmTrained  += d.tsetlin.g.trainedShots
  c.tmMisses   += d.tsetlin.g.traceMisses
  c.tmPredict  += d.tsetlin.g.predictCalls
  c.tmCorr     += d.tsetlin.g.correctionsNonzero
  c.selTrain   += d.tmsel.g.trainCalls
  c.selMiss    += d.tmsel.g.traceMisses
  c.selObs     += d.tmsel.g.totalObs
  c.selPredict += d.tmsel.g.predictCalls

# ── main ──────────────────────────────────────────────────────────────────────

const fixturesDir = currentSourcePath().parentDir.parentDir.parentDir / "tools" / "fixtures"

var
  accPoint: Audit
  accPath: Audit

proc runFile(path: string) =
  var fx = loadFixture(path)
  if fx.states.len < 50: return
  let MaxTicks = (if existsEnv("AUDIT_MAX_TICKS"): parseInt(getEnv("AUDIT_MAX_TICKS"))
                  else: 4000)
  if fx.states.len > MaxTicks:
    fx.states.setLen(MaxTicks)
    if fx.lastSeen.len > MaxTicks: fx.lastSeen.setLen(MaxTicks)
  let dPoint = initDrivers(seed = 1)
  var a = new(Audit)
  auditReplay(fx, dPoint, bmPoint, a, capPreds = false)
  for gi in 0..<NGun: accPoint.guns[gi].addGunAcc(a.guns[gi])

  let dPath = initDrivers(seed = 1)
  var b = new(Audit)
  auditReplay(fx, dPath, bmPath, b, capPreds = true)
  collectCounters(dPath, ctrPath)
  accPath.ticks += b.ticks
  accPath.fixtures += b.fixtures
  for gi in 0..<NGun:
    accPath.guns[gi].addGunAcc(b.guns[gi])
  for i in 0..<NGun:
    for j in 0..<NGun:
      accPath.pairCount[i][j] += b.pairCount[i][j]
      accPath.pairWithin20[i][j] += b.pairWithin20[i][j]
      accPath.pairWithin50[i][j] += b.pairWithin50[i][j]
      accPath.pairSumDist[i][j] += b.pairSumDist[i][j]
  echo fmt"# {path} ticks={fx.states.len}"

proc printCounters() =
  echo ""
  echo "=================== FEEDBACK INTEGRITY (bmPath pass) ==================="
  echo fmt"GF      pushes={ctrPath.gfPush:<8} starved={ctrPath.gfStarve}"
  echo fmt"DecayGF pushes={ctrPath.decPush:<8} starved={ctrPath.decStarve}"
  echo fmt"KNN     pushes={ctrPath.knnPush:<8} starved={ctrPath.knnStarve}"
  echo fmt"Tsetlin trained={ctrPath.tmTrained:<8} misses={ctrPath.tmMisses:<8} predicts={ctrPath.tmPredict:<8} corrNonzero={ctrPath.tmCorr}"
  echo fmt"TMSelect trained={ctrPath.selTrain:<8} misses={ctrPath.selMiss:<8} obs={ctrPath.selObs:<8} predicts={ctrPath.selPredict}"

proc main() =
  var files: seq[string]
  for i in 1..paramCount(): files.add paramStr(i)
  if files.len == 0:
    for f in walkFiles(fixturesDir / "tr_drussgt_*.jsonl"): files.add f
    for f in walkFiles(fixturesDir / "drussgt_*.jsonl"): files.add f
  files.sort()
  for f in files:
    runFile(f)

  printCounters()

  echo ""
  echo "=================== PER-GUN SUMMARY (all fixtures) ==================="
  echo fmt"fixtures={accPath.fixtures} ticks={accPath.ticks}"
  echo "gun          vpath%   point%  ratio  ident%  spread  leadPx  ptMiss"
  for gi in 0..<NGun:
    let g = addr accPath.guns[gi]
    let vp = rate(g.vpHits, g.vpCount)
    let pp = rate(accPoint.guns[gi].ptHits, accPoint.guns[gi].ptCount)
    let ratio = if pp > 0.0: vp / pp else: 0.0
    let ident = rate(g.identicalTicks, g.predTicks)
    let spread = if g.predTicks > 0: g.sumSpread / g.predTicks.float else: 0.0
    let ptMiss = if accPoint.guns[gi].ptCount > 0:
                   accPoint.guns[gi].ptSumMiss / accPoint.guns[gi].ptCount.float else: 0.0
    echo fmt"{GunNames[gi]:<12} {vp:6.1f}  {pp:6.1f}  {ratio:5.2f}  {ident:6.1f}  {spread:6.1f}  {ptMiss:6.1f}"
  echo "(vpath% = shipped bmPath virtual rate; point% = bmPoint arrival-accuracy rate;"
  echo " ident% = ticks whose 4 power-bin predictions are identical; spread = mean max"
  echo " pairwise bin distance px; ptMiss = mean bmPoint miss distance px)"

  echo ""
  echo "=================== TOP REDUNDANT PAIRS (pred within 20px / 50px) ==================="
  type Pair = tuple[pct20, pct50, meanD: float, i, j: int]
  var pairs: seq[Pair]
  for i in 0..<NGun:
    for j in i+1..<NGun:
      let n = accPath.pairCount[i][j]
      if n == 0: continue
      pairs.add (pct20: rate(accPath.pairWithin20[i][j], n),
                 pct50: rate(accPath.pairWithin50[i][j], n),
                 meanD: accPath.pairSumDist[i][j] / n.float, i: i, j: j)
  pairs.sort(proc(x, y: Pair): int = cmp(y.pct20, x.pct20))
  for p in pairs[0..<min(18, pairs.len)]:
    echo fmt"{GunNames[p.i]:<11} ~ {GunNames[p.j]:<11}  <20px {p.pct20:5.1f}%  <50px {p.pct50:5.1f}%  mean {p.meanD:6.1f}px"

  echo ""
  echo "===== FULL PAIR MATRIX (% ticks 4-bin preds within 20px) ====="
  var hdr = "            "
  for j in 0..<NGun: hdr.add fmt"{GunNames[j][0..<min(6,GunNames[j].len)]:>8}"
  echo hdr
  for i in 0..<NGun:
    var row = fmt"{GunNames[i]:<12}"
    for j in 0..<NGun:
      if i == j: row.add "      --"
      elif j < i: row.add "        "
      else:
        let n = accPath.pairCount[i][j]
        row.add fmt"{rate(accPath.pairWithin20[i][j], n):8.1f}"
    echo row

  echo ""
  echo "=================== CONDITIONAL bmPoint HIT RATE ==================="
  echo "gun          d<200    200-400   400-600   600+    v<1      v1-4    v4-7.5   v>=7.5  wall  open  rev   fwd"
  for gi in 0..<NGun:
    let g = addr accPoint.guns[gi]
    var line = fmt"{GunNames[gi]:<12}"
    for k in 0..<4: line.add fmt"{rate(g.dist[k].hits, g.dist[k].count):6.1f}  "
    for k in 0..<4: line.add fmt"{rate(g.speed[k].hits, g.speed[k].count):6.1f}  "
    line.add fmt"{rate(g.wall[1].hits, g.wall[1].count):6.1f} "
    line.add fmt"{rate(g.wall[0].hits, g.wall[0].count):6.1f} "
    line.add fmt"{rate(g.rev[1].hits, g.rev[1].count):6.1f} "
    line.add fmt"{rate(g.rev[0].hits, g.rev[0].count):6.1f}"
    echo line

  echo ""
  echo "=================== SAMPLE SIZES (bmPoint per-gun) ==================="
  for gi in 0..<NGun:
    let g = addr accPoint.guns[gi]
    echo fmt"{GunNames[gi]:<12} pointN={g.ptCount:<7} vpathN={accPath.guns[gi].vpCount:<7} predTicks={accPath.guns[gi].predTicks}"

when isMainModule:
  main()