SirRoboGarage/common_libs/tests/sweep_tm_pattern.nim

## Discrete-target Tsetlin gun sweep: tm_pattern vs its shuffled-feedback
## control, the current default Tsetlin gun, and Linear.
##
## Metric: EARLY virtual hit rate = resolutions in the first 100 ticks of each
## ROUND (the TM starts cold every round, so this is what "learns fast" means),
## plus first-300 and whole-round rates. Every fixture with a round sidecar is
## split into rounds and each round is replayed with a FRESH gun instance
## (cold every battle, overfit within the battle).
##
## Usage:
##   nim c -r --path:common_libs -d:release common_libs/tests/sweep_tm_pattern.nim \
##       --set=real --seeds=3
## Flags:
##   --set=real|range|synthetic   fixture set (default real)
##   --seeds=N                    seeds for stochastic variants (default 3)
##   --maxrounds=N                cap rounds per fixture
##   --variants=a,b,c             subset of: linear,tsetlin,tmpat,tmpat_shuf
##
## Emits per-run CSV then a COMPARE section with per-run means, ranges and a
## paired sign test (exact binomial, two-sided) for every variant pair.

import std/[os, strformat, strutils, json, tables, math, random, algorithm, sequtils]
import gun_harness/offline_range
import range_guns
import guns/tm_pattern
import guns/linear
import guns/lead_forecast

type
  LinearInsetGun = object  ## pre-fix TMPattern base: same forecast, BotRadius clamp

proc predict*(g: var LinearInsetGun, state: WorldState, bulletSpeed: float): GunPrediction =
  let f = forecastLinear(state, bulletSpeed)
  GunPrediction(x: clamp(f.x, BotRadius, state.arenaWidth - BotRadius),
                y: clamp(f.y, BotRadius, state.arenaHeight - BotRadius))

proc onResult*(g: var LinearInsetGun, e: FeedbackEvent) = discard

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

type
  GunStats = object
    obs, labelMiss, traceMiss: int
    labHist, choHist: array[TM_CLASSES, int]
    classCorrect, classTotal: int
    radCorrect, radTotal, revCorrect, revTotal: int
    confusion: array[TM_CLASSES, array[TM_CLASSES, int]]
    radConfusion: array[TM_CLASSES, array[TM_CLASSES, int]]

  Adapt = object
    h100, n100, h300, n300, hall, nall, f100, m100: int
    rounds: int
    st: GunStats

  RoundSpan = tuple[start, count: int]

  Row = object
    variant, fixture: string
    seed: int
    r: Adapt

  VariantKind = enum
    vLinear, vLinearOld, vTsetlin, vTmpat, vTmpatShuf, vTmpatBase,
    vTmpatRad, vTmpatRadShuf, vTmpatRev, vTmpatRevShuf

proc variantName(v: VariantKind): string =
  case v
  of vLinear: "Linear"
  of vLinearOld: "LinearOldClamp"
  of vTsetlin: "Tsetlin"
  of vTmpat: "TMPattern"
  of vTmpatShuf: "TMPatternShuf"
  of vTmpatBase: "TMPatternBase"
  of vTmpatRad: "TMRadial"
  of vTmpatRadShuf: "TMRadialShuf"
  of vTmpatRev: "TMReversal"
  of vTmpatRevShuf: "TMReversalShuf"

proc loadRounds(path: string): seq[RoundSpan] =
  let dir = path.parentDir
  let base = path.extractFilename
  var side = dir / "drussgt_meta" / (base & ".rounds.json")
  if not fileExists(side): side = dir / (base & ".rounds.json")
  if not fileExists(side): return @[]
  let node = parseJson(readFile(side))
  if not node.hasKey("rounds"): return @[]
  for r in node["rounds"]:
    result.add (r["startTick"].getInt(), r["count"].getInt())

