## TMComposites GATE — offline per-sample confidence + outcome recorder.
##
## Implements the measurement half of docs/tmcomposites_gate.md. It replays
## recorded fixtures through the SHIPPED rack exactly as `offline_range` does
## (same `VirtualTracker`, same `bmPath` virtual-bullet ground truth), but also
## captures, for every resolved virtual bullet, the gun's INTRINSIC per-sample
## `GunPrediction.confidence` (see gun_interface.nim) and the bullet's aim
## bearing relative to the fire-time line of sight.
##
## It also runs two COMPOSITE arms through the SAME tracker, so their hits are
## scored by the identical geometry as every member:
##   * Composite      — TMComposites Eq 8: each confident gun casts its
##                      alpha-normalised confidence into the angular bin of its
##                      own aim; the argmax bin wins.
##   * CompositeShuf  — the MANDATORY control: the same confidences are randomly
##                      permuted among the members within each sample, so the
##                      weighting distribution is preserved but competence is
##                      destroyed.
##
## alpha (Eq 7) is calibrated on the TRAIN fixtures (per member: max-min of its
## confidence over train) and then frozen for the TEST fixtures, so the test
## composite never sees test labels while choosing its weights.
##
## Output: one JSONL row per resolved bullet, to --out. Aggregation lives in
## common_libs/tests/analyze_tmcomposites.py.
##
## Usage:
##   nim c -r --nimcache:/tmp/nc_j104 common_libs/tests/measure_tmcomposites.nim \
##       --out /tmp/tmc.jsonl \
##       --train fx_a.jsonl fx_b.jsonl --test fx_c.jsonl

import std/[os, strformat, json, math, strutils, random, algorithm, tables, times]
import gun_harness/offline_range
import gun_harness/gun_interface
import gun_harness/virtual_bullets
import range_guns

const
  BinDeg      = 0.5
  HalfSpanDeg = 45.0
  NumBins     = int(2.0 * HalfSpanDeg / BinDeg)
  ## guns with a genuine intrinsic confidence signal (deterministic geometric
  ## guns leave GunPrediction.confidence at 0.0 and cast no composite vote).
  ConfidentGuns = ["Tsetlin", "GuessFactor", "Pattern", "DecayGF", "KNN"]

type
  Spawn = object
    relDeg: float
    range:  float
    conf:   float

proc wrapDeg(d: float): float =
  result = d
  while result > 180.0: result -= 360.0
  while result < -180.0: result += 360.0

proc binOf(relDeg: float): int =
  result = int(floor((relDeg + HalfSpanDeg) / BinDeg))
  if result < 0: result = 0
  elif result >= NumBins: result = NumBins - 1

var alphaByName: ref Table[string, float]

proc alphaFor(name: string): float =
  ## Eq 7 alpha_t for a member, keyed by GUN NAME so each composite's member
  ## order cannot mis-assign a scale. 1.0 until calibration writes the table.
  if alphaByName.isNil: return 1.0
  alphaByName[].getOrDefault(name, 1.0)

