## tm_diag/diagnostics.nim — THE DIAGNOSTICS KIT (Tasks 2 + 5).
##
## Everything runs OFFLINE: a trained `TmMachine` plus a labelled sample set.
## No battles, no harness, no Java. The seven groups are:
##   1. pre-flight DATA checks        (dataChecks, shuffledLabelControl)
##   2. clause introspection          (clauseInfo, clauseSummary, topClauses)
##   3. per-feature contribution      (featureContributions, deadInputs, ...)
##   4. accuracy diagnostics          (accuracyDiagnostics)
##   5. learning curve                (learningCurve)
##   6. ablation hooks                (ablateDropBlock, ablateScrambleFeature, ...)
##   7. AUTOMATA level (Task 5)       (settledness, clauseDiversity, churnTrace,
##                                     voteDisagreement, stateHistogram,
##                                     perInputConfidence, healthLine,
##                                     automataDiagnostics)
##
## A `DiagSample` carries the LITERAL vector (positive literals [0..nBits),
## negations [nBits..2*nBits)) exactly as the shipped TM stores it, plus the
## true label. This keeps the kit usable both on the reference core and on the
## literal vectors captured from the real gun.

import std/[math, strformat, algorithm, strutils]
import feature_spec
import tm_core

export feature_spec, tm_core

type
  DiagSample* = object
    lits*: seq[uint8]
    label*: int
    order*: int          ## optional capture/time index (for post-hoc curves)

# ── sample construction ──────────────────────────────────────────────────────

proc makeSample*(nBits: int, rawBits: openArray[int], label: int,
                 order = 0): DiagSample =
  ## Expand `nBits` raw bits into the pos/neg literal vector.
  doAssert rawBits.len >= nBits,
    "rawBits.len (" & $rawBits.len & ") < nBits (" & $nBits & ")"
  result.lits = newSeq[uint8](2 * nBits)
  for i in 0..<nBits:
    let v = uint8(if rawBits[i] != 0: 1 else: 0)
    result.lits[i] = v
    result.lits[i + nBits] = 1'u8 - v
  result.label = label
  result.order = order