proc addAdapt(dst: var Adapt, src: Adapt) =
  inc dst.rounds, src.rounds
  dst.h100 += src.h100; dst.n100 += src.n100
  dst.h300 += src.h300; dst.n300 += src.n300
  dst.hall += src.hall; dst.nall += src.nall
  dst.f100 += src.f100; dst.m100 += src.m100
  dst.st.obs += src.st.obs; dst.st.labelMiss += src.st.labelMiss
  dst.st.traceMiss += src.st.traceMiss
  dst.st.classCorrect += src.st.classCorrect
  dst.st.classTotal += src.st.classTotal
  dst.st.radCorrect += src.st.radCorrect
  dst.st.radTotal += src.st.radTotal
  dst.st.revCorrect += src.st.revCorrect
  dst.st.revTotal += src.st.revTotal
  for c in 0..<TM_CLASSES:
    dst.st.labHist[c] += src.st.labHist[c]
    dst.st.choHist[c] += src.st.choHist[c]
    for p in 0..<TM_CLASSES:
      dst.st.confusion[c][p] += src.st.confusion[c][p]
      dst.st.radConfusion[c][p] += src.st.radConfusion[c][p]

proc replayRound(states: seq[WorldState], lastSeen: seq[int], enemyId, baseTick: int,
                 driver: GunDriver, metric: BulletMetric,
                 obsBefore: GunStats, obsCount: proc(): GunStats): Adapt =
  var tracker = initTracker(1, metric)
  var res: Adapt
  inc res.rounds
  for si in 0..<states.len:
    let state = states[si]
    var preds: array[len(PowerBins), GunPrediction]
    for i in 0..<len(PowerBins):
      preds[i] = driver.predictCb(state, bulletSpeed(PowerBins[i]))
    let ready = if driver.readyCb == nil: true else: driver.readyCb()
    if ready:
      tracker.spawnBullets(0, preds, state, enemyId)

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

    let localTick = state.tick - baseTick
    tracker.tickBullets(state, enemyPositions,
      proc(gunId: GunId, binIdx: int, e: FeedbackEvent) =
        inc res.nall
        if e.hit: inc res.hall
        if localTick < 100:
          inc res.n100
          if e.hit: inc res.h100
        if localTick < 300:
          inc res.n300
          if e.hit: inc res.h300
        let fireTick = e.fireTick - baseTick
        if fireTick < 100:
          inc res.m100
          if e.hit: inc res.f100
        driver.resultCb(e))
  let after = obsCount()
  res.st.obs = after.obs - obsBefore.obs
  res.st.labelMiss = after.labelMiss - obsBefore.labelMiss
  res.st.traceMiss = after.traceMiss - obsBefore.traceMiss
  res.st.classCorrect = after.classCorrect - obsBefore.classCorrect
  res.st.classTotal = after.classTotal - obsBefore.classTotal
  res.st.radCorrect = after.radCorrect - obsBefore.radCorrect
  res.st.radTotal = after.radTotal - obsBefore.radTotal
  res.st.revCorrect = after.revCorrect - obsBefore.revCorrect
  res.st.revTotal = after.revTotal - obsBefore.revTotal
  for c in 0..<TM_CLASSES:
    res.st.labHist[c] = after.labHist[c] - obsBefore.labHist[c]
    res.st.choHist[c] = after.choHist[c] - obsBefore.choHist[c]
    for p in 0..<TM_CLASSES:
      res.st.confusion[c][p] = after.confusion[c][p] - obsBefore.confusion[c][p]
      res.st.radConfusion[c][p] = after.radConfusion[c][p] - obsBefore.radConfusion[c][p]
  result = res