proc makeComposite(name: string, members: seq[GunDriver],
                   shuffle: bool): GunDriver =
  ## One composite arm. `members` are the SAME driver closures the tracker uses
  ## for the member guns, so the composite reads their live state (calling
  ## predict twice in a tick is idempotent for every gun: GF/KNN guard their
  ## wave queue on (tick,bin), Pattern/Tsetlin/TMHorizon cache per tick).
  result.name = name
  let memberList = members
  result.predictCb = proc(state: WorldState, speed: float): GunPrediction =
    let n = memberList.len
    var rels = newSeq[float](n)
    var dists = newSeq[float](n)
    var confs = newSeq[float](n)
    let los = arctan2(state.enemyY - state.selfY, state.enemyX - state.selfX)
    for i in 0..<n:
      let p = memberList[i].predictCb(state, speed)
      rels[i]  = wrapDeg(radToDeg(arctan2(p.y - state.selfY, p.x - state.selfX) - los))
      dists[i] = hypot(p.x - state.selfX, p.y - state.selfY)
      confs[i] = p.confidence
    if shuffle:
      for i in countdown(n - 1, 1):
        let j = rand(i)
        swap(confs[i], confs[j])

    var votes = newSeq[float](NumBins)
    var total = 0.0
    for i in 0..<n:
      let a = max(1e-12, alphaFor(memberList[i].name))
      let w = confs[i] / a
      if w <= 0.0: continue
      votes[binOf(rels[i])] += w
      total += w
    if total <= 0.0:
      # cold: no member has any confidence yet. Make no claim.
      return GunPrediction(x: state.enemyX, y: state.enemyY, confidence: 0.0)

    var best = 0
    for b in 1..<NumBins:
      if votes[b] > votes[best]: best = b

    # aim distance = confidence-weighted mean distance of the winning bin's voters
    var dsum = 0.0
    var wsum = 0.0
    for i in 0..<n:
      if binOf(rels[i]) != best: continue
      let a = max(1e-12, alphaFor(memberList[i].name))
      let w = confs[i] / a
      dsum += dists[i] * w
      wsum += w
    let d = if wsum > 1e-12: dsum / wsum else: hypot(state.enemyX - state.selfX,
                                                     state.enemyY - state.selfY)
    let ang = los + degToRad(-HalfSpanDeg + (best.float + 0.5) * BinDeg)
    GunPrediction(x: state.selfX + cos(ang) * d,
                  y: state.selfY + sin(ang) * d,
                  confidence: votes[best])
  result.resultCb = proc(e: FeedbackEvent) = discard
  result.readyCb = nil

proc buildRack(): tuple[drivers: seq[GunDriver], memberLocal: seq[int]] =
  ## Fresh members + the composite arms, mirroring the standard per-fixture
  ## replay (guns start cold for every fixture, exactly as run_range does).
  ##
  ##   Composite      — all 5 confidence guns (Tsetlin, GF, Pattern, DecayGF, KNN)
  ##   CompositeShuf  — its within-sample confidence shuffle control
  ##   CompositeF     — only the guns the faithfulness test finds FAITHFUL
  ##                    (Pattern, DecayGF, KNN); GF is anti-faithful and Tsetlin
  ##                    useless, so this is the strongest reasonable variant
  ##   CompositeFShuf — its shuffle control
  let allDrivers = buildAllGunDrivers(seed = 1)
  var members: seq[GunDriver]
  var memberLocal: seq[int]
  var faithMembers: seq[GunDriver]
  for i, d in allDrivers:
    if d.name in ConfidentGuns:
      members.add d
      memberLocal.add i
    if d.name in ["Pattern", "DecayGF", "KNN"]:
      faithMembers.add d
  doAssert members.len == ConfidentGuns.len
  result.drivers = allDrivers
  result.drivers.add makeComposite("Composite", members, shuffle = false)
  result.drivers.add makeComposite("CompositeShuf", members, shuffle = true)
  result.drivers.add makeComposite("CompositeF", faithMembers, shuffle = false)
  result.drivers.add makeComposite("CompositeFShuf", faithMembers, shuffle = true)
  result.memberLocal = memberLocal