proc makeSample*(nBits: int, rawBits: openArray[uint8], label: int,
                 order = 0): DiagSample =
  result.lits = newSeq[uint8](2 * nBits)
  for i in 0..<nBits:
    let v = uint8(if rawBits[i] != 0'u8: 1 else: 0)
    result.lits[i] = v
    result.lits[i + nBits] = 1'u8 - v
  result.label = label
  result.order = order

proc machineFromTeams*(nBits, nClasses, nClauses, nStates: int, sValue: float,
                       teams: openArray[seq[int16]]): TmMachine =
  ## Build an introspectable machine around an externally-trained clause set
  ## (e.g. the shipped tm_pattern gun's `teams`).
  result = newMachine(nBits, nClasses, nClauses, nStates, sValue)
  doAssert teams.len == nClasses
  for c in 0..<nClasses:
    result.teams[c] = teams[c]

# ── shared training / scoring helpers (used by every group below) ────────────

proc shuffledIndices(n: int, rng: var TmRng): seq[int] =
  result = newSeq[int](n)
  for i in 0..<n: result[i] = i
  for i in countdown(n - 1, 1):
    let j = int(rng.nextU64() mod uint64(i + 1))
    swap(result[i], result[j])

proc trainModel*(tmpl: TmMachine, samples: openArray[DiagSample],
                 epochs = 1, seed = 777'u64): TmMachine =
  ## Fresh machine, `epochs` shuffled passes over `samples`.
  result = tmpl
  result.resetMachine(seed)
  for _ in 0..<epochs:
    let idx = shuffledIndices(samples.len, result.rng)
    for k in idx:
      result.trainSample(samples[k].lits, samples[k].label)

proc evalAcc*(m: TmMachine, samples: openArray[DiagSample]): float =
  var c = 0
  for s in samples:
    if m.predictClass(s.lits) == s.label: inc c
  if samples.len == 0: 0.0 else: c.float / samples.len.float

# ─────────────────────────────────────────────────────────────────────────────
# GROUP 1 — pre-flight DATA checks
# ─────────────────────────────────────────────────────────────────────────────

type
  DataCheck* = object
    n*: int
    nClasses*: int
    classCounts*: seq[int]
    classShares*: seq[float]
    majorityClass*: int
    majorityShare*: float
    majorityAccuracy*: float
    maxShare*: float
    overThreshold*: bool
    threshold*: float
    flags*: seq[string]

proc dataChecks*(labels: openArray[int], nClasses: int,
                 threshold = 0.30): DataCheck =
  ## Per-class label counts/shares, the majority-class share (which is also the
  ## majority-class accuracy baseline), and a loud flag when any class exceeds
  ## `threshold` of the samples — a too-large class compromises the accuracy
  ## test.
  result.threshold = threshold
  result.nClasses = nClasses
  result.n = labels.len
  result.classCounts = newSeq[int](nClasses)
  result.classShares = newSeq[float](nClasses)
  for l in labels:
    if l >= 0 and l < nClasses: inc result.classCounts[l]
  var maj = 0
  for c in 0..<nClasses:
    result.classShares[c] =
      if result.n > 0: result.classCounts[c].float / result.n.float else: 0.0
    if result.classCounts[c] > result.classCounts[maj]: maj = c
  result.majorityClass = maj
  result.majorityShare = result.classShares[maj]
  result.majorityAccuracy = result.majorityShare
  result.maxShare = result.classShares[maj]
  for c in 0..<nClasses:
    if result.classShares[c] > threshold:
      result.overThreshold = true
      result.flags.add &"MAJORITY HEAVY: class {c} = " &
        &"{result.classCounts[c]}/{result.n} = " &
        &"{result.classShares[c]*100.0:.1f}% > {threshold*100.0:.0f}%"

proc shuffledLabelControl*(tmpl: TmMachine, samples: openArray[DiagSample],
                           epochs = 1, seed = 4242'u64):
    tuple[acc, majority: float, n, nClasses: int, trainedObs: int] =
  ## PIPELINE LEAK TEST: retrain a fresh copy on SHUFFLED labels. A correct
  ## pipeline cannot beat the majority baseline on shuffled labels; a large
  ## margin means information is leaking (features from the future, duplicated
  ## samples, ...).
  var rng = seedRng(seed xor 0xabcdef12'u64)
  var shuf = newSeq[DiagSample](samples.len)
  for i, s in samples: shuf[i] = s
  var labels = newSeq[int](samples.len)
  for i in 0..<samples.len: labels[i] = samples[i].label
  for i in countdown(samples.len - 1, 1):
    let j = int(rng.nextU64() mod uint64(i + 1))
    swap(labels[i], labels[j])
  for i in 0..<samples.len: shuf[i].label = labels[i]
  let m = trainModel(tmpl, shuf, epochs, seed)
  result.acc = evalAcc(m, shuf)
  result.majority = dataChecks(labels, tmpl.nClasses).majorityShare
  result.n = samples.len
  result.nClasses = tmpl.nClasses
  result.trainedObs = epochs * samples.len

# ─────────────────────────────────────────────────────────────────────────────
# GROUP 2 — clause introspection
# ─────────────────────────────────────────────────────────────────────────────

type
  ClauseInfo* = object
    cls*: int
    index*: int
    polarity*: int
    length*: int
    lits*: seq[int]
    votes*: int        ## samples in which the clause fires (casts a vote)
    text*: string      ## describeClause rendering

  ClauseSummary* = object
    totalClauses*: int
    emptyClauses*: int
    nonEmpty*: int
    firedAtLeastOnce*: int
    neverFired*: int
    posFired*: int
    negFired*: int
    lengthHist*: seq[int]   ## index = clause length, value = count (non-empty)
    meanLength*: float
    maxLength*: int

proc clauseInfo*(m: TmMachine, samples: openArray[DiagSample],
                 spec: FeatureSpec): seq[ClauseInfo] =
  for c in 0..<m.nClasses:
    for cl in 0..<m.nClauses:
      let ls = m.clauseLits(c, cl)
      var votes = 0
      for s in samples:
        if m.clauseFires(c, cl, s.lits): inc votes
      result.add ClauseInfo(cls: c, index: cl,
        polarity: (if cl < m.half: 1 else: -1), length: ls.len, lits: ls,
        votes: votes, text: spec.describeClause(ls, c))

proc clauseSummary*(infos: openArray[ClauseInfo]): ClauseSummary =
  var maxLen = 0
  var lenSum = 0
  for inf in infos:
    inc result.totalClauses
    if inf.length == 0:
      inc result.emptyClauses
    else:
      inc result.nonEmpty
      lenSum += inf.length
      if inf.length > maxLen: maxLen = inf.length
      if inf.votes > 0:
        inc result.firedAtLeastOnce
        if inf.polarity > 0: inc result.posFired
        else: inc result.negFired
      else:
        inc result.neverFired
  result.maxLength = maxLen
  result.lengthHist = newSeq[int](maxLen + 1)
  for inf in infos:
    if inf.length > 0: inc result.lengthHist[inf.length]
  result.meanLength =
    if result.nonEmpty > 0: lenSum.float / result.nonEmpty.float else: 0.0

type
  ClassClauseBalance* = object
    cls*: int
    empty*: int
    nonEmpty*: int
    posFired*: int
    negFired*: int
    neverFired*: int

proc clauseBalanceByClass*(infos: openArray[ClauseInfo],
                           nClasses: int): seq[ClassClauseBalance] =
  ## Positive/negative firing balance PER CLASS, plus empty and never-fired
  ## counts. A healthy class should use both polarities; a class with 0 firing
  ## positive clauses has no readable rule.
  result = newSeq[ClassClauseBalance](nClasses)
  for c in 0..<nClasses: result[c].cls = c
  for inf in infos:
    if inf.cls < 0 or inf.cls >= nClasses: continue
    if inf.length == 0:
      inc result[inf.cls].empty
      continue
    inc result[inf.cls].nonEmpty
    if inf.votes == 0:
      inc result[inf.cls].neverFired
    elif inf.polarity > 0:
      inc result[inf.cls].posFired
    else:
      inc result[inf.cls].negFired

proc topClauses*(infos: openArray[ClauseInfo], k = 10,
                 minVotes = 1): seq[ClauseInfo] =
  ## The clauses that actually fire, ranked by vote count (ties: longer first,
  ## then class/index).
  for inf in infos:
    if inf.votes >= minVotes and inf.length > 0: result.add inf
  result.sort(proc(a, b: ClauseInfo): int =
    result = cmp(b.votes, a.votes)
    if result == 0: result = cmp(b.length, a.length)
    if result == 0: result = cmp(a.cls, b.cls)
    if result == 0: result = cmp(a.index, b.index))
  if result.len > k: result.setLen(k)

proc topClausesByPolarity*(infos: openArray[ClauseInfo], polarity: int,
                           k = 10): seq[ClauseInfo] =
  ## Firing clauses of one polarity only (positive = the rule-encoding clauses,
  ## negative = the discriminators), ranked by votes.
  var filtered: seq[ClauseInfo]
  for inf in infos:
    if inf.polarity == polarity: filtered.add inf
  result = topClauses(filtered, k)

proc necessaryLiterals*(m: TmMachine, samples: openArray[DiagSample],
                        cls: int): seq[int] =
  ## The literals common to EVERY positive-polarity clause of `cls` that fires
  ## at least once. A converged multiclass TM pads its clauses with redundant
  ## literals; the intersection strips the padding and recovers the class rule
  ## (e.g. the planted `A AND B`).
  var first = true
  for cl in 0..<m.half:              # positive polarity only
    let ls = m.clauseLits(cls, cl)
    if ls.len == 0: continue
    var fires = false
    for s in samples:
      if m.clauseFires(cls, cl, s.lits):
        fires = true
        break
    if not fires: continue
    if first:
      result = ls
      first = false
    else:
      var keep: seq[int]
      for l in result:
        if l in ls: keep.add l
      result = keep
  result.sort()

# ─────────────────────────────────────────────────────────────────────────────
# GROUP 3 — per-feature contribution + DEAD-INPUT LIST
# ─────────────────────────────────────────────────────────────────────────────

type
  FeatureContribution* = object
    bit*: int
    name*: string
    appearances*: int     ## times the bit appears in a clause that casts a vote
    weighted*: float      ## sum of clause vote-shares (1/nCasting per fire)

proc featureContributions*(m: TmMachine, samples: openArray[DiagSample],
                           spec: FeatureSpec): seq[FeatureContribution] =
  ## For every raw bit: how often its positive or negated literal appears in a
  ## clause that ACTUALLY CASTS A VOTE (non-empty + fires on a sample), weighted
  ## by that clause's share of the sample's casting votes.
  result = newSeq[FeatureContribution](m.nBits)
  for b in 0..<m.nBits:
    result[b].bit = b
    result[b].name = spec.describe(b)

  var allLits = newSeq[seq[int]](m.nClasses * m.nClauses)
  for c in 0..<m.nClasses:
    for cl in 0..<m.nClauses:
      allLits[c * m.nClauses + cl] = m.clauseLits(c, cl)

  for s in samples:
    var casting: seq[int]
    for k in 0..<allLits.len:
      let ls = allLits[k]
      if ls.len == 0: continue
      var fires = true
      for lit in ls:
        if s.lits[lit] == 0'u8:
          fires = false
          break
      if fires: casting.add k
    if casting.len == 0: continue
    let w = 1.0 / float(casting.len)
    for k in casting:
      for lit in allLits[k]:
        let b = if lit < m.nBits: lit else: lit - m.nBits
        inc result[b].appearances
        result[b].weighted += w

proc deadInputs*(contribs: openArray[FeatureContribution],
                 relThreshold = 0.05): seq[int] =
  ## THE DEAD-INPUT LIST. A bit is dead if it never appears in a voting clause
  ## (`appearances == 0`) OR its weighted vote-share contribution is below
  ## `relThreshold` x the largest contribution.
  ##
  ## The relative form is the practical one: a CONVERGED multiclass TM keeps
  ## redundant literals inside otherwise-correct clauses (each fires rarely, so
  ## its vote share is tiny). Strict `appearances == 0` therefore misses bits
  ## that are effectively dead; pass `relThreshold = 0.0` for the strict list.
  var mx = 0.0
  for c in contribs:
    if c.weighted > mx: mx = c.weighted
  let cut = relThreshold * mx
  for c in contribs:
    if c.appearances == 0 or c.weighted < cut:
      result.add c.bit

proc neverUsedInputs*(contribs: openArray[FeatureContribution]): seq[int] =
  ## Strict form: raw bits that never appear in a clause that casts a vote.
  for c in contribs:
    if c.appearances == 0: result.add c.bit

proc constantInputs*(samples: openArray[DiagSample], nBits: int): seq[int] =
  ## Raw bits with zero variance across the sample set (information-free even if
  ## a clause happens to include them).
  if samples.len == 0: return
  for b in 0..<nBits:
    let v0 = samples[0].lits[b]
    var constant = true
    for i in 1..<samples.len:
      if samples[i].lits[b] != v0:
        constant = false
        break
    if constant: result.add b

proc rankedInputs*(contribs: seq[FeatureContribution]): seq[FeatureContribution] =
  ## Most valuable first.
  result = contribs
  result.sort(proc(a, b: FeatureContribution): int =
    result = cmp(b.weighted, a.weighted)
    if result == 0: result = cmp(b.appearances, a.appearances)
    if result == 0: result = cmp(a.bit, b.bit))

# ─────────────────────────────────────────────────────────────────────────────
# GROUP 4 — accuracy diagnostics
# ─────────────────────────────────────────────────────────────────────────────

type
  AccDiag* = object
    n*: int
    correct*: int
    acc*: float
    confusion*: seq[seq[int]]     ## [true][pred]
    majorityClass*: int
    majorityShare*: float
    majorityBaseline*: float
    margin*: float                ## acc - majority baseline (pp as fraction)
    perClassRecall*: seq[float]
    perClassPrecision*: seq[float]
    predCounts*: seq[int]
    predMajorityShare*: float     ## how often the model predicts the majority class

proc accuracyDiagnostics*(m: TmMachine,
                          samples: openArray[DiagSample]): AccDiag =
  result.confusion = newSeq[seq[int]](m.nClasses)
  for c in 0..<m.nClasses: result.confusion[c] = newSeq[int](m.nClasses)
  result.predCounts = newSeq[int](m.nClasses)
  result.perClassRecall = newSeq[float](m.nClasses)
  result.perClassPrecision = newSeq[float](m.nClasses)
  result.n = samples.len
  for s in samples:
    let p = m.predictClass(s.lits)
    if s.label >= 0 and s.label < m.nClasses and p >= 0 and p < m.nClasses:
      inc result.confusion[s.label][p]
      inc result.predCounts[p]
      if p == s.label: inc result.correct
  result.acc =
    if result.n > 0: result.correct.float / result.n.float else: 0.0
  var rowSum = newSeq[int](m.nClasses)
  for c in 0..<m.nClasses:
    for p in 0..<m.nClasses: rowSum[c] += result.confusion[c][p]
  var maj = 0
  for c in 0..<m.nClasses:
    if rowSum[c] > rowSum[maj]: maj = c
  result.majorityClass = maj
  result.majorityShare =
    if result.n > 0: rowSum[maj].float / result.n.float else: 0.0
  result.majorityBaseline = result.majorityShare
  result.margin = result.acc - result.majorityBaseline
  for c in 0..<m.nClasses:
    result.perClassRecall[c] =
      if rowSum[c] > 0: result.confusion[c][c].float / rowSum[c].float else: 0.0
    result.perClassPrecision[c] =
      if result.predCounts[c] > 0:
        result.confusion[c][c].float / result.predCounts[c].float
      else: 0.0
  result.predMajorityShare =
    if result.n > 0: result.predCounts[maj].float / result.n.float else: 0.0

# ─────────────────────────────────────────────────────────────────────────────
# GROUP 5 — learning curve
# ─────────────────────────────────────────────────────────────────────────────

type
  LearningCurve* = object
    points*: seq[int]     ## samples seen
    accs*: seq[float]     ## accuracy on the (held-out) eval set
    trend*: string

proc classifyTrend*(accs: seq[float], eps = 0.02): string =
  if accs.len < 3: return "insufficient"
  let first = accs[0]
  let last = accs[^1]
  var mx = first
  var mxIdx = 0
  for i, a in accs:
    if a > mx:
      mx = a
      mxIdx = i
  if mx - last > eps and mxIdx > 0 and mxIdx < accs.len - 1:
    return "rising-then-falling (noise fitting)"
  if last - first > eps: return "rising"
  if first - last > eps: return "falling"
  "flat"

proc learningCurve*(tmpl: TmMachine, train, eval: openArray[DiagSample],
                    nPoints = 10, seed = 777'u64): LearningCurve =
  ## Train a fresh machine on growing prefixes of `train` and score each prefix
  ## on `eval`, so fast vs slow adaptation is directly visible.
  var m = tmpl
  m.resetMachine(seed)
  let np = max(2, nPoints)
  let step = max(1, train.len div np)
  var seen = 0
  var i = 0
  while i < np:
    let upto = min(train.len, (i + 1) * step)
    var k = seen
    while k < upto:
      m.trainSample(train[k].lits, train[k].label)
      inc k
    seen = upto
    var correct = 0
    for s in eval:
      if m.predictClass(s.lits) == s.label: inc correct
    result.points.add upto
    result.accs.add (if eval.len > 0: correct.float / eval.len.float else: 0.0)
    inc i
    if upto >= train.len:
      while i < np:
        result.points.add train.len
        result.accs.add result.accs[^1]
        inc i
      break
  result.trend = classifyTrend(result.accs)

# ─────────────────────────────────────────────────────────────────────────────
# GROUP 6 — ablation hooks
# ─────────────────────────────────────────────────────────────────────────────

type
  AblationResult* = object
    name*: string
    kind*: string        ## "baseline" | "drop-block" | "scramble-feature"
    baselineAcc*: float
    ablatedAcc*: float
    delta*: float        ## ablated - baseline

proc ablateAdd*(a, b: AblationResult): AblationResult =
  ## Accumulate ablations across folds/fixtures.
  result.name = a.name
  result.kind = a.kind
  let n = 2.0
  result.baselineAcc = (a.baselineAcc + b.baselineAcc) / n
  result.ablatedAcc = (a.ablatedAcc + b.ablatedAcc) / n
  result.delta = (a.delta + b.delta) / n

proc dropBits(s: DiagSample, nBits, first, count: int): DiagSample =
  result = s
  result.lits = s.lits
  for b in first..<first + count:
    result.lits[b] = 0'u8
    result.lits[b + nBits] = 1'u8

proc scrambleBit(samples: var seq[DiagSample], nBits, bit: int,
                 rng: var TmRng) =
  let n = samples.len
  if n < 2: return
  var vals = newSeq[uint8](n)
  for i in 0..<n: vals[i] = samples[i].lits[bit]
  for i in countdown(n - 1, 1):
    let j = int(rng.nextU64() mod uint64(i + 1))
    swap(vals[i], vals[j])
  for i in 0..<n:
    samples[i].lits[bit] = vals[i]
    samples[i].lits[bit + nBits] = 1'u8 - vals[i]

proc ablateDropBlock*(tmpl: TmMachine, train, eval: openArray[DiagSample],
                      spec: FeatureSpec, blockIdx: int, epochs = 1,
                      seed = 777'u64, baselineAcc = -1.0): AblationResult =
  ## Retrain from scratch with one whole block held at 0 (feature removed) and
  ## report the accuracy delta. ~zero => the block does not earn its bits.
  ## Pass `baselineAcc` (from ablateBaseline) to skip recomputing it per block.
  doAssert blockIdx >= 0 and blockIdx < spec.blocks.len
  let b = spec.blocks[blockIdx]
  var tr = newSeq[DiagSample](train.len)
  for i, s in train: tr[i] = dropBits(s, spec.nBits, b.first, b.count)
  var ev = newSeq[DiagSample](eval.len)
  for i, s in eval: ev[i] = dropBits(s, spec.nBits, b.first, b.count)
  let base =
    if baselineAcc >= 0.0: baselineAcc
    else: evalAcc(trainModel(tmpl, train, epochs, seed), eval)
  let abl = evalAcc(trainModel(tmpl, tr, epochs, seed), ev)
  AblationResult(name: b.name, kind: "drop-block",
                 baselineAcc: base, ablatedAcc: abl, delta: abl - base)

proc ablateScrambleFeature*(tmpl: TmMachine, train, eval: openArray[DiagSample],
                            spec: FeatureSpec, bit: int, epochs = 1,
                            seed = 777'u64, baselineAcc = -1.0): AblationResult =
  ## Retrain from scratch with ONE raw bit permuted across samples in both train
  ## and eval. A negative delta => the bit is genuinely informative.
  var rng = seedRng(seed xor 0x5bd1e995'u64)
  var tr = newSeq[DiagSample](train.len)
  for i, s in train: tr[i] = s
  var ev = newSeq[DiagSample](eval.len)
  for i, s in eval: ev[i] = s
  scrambleBit(tr, spec.nBits, bit, rng)
  scrambleBit(ev, spec.nBits, bit, rng)
  let base =
    if baselineAcc >= 0.0: baselineAcc
    else: evalAcc(trainModel(tmpl, train, epochs, seed), eval)
  let abl = evalAcc(trainModel(tmpl, tr, epochs, seed), ev)
  AblationResult(name: spec.describe(bit), kind: "scramble-feature",
                 baselineAcc: base, ablatedAcc: abl, delta: abl - base)

proc ablateDropAllBlocks*(tmpl: TmMachine, train, eval: openArray[DiagSample],
                          spec: FeatureSpec, epochs = 1, seed = 777'u64):
    seq[AblationResult] =
  for i in 0..<spec.blocks.len:
    result.add ablateDropBlock(tmpl, train, eval, spec, i, epochs, seed)

proc ablateBaseline*(tmpl: TmMachine, train, eval: openArray[DiagSample],
                     epochs = 1, seed = 777'u64): AblationResult =
  AblationResult(name: "baseline", kind: "baseline",
                 baselineAcc: evalAcc(trainModel(tmpl, train, epochs, seed), eval),
                 ablatedAcc: evalAcc(trainModel(tmpl, train, epochs, seed), eval),
                 delta: 0.0)

# ─────────────────────────────────────────────────────────────────────────────
# GROUP 7 — automata-level metrics
#   settledness / clause diversity / churn / vote disagreement
#   plus the state histogram, the per-input confidence table and a health line.
#
# STATE CONVENTIONS — derived from the ACTUAL code (tm_core.nim and
# guns/tm_pattern.nim), not from the textbook formulation:
#   * every clause owns one automaton per literal (2 * nBits of them);
#   * states are int16 in [-nStates, nStates];
#   * INCLUDE iff state > 0, EXCLUDE iff state <= 0 (tmEval / clauseLits);
#   * both cores initialise every state to 0 ("the Exclude boundary"), so the
#     flip boundary sits BETWEEN state 0 (Exclude) and state 1 (Include) — the
#     middle of the range.
#   * commitment(st) = |st| / nStates in [0, 1]: 0 exactly on the boundary,
#     1 at either extreme. It is symmetric for the two sides.
# ─────────────────────────────────────────────────────────────────────────────

proc stateCommitment*(nStates, st: int): float {.inline.} =
  ## Normalised distance of an automaton state from its flip boundary, in
  ## [0, 1]. 0 == on the boundary (state 0), 1 == at either extreme (+/-nStates).
  min(abs(st).float / max(1, nStates).float, 1.0)

proc stateIncluded*(st: int): bool {.inline.} =
  ## The kit's include predicate, matching `tmEval` / `clauseLits` exactly.
  st > 0

# ── 1. SETTLEDNESS ───────────────────────────────────────────────────────────

type
  ClauseSettledness* = object
    cls*: int
    index*: int
    polarity*: int          ## +1 positive clauses, -1 negative clauses
    nAutomata*: int
    meanCommitment*: float
    settledFraction*: float

  SettlednessResult* = object
    threshold*: float
    nAutomata*: int
    overallMean*: float
    overallSettledFraction*: float
    clauses*: seq[ClauseSettledness]

proc settledness*(m: TmMachine, threshold = 0.5): SettlednessResult =
  ## Per clause and overall: the mean automaton `commitment` and the fraction of
  ## automata whose commitment is >= `threshold`. A clause full of automata
  ## parked near the boundary is wavering noise; one with automata pushed to an
  ## extreme has committed to a rule.
  result.threshold = threshold
  var sum = 0.0
  var settled = 0
  for c in 0..<m.nClasses:
    for cl in 0..<m.nClauses:
      let base = cl * m.nLiterals
      var csum = 0.0
      var csettled = 0
      for lit in 0..<m.nLiterals:
        let cm = stateCommitment(m.nStates, int(m.teams[c][base + lit]))
        csum += cm
        if cm >= threshold: inc csettled
      result.clauses.add ClauseSettledness(
        cls: c, index: cl,
        polarity: (if cl < m.half: 1 else: -1),
        nAutomata: m.nLiterals,
        meanCommitment: csum / m.nLiterals.float,
        settledFraction: csettled.float / m.nLiterals.float)
      sum += csum
      settled += csettled
      inc result.nAutomata, m.nLiterals
  result.overallMean =
    if result.nAutomata > 0: sum / result.nAutomata.float else: 0.0
  result.overallSettledFraction =
    if result.nAutomata > 0: settled.float / result.nAutomata.float else: 0.0

proc settlednessTrend*(early, late: SettlednessResult): string =
  ## Classify whether settledness ROSE between two snapshots (the healthy sign).
  let d = late.overallMean - early.overallMean
  if d > 0.05: "rising"
  elif d < -0.05: "falling"
  else: "flat"

# ── state histogram (the raw picture behind settledness) ─────────────────────

type
  StateHistogram* = object
    nStates*: int
    nBins*: int
    binLo*: seq[int]        ## inclusive low state observed in the bin
    binHi*: seq[int]        ## inclusive high state observed in the bin
    counts*: seq[int]
    total*: int

proc stateHistogram*(m: TmMachine, nBins = 9): StateHistogram =
  ## Distribution of every automaton state across [-nStates, nStates].
  let nb = max(1, nBins)
  let span = 2 * m.nStates + 1
  result.nStates = m.nStates
  result.nBins = nb
  result.counts = newSeq[int](nb)
  result.binLo = newSeq[int](nb)
  result.binHi = newSeq[int](nb)
  for b in 0..<nb:
    result.binLo[b] = high(int)
    result.binHi[b] = low(int)
  for c in 0..<m.nClasses:
    for i in 0..<m.teams[c].len:
      let st = int(m.teams[c][i])
      var bin = ((st + m.nStates) * nb) div span
      if bin < 0: bin = 0
      if bin >= nb: bin = nb - 1
      inc result.counts[bin]
      if st < result.binLo[bin]: result.binLo[bin] = st
      if st > result.binHi[bin]: result.binHi[bin] = st
      inc result.total
  for b in 0..<nb:
    if result.counts[b] == 0:
      result.binLo[b] = 0
      result.binHi[b] = 0

proc histogramText*(h: StateHistogram): string =
  ## One line per bin, e.g. `[-64..-57]  1234 #####`.
  for b in 0..<h.nBins:
    let bar = "#".repeat(min(40, h.counts[b] div max(1, h.total div 200 + 1)))
    result.add &"  [{h.binLo[b]:>4}..{h.binHi[b]:>4}] {h.counts[b]:>8} {bar}\n"

# ── per-input automata confidence table ──────────────────────────────────────

type
  InputConfidence* = object
    bit*: int
    name*: string
    nAutomata*: int
    meanCommitment*: float
    settledFraction*: float
    constant*: bool   ## zero variance across the sample set (information-free)

proc perInputConfidence*(m: TmMachine, spec: FeatureSpec,
                         samples: openArray[DiagSample],
                         threshold = 0.5): seq[InputConfidence] =
  ## For every raw bit: the mean commitment of the automata of its positive
  ## literal and its negation, across every clause and class. Inputs the TM is
  ## confident about sit high; a bit whose automata are all parked at the
  ## boundary (e.g. a constant input) sits near 0. `constant` flags zero-variance
  ## bits so a high commitment on a never-written input is not mistaken for
  ## learning (the shipped gun's bits 38/39 are the canonical example).
  result = newSeq[InputConfidence](m.nBits)
  var consts: seq[int]
  if samples.len > 0:
    for b in 0..<m.nBits:
      let v0 = samples[0].lits[b]
      var cst = true
      for i in 1..<samples.len:
        if samples[i].lits[b] != v0:
          cst = false
          break
      if cst: consts.add b
  for b in 0..<m.nBits:
    result[b].bit = b
    result[b].name = spec.describe(b)
    result[b].constant = b in consts
    var sum = 0.0
    var settled = 0
    var n = 0
    for c in 0..<m.nClasses:
      for cl in 0..<m.nClauses:
        let base = cl * m.nLiterals
        let cmPos = stateCommitment(m.nStates, int(m.teams[c][base + b]))
        let cmNeg = stateCommitment(m.nStates,
                                    int(m.teams[c][base + b + m.nBits]))
        sum += cmPos + cmNeg
        if cmPos >= threshold: inc settled
        if cmNeg >= threshold: inc settled
        inc n, 2
    result[b].nAutomata = n
    result[b].meanCommitment = if n > 0: sum / n.float else: 0.0
    result[b].settledFraction = if n > 0: settled.float / n.float else: 0.0

proc rankedInputConfidence*(conf: seq[InputConfidence]): seq[InputConfidence] =
  ## Least confident first — the inputs whose automata are still undecided.
  result = conf
  result.sort(proc(a, b: InputConfidence): int =
    result = cmp(a.meanCommitment, b.meanCommitment)
    if result == 0: result = cmp(a.bit, b.bit))

# ── 2. CLAUSE DIVERSITY ──────────────────────────────────────────────────────

type
  DiversityResult* = object
    jaccardPos*: float      ## mean pairwise Jaccard among positive clauses
    jaccardNeg*: float      ## ... among negative clauses
    jaccardOverall*: float  ## pooled over both polarities
    nPosClauses*: int
    nNegClauses*: int
    nPairs*: int
    posValid*: bool
    negValid*: bool
    valid*: bool
    skipEmpty*: bool

proc jaccard*(a, b: openArray[int]): float =
  ## |a n b| / |a u b| for two sorted, unique literal sets. Two empty sets are
  ## identical (1.0); an empty against a non-empty set share nothing (0.0).
  if a.len == 0 and b.len == 0: return 1.0
  if a.len == 0 or b.len == 0: return 0.0
  var i = 0
  var j = 0
  var inter = 0
  while i < a.len and j < b.len:
    if a[i] == b[j]: inc inter; inc i; inc j
    elif a[i] < b[j]: inc i
    else: inc j
  let uni = a.len + b.len - inter
  if uni == 0: 1.0 else: inter.float / uni.float

proc polarityDiversity(m: TmMachine, polarity: int, skipEmpty: bool):
    tuple[sum: float, npairs, nclauses: int] =
  var sets: seq[seq[int]]
  let cl0 = if polarity > 0: 0 else: m.half
  let cl1 = if polarity > 0: m.half else: m.nClauses
  for c in 0..<m.nClasses:
    for cl in cl0..<cl1:
      let ls = m.clauseLits(c, cl)
      if skipEmpty and ls.len == 0: continue
      sets.add ls
  result.nclauses = sets.len
  for i in 0..<sets.len:
    for j in (i + 1)..<sets.len:
      result.sum += jaccard(sets[i], sets[j])
      inc result.npairs

proc clauseDiversity*(m: TmMachine, skipEmpty = true): DiversityResult =
  ## Mean pairwise Jaccard overlap of the INCLUDED-literal sets, within the same
  ## polarity (positive vs positive, negative vs negative). Low-to-moderate is
  ## healthy; near 1.0 means the clauses are the same rule wearing 50 hats;
  ## near 0 with few samples usually means memorising individual ticks. Empty
  ## clauses carry no rule and are skipped by default (`skipEmpty`).
  result.skipEmpty = skipEmpty
  let p = polarityDiversity(m, 1, skipEmpty)
  let n = polarityDiversity(m, -1, skipEmpty)
  result.nPosClauses = p.nclauses
  result.nNegClauses = n.nclauses
  result.posValid = p.npairs > 0
  result.negValid = n.npairs > 0
  result.jaccardPos = if p.npairs > 0: p.sum / p.npairs.float else: 0.0
  result.jaccardNeg = if n.npairs > 0: n.sum / n.npairs.float else: 0.0
  let pairs = p.npairs + n.npairs
  result.nPairs = pairs
  result.valid = pairs > 0
  result.jaccardOverall =
    if pairs > 0: (p.sum + n.sum) / pairs.float else: 0.0

# ── 3. CHURN ─────────────────────────────────────────────────────────────────

type
  ChurnResult* = object
    samples*: int
    epochs*: int
    nAutomata*: int
    nClauses*: int
    totalFlips*: int
    totalClauseChanges*: int
    flipRate*: float           ## automaton flips per training sample (fraction)
    flipRatePer100*: float     ## same, x100 (readable)
    clauseChangeRate*: float   ## clause-set changes per training sample
    clauseChangePer100*: float
    flipFirst*: float
    flipLast*: float
    clauseFirst*: float
    clauseLast*: float
    flipTrend*: string
    clauseTrend*: string
    windowSize*: int
    flipWindows*: seq[float]   ## mean per-sample flip rate per window
    clauseWindows*: seq[float]
    trained*: TmMachine

proc binMeans*(s: seq[float], window: int): seq[float] =
  if s.len == 0: return
  let w = max(1, window)
  var i = 0
  while i < s.len:
    let j = min(s.len, i + w)
    var sum = 0.0
    for k in i..<j: sum += s[k]
    result.add sum / (j - i).float
    i = j

proc firstLastMeans*(s: seq[float]): tuple[first, last: float] =
  if s.len == 0: return (0.0, 0.0)
  let q = max(1, s.len div 5)
  var a = 0.0
  var b = 0.0
  for i in 0..<q: a += s[i]
  for i in (s.len - q)..<s.len: b += s[i]
  (a / q.float, b / q.float)

proc classifyChurnTrend*(s: seq[float]): string =
  ## "falling" is the healthy sign; "flat"/"rising" at a high rate is
  ## fidgeting; "frozen" means it never moved at all.
  if s.len < 4: return "insufficient"
  let (first, last) = firstLastMeans(s)
  if first <= 1e-12 and last <= 1e-12: return "frozen"
  if first <= 1e-12: return "rising"
  let ratio = last / first
  if ratio < 0.6: "falling"
  elif ratio > 1.6: "rising"
  else: "flat"

proc churnTrace*(tmpl: TmMachine, samples: openArray[DiagSample],
                 epochs = 1, seed = 777'u64, window = 100,
                 shuffle = true): ChurnResult =
  ## Train a fresh machine while measuring, per training sample, both churn
  ## levels: the fraction of AUTOMATA crossing the flip boundary, and the
  ## fraction of CLAUSES whose included-literal set changed. `shuffle = false`
  ## replays the samples in their given (temporal) order once per epoch, which
  ## mirrors a live gun training as bullets resolve.
  result.epochs = max(1, epochs)
  result.samples = samples.len
  result.windowSize = max(1, window)
  var m = tmpl
  m.resetMachine(seed)
  result.nAutomata = m.nClasses * m.nClauses * m.nLiterals
  result.nClauses = m.nClasses * m.nClauses
  if result.nAutomata == 0:
    result.trained = m
    return
  var prevInc = newSeq[uint8](result.nAutomata)
  for c in 0..<m.nClasses:
    let team = m.teams[c]
    let gbase = c * m.nClauses * m.nLiterals
    for i in 0..<team.len:
      prevInc[gbase + i] = uint8(team[i] > 0)
  var flipSeries = newSeq[float]()
  var clauseSeries = newSeq[float]()
  for _ in 0..<result.epochs:
    let idx =
      if shuffle: shuffledIndices(samples.len, m.rng)
      else: (block:
        var id = newSeq[int](samples.len)
        for i in 0..<samples.len: id[i] = i
        id)
    for k in idx:
      m.trainSample(samples[k].lits, samples[k].label)
      var flips = 0
      var clauseChanges = 0
      for c in 0..<m.nClasses:
        let team = m.teams[c]
        let gbase = c * m.nClauses * m.nLiterals
        for cl in 0..<m.nClauses:
          var changed = false
          let base = cl * m.nLiterals
          let gb = gbase + base
          for lit in 0..<m.nLiterals:
            let nowInc = uint8(team[base + lit] > 0)
            if nowInc != prevInc[gb + lit]:
              inc flips
              changed = true
              prevInc[gb + lit] = nowInc
          if changed: inc clauseChanges
      result.totalFlips += flips
      result.totalClauseChanges += clauseChanges
      flipSeries.add flips.float / result.nAutomata.float
      clauseSeries.add clauseChanges.float / result.nClauses.float
  result.trained = m
  let nObs = flipSeries.len
  if nObs > 0:
    result.flipRate = result.totalFlips.float /
      (nObs.float * result.nAutomata.float)
    result.clauseChangeRate = result.totalClauseChanges.float /
      (nObs.float * result.nClauses.float)
  result.flipRatePer100 = result.flipRate * 100.0
  result.clauseChangePer100 = result.clauseChangeRate * 100.0
  let ff = firstLastMeans(flipSeries)
  let cf = firstLastMeans(clauseSeries)
  result.flipFirst = ff.first
  result.flipLast = ff.last
  result.clauseFirst = cf.first
  result.clauseLast = cf.last
  result.flipTrend = classifyChurnTrend(flipSeries)
  result.clauseTrend = classifyChurnTrend(clauseSeries)
  result.flipWindows = binMeans(flipSeries, result.windowSize)
  result.clauseWindows = binMeans(clauseSeries, result.windowSize)

# ── 4. VOTE DISAGREEMENT ─────────────────────────────────────────────────────

type
  VoteDisagreementResult* = object
    overall*: float
    perClass*: seq[float]
    perClassSamples*: seq[int]
    nPairs*: int             ## (class, sample) observations that had a vote

proc voteDisagreement*(m: TmMachine, samples: openArray[DiagSample]):
    VoteDisagreementResult =
  ## For each class, over the samples where it casts at least one clause vote:
  ## the fraction of FIRING clauses whose own polarity disagrees with the sign
  ## of the class's total vote. A class whose clauses always pull together has
  ## disagreement ~0; a fuzzy class boundary or too few clauses makes it high.
  ## A tied class vote (mixed signs) counts every firing clause as disagreeing.
  result.perClass = newSeq[float](m.nClasses)
  result.perClassSamples = newSeq[int](m.nClasses)
  var classSum = newSeq[float](m.nClasses)
  for c in 0..<m.nClasses:
    for s in samples:
      var vote = 0.0
      var firing: seq[int]
      for cl in 0..<m.nClauses:
        let o = m.tmEval(m.teams[c], s.lits, cl, learning = false)
        if o == 1'u8:
          vote += m.tmPolarity(cl)
          firing.add cl
      if firing.len == 0: continue
      let majSign = if vote > 0.0: 1 elif vote < 0.0: -1 else: 0
      var disagree = 0
      for cl in firing:
        let pol = if m.tmPolarity(cl) > 0.0: 1 else: -1
        if pol != majSign: inc disagree
      classSum[c] += disagree.float / firing.len.float
      inc result.perClassSamples[c]
      inc result.nPairs
  var totalSum = 0.0
  var totalN = 0
  for c in 0..<m.nClasses:
    if result.perClassSamples[c] > 0:
      result.perClass[c] = classSum[c] / result.perClassSamples[c].float
      totalSum += result.perClass[c]
      inc totalN
  result.overall = if totalN > 0: totalSum / totalN.float else: 0.0

# ── the convenience driver + the one-line health summary ─────────────────────

type
  AutomataDiag* = object
    machine*: TmMachine
    settledness*: SettlednessResult
    diversity*: DiversityResult
    churn*: ChurnResult
    disagreement*: VoteDisagreementResult
    histogram*: StateHistogram
    inputConfidence*: seq[InputConfidence]
    summary*: string

proc settlednessVerdict*(s: SettlednessResult): string =
  if s.overallMean >= 0.65: "settled"
  elif s.overallMean >= 0.40: "settling"
  else: "wavering"

proc diversityVerdict*(d: DiversityResult): string =
  if not d.valid: "n/a"
  elif d.jaccardOverall < 0.15: "diverse"
  elif d.jaccardOverall <= 0.70: "moderate"
  else: "redundant"

proc churnVerdict*(c: ChurnResult): string =
  case c.flipTrend
  of "falling": "settling"
  of "rising": "fidgeting (rising)"
  of "frozen": "frozen"
  of "insufficient", "not-measured": c.flipTrend
  else:
    if c.flipRatePer100 > 5.0: "fidgeting" else: "quiet"

proc disagreementVerdict*(d: VoteDisagreementResult): string =
  if d.overall < 0.15: "coherent"
  elif d.overall < 0.35: "fuzzy"
  else: "fragmented"

proc healthLine*(d: AutomataDiag): string =
  ## ONE line to glance at: settledness / diversity / churn trend / disagreement,
  ## each with its own verdict word.
  &"settledness={d.settledness.overallMean:.3f} ({settlednessVerdict(d.settledness)}) | " &
  &"diversity={d.diversity.jaccardOverall:.3f} ({diversityVerdict(d.diversity)}) | " &
  &"churn={d.churn.flipRatePer100:.3f}/100 {d.churn.flipTrend} " &
  &"({churnVerdict(d.churn)}) | " &
  &"disagreement={d.disagreement.overall:.3f} " &
  &"({disagreementVerdict(d.disagreement)})"

proc automataDiagnostics*(tmpl: TmMachine, samples: openArray[DiagSample],
                          spec: FeatureSpec, epochs = 1, seed = 777'u64,
                          settleThreshold = 0.5, nHistBins = 9,
                          window = 100, measureChurn = true): AutomataDiag =
  ## Train a fresh machine on `samples` (measuring churn along the way unless
  ## `measureChurn` is false) and compute every automata-level metric on it.
  if measureChurn:
    result.churn = churnTrace(tmpl, samples, epochs, seed, window)
    result.machine = result.churn.trained
  else:
    result.machine = trainModel(tmpl, samples, epochs, seed)
    result.churn.epochs = max(1, epochs)
    result.churn.samples = samples.len
    result.churn.trained = result.machine
    result.churn.flipTrend = "not-measured"
    result.churn.clauseTrend = "not-measured"
  result.settledness = settledness(result.machine, settleThreshold)
  result.diversity = clauseDiversity(result.machine)
  result.disagreement = voteDisagreement(result.machine, samples)
  result.histogram = stateHistogram(result.machine, nHistBins)
  result.inputConfidence = perInputConfidence(result.machine, spec, samples,
                                              settleThreshold)
  result.summary = healthLine(result)

proc automataVerdict*(d: AutomataDiag): string =
  ## One word for the whole trajectory: settling / fidgeting / collapsed / other.
  if d.churn.flipTrend == "falling" and d.settledness.overallMean >= 0.40:
    "settling"
  elif d.churn.flipTrend in ["flat", "rising"] and d.churn.flipRatePer100 > 5.0:
    "fidgeting"
  elif d.settledness.overallMean < 0.25 and d.churn.flipTrend == "frozen":
    "collapsed"
  elif d.disagreement.overall > 0.35 and d.settledness.overallMean < 0.40:
    "collapsed"
  elif d.settledness.overallMean >= 0.65 and d.churn.flipTrend == "frozen":
    "converged-static"
  else:
    "mixed"

proc formatAutomataReport*(d: AutomataDiag): string =
  ## The full readout: health line, verdict, state histogram, per-input
  ## confidence (least confident first) and per-class disagreement.
  result.add "health: " & d.summary & "\n"
  result.add "verdict: " & automataVerdict(d) & "\n"
  result.add &"automata: n={d.settledness.nAutomata} settled>={d.settledness.threshold:.2f} " &
    &"frac={d.settledness.overallSettledFraction:.3f}\n"
  result.add &"churn: flips/sample={d.churn.flipRate:.6f} ({d.churn.flipRatePer100:.3f}/100, " &
    &"{d.churn.flipTrend}: {d.churn.flipFirst:.6f}->{d.churn.flipLast:.6f}) " &
    &"clause-change/sample={d.churn.clauseChangeRate:.6f} " &
    &"({d.churn.clauseChangePer100:.3f}/100, {d.churn.clauseTrend})\n"
  result.add &"diversity: pos={d.diversity.jaccardPos:.3f} " &
    &"neg={d.diversity.jaccardNeg:.3f} overall={d.diversity.jaccardOverall:.3f} " &
    &"(pairs={d.diversity.nPairs}, posClauses={d.diversity.nPosClauses}, " &
    &"negClauses={d.diversity.nNegClauses})\n"
  result.add "state histogram:\n" & histogramText(d.histogram)
  result.add "per-input confidence (least confident first):\n"
  let ranked = rankedInputConfidence(d.inputConfidence)
  for i in 0..<min(12, ranked.len):
    let tag = if ranked[i].constant: " CONST" else: ""
    result.add &"  bit{ranked[i].bit:>2} {ranked[i].name:<28} " &
      &"meanCommit={ranked[i].meanCommitment:.3f} settled={ranked[i].settledFraction:.3f}{tag}\n"
  var cbits: seq[int]
  for ic in d.inputConfidence:
    if ic.constant: cbits.add ic.bit
  result.add &"constant inputs (zero variance): {cbits}\n"
  result.add "disagreement per class:"
  for c in 0..<d.disagreement.perClass.len:
    result.add &" c{c}={d.disagreement.perClass[c]:.3f}"
  result.add "\n"