proc replayFixture(fx: Fixture, path: string, driver: GunDriver, metric: BulletMetric,
                   obsBefore: GunStats, obsCount: proc(): GunStats,
                   maxRounds = 0): Adapt =
  var spans =
    if fx.meta.source == "synthetic": @[(start: 0, count: fx.states.len)]
    else: loadRounds(path)
  if maxRounds > 0 and spans.len > maxRounds: spans.setLen(maxRounds)
  if spans.len == 0:
    return replayRound(fx.states, fx.lastSeen, fx.enemyId, 0, driver, metric,
                       obsBefore, obsCount)
  for sp in spans:
    var st: seq[WorldState]
    var ls: seq[int]
    for i in 0..<fx.states.len:
      let t = fx.states[i].tick
      if t >= sp.start and t < sp.start + sp.count:
        st.add fx.states[i]
        ls.add(if i < fx.lastSeen.len: fx.lastSeen[i] else: -1)
    if st.len == 0: continue
    addAdapt(result, replayRound(st, ls, fx.enemyId, sp.start, driver, metric,
                                 obsBefore, obsCount))

proc emptyStats(): GunStats = GunStats()

proc makeTmpatDriver(seed: int, shuffle: bool, forceBase = false,
                     mode = tmGF):
    tuple[driver: GunDriver, gun: ref TmPatternGun] =
  let g = new(TmPatternGun)
  g[] = initTmPatternGun()
  g[].shuffleLabels = shuffle
  g[].forceBase = forceBase
  g[].targetMode = mode
  if seed >= 0: randomize(seed)
  result.gun = g
  result.driver = GunDriver(
    name: (if forceBase: "TMPatternBase" elif shuffle: "TMPatternShuf" else: "TMPattern"),
    predictCb: proc(state: WorldState, bulletSpeed: float): GunPrediction =
      g[].predict(state, bulletSpeed),
    resultCb: proc(e: FeedbackEvent) = g[].onResult(e),
    readyCb: proc(): bool = g[].isWarmedUp())

proc tmpatStats(g: ref TmPatternGun): GunStats =
  result.obs = g[].totalObs
  result.labelMiss = g[].labelMisses
  result.traceMiss = g[].traceMisses
  result.labHist = g[].labelHist
  result.choHist = g[].chosenHist
  result.classCorrect = g[].classCorrect
  result.classTotal = g[].classTotal
  result.confusion = g[].confusion
  result.radConfusion = g[].radConfusion
  result.radCorrect = g[].radCorrect
  result.radTotal = g[].radTotal
  result.revCorrect = g[].revCorrect
  result.revTotal = g[].revTotal

proc fixtureSet(name: string): seq[string] =
  case name
  of "synthetic", "":
    for n in SyntheticFixtureNames: result.add n
  of "real":
    for n in ["drussgt_vs_crazy", "drussgt_vs_spinbot", "drussgt_vs_drussgt",
              "tr_drussgt_vs_crazy", "tr_drussgt_vs_spinbot",
              "tr_drussgt_vs_modularbot"]:
      result.add(fixturesDir / (n & ".jsonl"))
  of "range":
    for n in SyntheticFixtureNames: result.add n
    for n in ["drussgt_vs_crazy", "drussgt_vs_spinbot", "tr_drussgt_vs_crazy"]:
      result.add(fixturesDir / (n & ".jsonl"))
  else: discard

proc resolve(name: string): tuple[fx: Fixture, path: string] =
  let p = if fileExists(name): name else: fixturesDir / (name & ".jsonl")
  (loadFixture(p), p)

# ── stats ────────────────────────────────────────────────────────────────────

proc rateStr(h, n: int): string =
  if n == 0: "  n/a " else: &"{h.float / n.float * 100.0:5.1f}%"

proc binomPmf(k, n: int): float =
  if k < 0 or k > n: return 0.0
  var lg = 0.0
  for i in 1..k: lg += ln(float(n - k + i)) - ln(float(i))
  exp(lg - float(n) * ln(2.0))

proc signTestP(wins, n: int): float =
  ## exact two-sided binomial p under p=0.5
  if n == 0: return 1.0
  let lo = min(wins, n - wins)
  var s = 0.0
  for k in 0..lo: s += binomPmf(k, n)
  min(1.0, 2.0 * s)