proc runFixture(fx: Fixture, fixtureName, split: string,
                calMin, calMax: ref seq[float], outFile: File) =
  let (drivers, memberLocal) = buildRack()
  var byTick = initTable[int, WorldState]()
  for s in fx.states: byTick[s.tick] = s
  var spawns = initTable[tuple[gunId, tick, bin: int], Spawn]()
  var tracker = initTracker(drivers.len, ActiveMetric)

  for si in 0..<fx.states.len:
    let state = fx.states[si]
    for gi in 0..<drivers.len:
      var preds: array[len(PowerBins), GunPrediction]
      var sps:   array[len(PowerBins), Spawn]
      let los = arctan2(state.enemyY - state.selfY, state.enemyX - state.selfX)
      for i in 0..<len(PowerBins):
        preds[i] = drivers[gi].predictCb(state, bulletSpeed(PowerBins[i]))
        sps[i] = Spawn(
          relDeg: wrapDeg(radToDeg(arctan2(preds[i].y - state.selfY,
                                           preds[i].x - state.selfX) - los)),
          range:  hypot(preds[i].x - state.selfX, preds[i].y - state.selfY),
          conf:   preds[i].confidence)
      let ready = if drivers[gi].readyCb == nil: true else: drivers[gi].readyCb()
      if ready:
        for i in 0..<len(PowerBins):
          spawns[(gi, state.tick, i)] = sps[i]
        tracker.spawnBullets(gi, preds, state, fx.enemyId)

    var enemyPositions: Table[int, tuple[x, y: float, lastSeenTick: int, alive: bool]]
    if state.enemies.len > 0:
      for e in state.enemies:
        enemyPositions[e.id] = (x: e.x, y: e.y, lastSeenTick: e.lastSeenTick, alive: true)
    else:
      enemyPositions[fx.enemyId] = (x: state.enemyX, y: state.enemyY,
                                    lastSeenTick: state.tick, alive: true)

    tracker.tickBullets(state, enemyPositions,
      proc(gunId: int, binIdx: int, e: FeedbackEvent) =
        # The gun must LEARN from its own resolved bullet (exactly as
        # offline_range.replayFixture does); the recorder is an extra hook.
        drivers[gunId].resultCb(e)
        let key = (gunId, e.fireTick, binIdx)
        if not spawns.hasKey(key): return
        let sp = spawns[key]
        spawns.del(key)
        if split == "train":
          for mi, gid in memberLocal:
            if gunId == gid:
              calMin[][mi] = min(calMin[][mi], sp.conf)
              calMax[][mi] = max(calMax[][mi], sp.conf)
        outFile.writeLine($(%*{
          "fixture": fixtureName,
          "split":   split,
          "gun":     drivers[gunId].name,
          "tick":    e.fireTick,
          "bin":     binIdx,
          "conf":    sp.conf,
          "hit":     e.hit,
          "relDeg":  sp.relDeg,
          "range":   sp.range,
          "missPx":  e.missDistance,
        }))
    )

proc loadAll(paths: seq[string]): seq[tuple[name: string, fx: Fixture]] =
  for p in paths:
    let fx = loadFixture(p)
    let name = extractFilename(p).replace(".jsonl", "")
    result.add (name: name, fx: fx)

proc main() =
  var outPath = "/tmp/tmc.jsonl"
  var trainPaths, testPaths: seq[string]
  var args: seq[string]
  for i in 1..paramCount(): args.add paramStr(i)
  var mode = ""
  for a in args:
    if a == "--out": mode = "out"; continue
    if a == "--train": mode = "train"; continue
    if a == "--test": mode = "test"; continue
    case mode
    of "out": outPath = a
    of "train": trainPaths.add a
    of "test": testPaths.add a
    else: discard

  var calMinRef = new(seq[float])
  var calMaxRef = new(seq[float])
  calMinRef[] = newSeq[float](ConfidentGuns.len)
  calMaxRef[] = newSeq[float](ConfidentGuns.len)
  for i in 0..<ConfidentGuns.len:
    calMinRef[][i] = 1e18
    calMaxRef[][i] = -1e18

  let trainFx = loadAll(trainPaths)
  let testFx = loadAll(testPaths)
  let outFile = open(outPath, fmWrite)
  defer: outFile.close()

  randomize(20250925)

  echo fmt"# TMComposites recorder: {trainFx.len} train fixtures, {testFx.len} test fixtures"
  echo fmt"# members: {ConfidentGuns}"

  var t0 = epochTime()
  for (name, fx) in trainFx:
    let a = epochTime()
    runFixture(fx, name, "train", calMinRef, calMaxRef, outFile)
    echo fmt"  [train] {name:<28} ticks={fx.states.len:>6}  {epochTime()-a:6.1f}s"

  # Eq 7 alpha_t = max-min over the training input set; frozen for test.
  alphaByName = new(Table[string, float])
  for i in 0..<ConfidentGuns.len:
    alphaByName[][ConfidentGuns[i]] =
      max(1e-9, calMaxRef[][i] - min(calMinRef[][i], 0.0))
  echo "# alpha_t (Eq 7, from train): " &
       (block:
          var s = ""
          for i in 0..<ConfidentGuns.len:
            s.add fmt"{ConfidentGuns[i]}={alphaByName[][ConfidentGuns[i]]:.4g} "
          s)

  for (name, fx) in testFx:
    let a = epochTime()
    runFixture(fx, name, "test", calMinRef, calMaxRef, outFile)
    echo fmt"  [test ] {name:<28} ticks={fx.states.len:>6}  {epochTime()-a:6.1f}s"

  echo fmt"# done in {epochTime()-t0:.1f}s -> {outPath}"

when isMainModule:
  main()