proc printConfusion(variant: string, cm: array[TM_CLASSES, array[TM_CLASSES, int]]) =
  ## Task 1 deliverable: majority-class baseline vs online accuracy on the SAME
  ## (warm) samples, plus per-class precision/recall. `cm[true][pred]`.
  var total = 0
  var diag = 0
  var rowSum: array[TM_CLASSES, int]
  var colSum: array[TM_CLASSES, int]
  for c in 0..<TM_CLASSES:
    for p in 0..<TM_CLASSES:
      total += cm[c][p]
      if c == p: diag += cm[c][p]
      rowSum[c] += cm[c][p]
      colSum[p] += cm[c][p]
  var maj = 0
  for c in 1..<TM_CLASSES:
    if rowSum[c] > rowSum[maj]: maj = c
  let majShare = if total > 0: rowSum[maj].float / total.float else: 0.0
  let acc = if total > 0: diag.float / total.float else: 0.0
  var hs = ""
  for c in 0..<TM_CLASSES: hs.add &"{rowSum[c]},"
  echo &"# {variant}: warmTotal={total} warmLabelHist=[{hs}] majority=class{maj} " &
       &"majShare={majShare*100:.1f}% acc={diag}/{total}={acc*100:.1f}% " &
       &"margin={((acc-majShare)*100):+.1f}pp"
  for c in 0..<TM_CLASSES:
    let rec = if rowSum[c] > 0: cm[c][c].float / rowSum[c].float else: 0.0
    let prec = if colSum[c] > 0: cm[c][c].float / colSum[c].float else: 0.0
    echo &"#   class{c}: trueN={rowSum[c]} predN={colSum[c]} TP={cm[c][c]} " &
         &"recall={rec*100:.1f}% precision={prec*100:.1f}%"

proc main() =
  var nSeeds = 3
  var set = "real"
  var maxRounds = 0
  var metricName = "path"
  var variants: seq[VariantKind] = @[vLinear, vTsetlin, vTmpat, vTmpatShuf]
  for i in 1..paramCount():
    let a = paramStr(i)
    if a.startsWith("--seeds="): nSeeds = parseInt(a[8..^1])
    elif a.startsWith("--set="): set = a[6..^1]
    elif a.startsWith("--metric="): metricName = a[9..^1]
    elif a.startsWith("--maxrounds="): maxRounds = parseInt(a[12..^1])
    elif a.startsWith("--variants="):
      variants = @[]
      for tok in a[11..^1].split(','):
        case tok.strip()
        of "linear": variants.add vLinear
        of "linear_old": variants.add vLinearOld
        of "tsetlin": variants.add vTsetlin
        of "tmpat": variants.add vTmpat
        of "tmpat_shuf": variants.add vTmpatShuf
        of "tmbase": variants.add vTmpatBase
        of "tmrad": variants.add vTmpatRad
        of "tmrad_shuf": variants.add vTmpatRadShuf
        of "tmrev": variants.add vTmpatRev
        of "tmrev_shuf": variants.add vTmpatRevShuf
        else: discard
  let metric = if metricName == "point": bmPoint else: bmPath
  let names = fixtureSet(set)
  echo "# set=", set, " seeds=", nSeeds, " metric=", metricName, " variants=", variants.mapIt(variantName(it)).join(",")
  echo "variant,fixture,seed,h100,n100,h300,n300,hall,nall,f100,m100,rounds,obs,labelMiss,traceMiss"

  var rows: seq[Row]
  var revLabTab = initTable[string, array[2, int]]()
  var radLabTab = initTable[string, array[TM_CLASSES, int]]()
  for name in names:
    let (fx, path) = resolve(name)
    let fxName = path.extractFilename.replace(".jsonl", "")
    for v in variants:
      let nIter = if v in [vLinear, vLinearOld, vTmpatBase]: 1 else: nSeeds
      for seed in 1..nIter:
        var drv: GunDriver
        var gun: ref TmPatternGun
        var obsCount: proc(): GunStats = emptyStats
        case v
        of vLinear:
          drv = makeDriver("Linear", LinearGun())
        of vLinearOld:
          drv = makeDriver("LinearOldClamp", LinearInsetGun())
        of vTsetlin:
          let pair = makeTsetlinDriver(seed = seed)
          drv = pair.driver
        of vTmpat, vTmpatShuf, vTmpatBase, vTmpatRad, vTmpatRadShuf,
           vTmpatRev, vTmpatRevShuf:
          let mode =
            case v
            of vTmpatRad, vTmpatRadShuf: tmRadial
            of vTmpatRev, vTmpatRevShuf: tmReversal
            else: tmGF
          let shuf = v in [vTmpatShuf, vTmpatRadShuf, vTmpatRevShuf]
          let pair = makeTmpatDriver(seed = seed, shuffle = shuf,
                                     forceBase = (v == vTmpatBase), mode = mode)
          drv = pair.driver
          gun = pair.gun
          obsCount = proc(): GunStats = tmpatStats(gun)
        let r = replayFixture(fx, path, drv, metric, emptyStats(), obsCount, maxRounds)
        rows.add Row(variant: variantName(v), fixture: fxName, seed: seed, r: r)
        if gun != nil:
          if variantName(v) notin revLabTab: revLabTab[variantName(v)] = [0, 0]
          if variantName(v) notin radLabTab:
            radLabTab[variantName(v)] = default(array[TM_CLASSES, int])
          for c in 0..<2: revLabTab[variantName(v)][c] += gun[].revLabelHist[c]
          for c in 0..<TM_CLASSES: radLabTab[variantName(v)][c] += gun[].radLabelHist[c]
        echo &"{variantName(v)},{fxName},{seed},{r.h100},{r.n100},{r.h300},{r.n300}," &
             &"{r.hall},{r.nall},{r.f100},{r.m100},{r.rounds},{r.st.obs},{r.st.labelMiss},{r.st.traceMiss}"

  # ── per-variant pooled summary ──
  echo "\n# ── pooled summary ──"
  echo "variant,runs,h100,n100,early%,h300,n300,early300%,hall,nall,overall%,obs,labelMiss,traceMiss"
  var pooled = initTable[string, Adapt]()
  for v in variants: pooled[variantName(v)] = Adapt()
  for row in rows: addAdapt(pooled[row.variant], row.r)
  for v in variants:
    let a = pooled[variantName(v)]
    echo &"{variantName(v)},{a.rounds},{a.h100},{a.n100},{rateStr(a.h100, a.n100)}," &
         &"{a.h300},{a.n300},{rateStr(a.h300, a.n300)},{a.hall},{a.nall}," &
         &"{rateStr(a.hall, a.nall)},{a.st.obs},{a.st.labelMiss},{a.st.traceMiss}"

  # ── label vs chosen class histogram (all TM variants) ──
  for v in [vTmpat, vTmpatShuf, vTmpatRad, vTmpatRadShuf, vTmpatRev, vTmpatRevShuf]:
    if v in variants:
      let a = pooled[variantName(v)]
      var ls, cs: string
      for c in 0..<TM_CLASSES:
        ls.add &"{a.st.labHist[c]},"
        cs.add &"{a.st.choHist[c]},"
      echo &"\n# class histogram {variantName(v)}: labels=[{ls}] chosen=[{cs}] " &
           &"onlineAcc={a.st.classCorrect}/{a.st.classTotal} " &
           &"radAcc={a.st.radCorrect}/{a.st.radTotal} " &
           &"revAcc={a.st.revCorrect}/{a.st.revTotal}"
      if variantName(v) in revLabTab:
        var rl, rdl: string
        for c in 0..<2: rl.add &"{revLabTab[variantName(v)][c]},"
        for c in 0..<TM_CLASSES: rdl.add &"{radLabTab[variantName(v)][c]},"
        echo &"#   label hist {variantName(v)}: radial=[{rdl}] rev=[{rl}]"

  # ── TASK 1: majority-class baseline vs accuracy + per-class P/R ──
  echo "\n# ── GF head majority-class test (warm samples only) ──"
  for v in [vTmpat, vTmpatShuf]:
    if v in variants: printConfusion(variantName(v), pooled[variantName(v)].st.confusion)
  echo "\n# ── radial head confusion (context) ──"
  for v in [vTmpatRad, vTmpatRadShuf]:
    if v in variants: printConfusion(variantName(v), pooled[variantName(v)].st.radConfusion)

  # ── per-run distributions (a "run" = one fixture × one seed) ──
  # Linear is deterministic: replicate its one row per fixture across seeds so a
  # paired comparison against a stochastic variant has a partner per run.
  type Key = tuple[fixture: string, seed: int]
  var byVariant = initTable[string, Table[Key, float]]()   # early rate
  var byVariantAll = initTable[string, Table[Key, float]]() # overall rate
  for v in variants:
    byVariant[variantName(v)] = initTable[Key, float]()
    byVariantAll[variantName(v)] = initTable[Key, float]()
  for row in rows:
    let early = if row.r.n100 > 0: row.r.h100.float / row.r.n100.float else: 0.0
    let overall = if row.r.nall > 0: row.r.hall.float / row.r.nall.float else: 0.0
    byVariant[row.variant][(row.fixture, row.seed)] = early
    byVariantAll[row.variant][(row.fixture, row.seed)] = overall
  if vLinear in variants:
    for row in rows:
      if row.variant == "Linear":
        for s in 2..nSeeds:
          byVariant["Linear"][(row.fixture, s)] = byVariant["Linear"][(row.fixture, 1)]
          byVariantAll["Linear"][(row.fixture, s)] = byVariantAll["Linear"][(row.fixture, 1)]
  if vLinearOld in variants:
    for row in rows:
      if row.variant == "LinearOldClamp":
        for s in 2..nSeeds:
          byVariant["LinearOldClamp"][(row.fixture, s)] = byVariant["LinearOldClamp"][(row.fixture, 1)]
          byVariantAll["LinearOldClamp"][(row.fixture, s)] = byVariantAll["LinearOldClamp"][(row.fixture, 1)]
  if vTmpatBase in variants:
    for row in rows:
      if row.variant == "TMPatternBase":
        for s in 2..nSeeds:
          byVariant["TMPatternBase"][(row.fixture, s)] = byVariant["TMPatternBase"][(row.fixture, 1)]
          byVariantAll["TMPatternBase"][(row.fixture, s)] = byVariantAll["TMPatternBase"][(row.fixture, 1)]

  echo "\n# ── per-run early-rate distribution (mean / min / max, n runs) ──"
  echo "variant,earlyMean%,earlyMin%,earlyMax%,overallMean%,overallMin%,overallMax%,n"
  for v in variants:
    var es: seq[float]
    var os: seq[float]
    for r in byVariant[variantName(v)].values: es.add r
    for r in byVariantAll[variantName(v)].values: os.add r
    if es.len == 0: continue
    es.sort(); os.sort()
    echo &"{variantName(v)},{es.sum/float(es.len)*100:.2f},{es[0]*100:.2f},{es[^1]*100:.2f}," &
         &"{os.sum/float(os.len)*100:.2f},{os[0]*100:.2f},{os[^1]*100:.2f},{es.len}"

  echo "\n# ── pairwise paired sign tests (rows = fixture×seed) ──"
  echo "A,B,metric,nA>B,nB>A,ties,p"
  for i in 0..<variants.len:
    for j in 0..<variants.len:
      if i == j: continue
      let aName = variantName(variants[i])
      let bName = variantName(variants[j])
      for (label, tab) in [("early", byVariant), ("overall", byVariantAll)]:
        var winsA, winsB, ties, n = 0
        for k, va in tab[aName].pairs:
          if k notin tab[bName]: continue
          let vb = tab[bName][k]
          inc n
          if va > vb: inc winsA
          elif vb > va: inc winsB
          else: inc ties
        if n == 0: continue
        echo &"{aName},{bName},{label},{n},{winsA},{winsB},{ties},{signTestP(winsA, n - ties):.4f}"

when isMainModule:
  main()