## Pure unit tests for the AUTOMATA-LEVEL metrics added to tm_diag:
## settledness, clause diversity, churn, vote disagreement, the state
## histogram, per-input confidence and the health/verdict helpers.
##
## Run: nim c -r --path:common_libs common_libs/tests/test_tm_automata_diag.nim

import tm_diag/diagnostics
import std/strutils

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

# ── state conventions ────────────────────────────────────────────────────────

proc testStateConventions() =
  check "commitment(0) = 0 (on the flip boundary)",
        stateCommitment(64, 0) == 0.0
  check "commitment(+64) = 1 (extreme Include)",
        stateCommitment(64, 64) == 1.0
  check "commitment(-64) = 1 (extreme Exclude)",
        stateCommitment(64, -64) == 1.0
  check "commitment(+32) = 0.5", abs(stateCommitment(64, 32) - 0.5) < 1e-12
  check "commitment is symmetric", stateCommitment(64, 17) == stateCommitment(64, -17)
  check "include predicate matches state>0",
        stateIncluded(1) and not stateIncluded(0) and not stateIncluded(-1)

# ── settledness ──────────────────────────────────────────────────────────────

proc settledMachine(): TmMachine =
  ## nBits=4, 2 classes, 4 clauses/class, 8 literals.
  ## class0/clause0 = [+64, -64, 0, +32, 0,0,0,0] -> mean commitment 0.625,
  ## 3 of 8 automata settled at threshold 0.5.
  result = newMachine(4, 2, 4, 64, 3.0, 1)
  result.teams[0][0] = 64
  result.teams[0][1] = -64
  result.teams[0][2] = 0
  result.teams[0][3] = 32

proc testSettledness() =
  let m = settledMachine()
  let s = settledness(m, 0.5)
  check "settledness covers every clause",
        s.clauses.len == m.nClasses * m.nClauses
  check "settledness nAutomata = classes*clauses*literals",
        s.nAutomata == m.nClasses * m.nClauses * m.nLiterals
  var c0: ClauseSettledness
  for c in s.clauses:
    if c.cls == 0 and c.index == 0: c0 = c
  check "clause0 mean commitment = 0.3125",
        abs(c0.meanCommitment - 0.3125) < 1e-12
  check "clause0 settled fraction = 3/8",
        abs(c0.settledFraction - 3.0 / 8.0) < 1e-12
  check "clause0 positive polarity", c0.polarity == 1
  check "overall settled fraction = 3/64",
        abs(s.overallSettledFraction - 3.0 / 64.0) < 1e-12
  check "overall mean commitment = 2.5/64",
        abs(s.overallMean - 2.5 / 64.0) < 1e-12
  # trend helper
  var lo = s
  lo.overallMean = 0.2
  var hi = s
  hi.overallMean = 0.8
  check "settlednessTrend rising", settlednessTrend(lo, hi) == "rising"
  check "settlednessTrend falling", settlednessTrend(hi, lo) == "falling"
  check "settlednessTrend flat", settlednessTrend(s, s) == "flat"

# ── state histogram ──────────────────────────────────────────────────────────

proc testHistogram() =
  let m = settledMachine()
  let h = stateHistogram(m, 9)
  var total = 0
  for c in h.counts: total += c
  check "histogram counts sum to every automaton",
        total == m.nClasses * m.nClauses * m.nLiterals
  check "the -64 state lands in the first bin", h.counts[0] >= 1
  check "the +64 state lands in the last bin", h.counts[^1] >= 1
  check "histogram text is non-empty", histogramText(h).len > 0

# ── per-input confidence ─────────────────────────────────────────────────────

proc testInputConfidence() =
  let m = settledMachine()
  let spec = draftTMSpec()
  let samples = @[
    makeSample(49, (block:
      var r = newSeq[int](49)
      for i in 0..<49: r[i] = 1
      r), 0),
    makeSample(49, (block:
      var r = newSeq[int](49)
      for i in 0..<49: r[i] = 1
      r), 0),
  ]
  # all-ones samples make every bit constant, so all should be flagged.
  let conf = perInputConfidence(m, spec, samples, 0.5)
  check "confidence table covers every bit", conf.len == m.nBits
  check "confidence nAutomata per bit = classes*clauses*2",
        conf[0].nAutomata == m.nClasses * m.nClauses * 2
  check "a zero-variance bit is flagged constant", conf[0].constant
  let ranked = rankedInputConfidence(conf)
  check "ranked confidence is ascending",
        ranked.len == conf.len and
        ranked[0].meanCommitment <= ranked[^1].meanCommitment

# ── diversity / jaccard ──────────────────────────────────────────────────────

proc testJaccard() =
  check "identical sets -> 1", jaccard(@[1, 2, 3], @[1, 2, 3]) == 1.0
  check "disjoint sets -> 0", jaccard(@[1, 2], @[3, 4]) == 0.0
  check "two empty sets -> 1", jaccard(@[], @[]) == 1.0
  check "empty vs non-empty -> 0", jaccard(@[], @[1]) == 0.0
  check "partial overlap -> |n|/|u|",
        abs(jaccard(@[1, 2], @[2, 3]) - 1.0 / 3.0) < 1e-12

proc diversityMachine(): TmMachine =
  ## nBits=4, 2 classes, 6 clauses/class (half=3 positive, 3 negative).
  ## class0 positive clauses: {0,1}, {0,1}, {2,3}; negative clauses empty.
  result = newMachine(4, 2, 6, 64, 3.0, 1)
  for lit in [0, 1]:
    result.teams[0][0 * 8 + lit] = 1
    result.teams[0][1 * 8 + lit] = 1
  for lit in [2, 3]:
    result.teams[0][2 * 8 + lit] = 1

proc testDiversity() =
  let m = diversityMachine()
  let d = clauseDiversity(m)
  check "positive clauses: {0,1},{0,1},{2,3} mean Jaccard = 1/3",
        abs(d.jaccardPos - 1.0 / 3.0) < 1e-12
  check "positive diversity is valid", d.posValid
  check "empty negative clauses are skipped", not d.negValid
  check "overall pools the same 3 pairs", d.nPairs == 3
  check "overall Jaccard = 1/3",
        abs(d.jaccardOverall - 1.0 / 3.0) < 1e-12

# ── churn ────────────────────────────────────────────────────────────────────

proc testChurnHelpers() =
  check "binMeans bins [1,1,3,3,5] with window 2 -> [1,3,5]",
        binMeans(@[1.0, 1.0, 3.0, 3.0, 5.0], 2) == @[1.0, 3.0, 5.0]
  let fl = firstLastMeans(@[0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0])
  check "firstLastMeans first ~0", fl.first < 1e-9
  check "firstLastMeans last ~1", abs(fl.last - 1.0) < 1e-9
  check "churn trend falling",
        classifyChurnTrend(@[0.9, 0.8, 0.7, 0.6, 0.5, 0.1, 0.05, 0.01]) == "falling"
  check "churn trend flat",
        classifyChurnTrend(@[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]) == "flat"
  check "churn trend rising",
        classifyChurnTrend(@[0.01, 0.02, 0.03, 0.1, 0.5, 0.6, 0.7, 0.9]) == "rising"
  check "churn trend frozen",
        classifyChurnTrend(@[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) == "frozen"

proc testChurnTrace() =
  var samples: seq[DiagSample]
  for i in 0..<120:
    var raw = newSeq[int](49)
    for b in 0..<49: raw[b] = (i + b) mod 2
    raw[0] = 1
    raw[45] = if i mod 2 == 0: 1 else: 0
    samples.add makeSample(49, raw, if raw[45] == 1: 2 else: 1, i)
  let tmpl = newMachine(49, 3, 40, 64, 3.0, 1)
  let c = churnTrace(tmpl, samples, epochs = 1, seed = 5, window = 30)
  check "churn nAutomata = classes*clauses*literals",
        c.nAutomata == 3 * 40 * 98
  check "churn nClauses = classes*clauses", c.nClauses == 3 * 40
  check "churn rate is non-negative", c.flipRate >= 0.0
  check "churn per100 = rate*100",
        abs(c.flipRatePer100 - c.flipRate * 100.0) < 1e-12
  check "churn windows are non-empty", c.flipWindows.len > 0
  check "churn trained machine has the right shape",
        c.trained.teams.len == 3 and c.trained.teams[0].len == 40 * 98

# ── vote disagreement ────────────────────────────────────────────────────────

proc testVoteDisagreement() =
  # all-agree machine: only positive clauses fire -> disagreement 0
  var agree = newMachine(4, 2, 4, 64, 3.0, 1)
  agree.teams[0][0 * 8 + 0] = 1
  agree.teams[0][1 * 8 + 1] = 1
  let allOnes = @[makeSample(4, @[1, 1, 1, 1], 0)]
  let da = voteDisagreement(agree, allOnes)
  check "all-positive firing clauses -> disagreement 0",
        abs(da.overall - 0.0) < 1e-12
  check "one class observed", da.perClassSamples[0] == 1

  # split machine: one positive and one negative clause fire -> tied vote -> 1
  var split = newMachine(4, 2, 4, 64, 3.0, 1)
  split.teams[0][0 * 8 + 0] = 1
  split.teams[0][2 * 8 + 2] = 1
  let ds = voteDisagreement(split, allOnes)
  check "a tied vote counts every firing clause as disagreeing",
        abs(ds.overall - 1.0) < 1e-12

# ── health line / verdict / report ───────────────────────────────────────────

proc testHealthLine() =
  let spec = draftTMSpec()
  var samples: seq[DiagSample]
  for i in 0..<200:
    var raw = newSeq[int](49)
    for b in 0..<49: raw[b] = if (i * 7 + b) mod 3 == 0: 1 else: 0
    raw[0] = if i mod 2 == 0: 1 else: 0
    raw[45] = if i mod 3 == 0: 1 else: 0
    let label = if raw[0] == 1 and raw[45] == 1: 2 elif raw[0] == 1: 1 else: 0
    samples.add makeSample(49, raw, label, i)
  let tmpl = newMachine(49, 3, 40, 64, 3.0, 1)
  let d = automataDiagnostics(tmpl, samples, spec, epochs = 3, seed = 7,
                              settleThreshold = 0.5, nHistBins = 9)
  check "health line is non-empty and mentions every metric",
        d.summary.len > 0 and "settledness" in d.summary and
        "diversity" in d.summary and "churn" in d.summary and
        "disagreement" in d.summary
  check "automataVerdict is non-empty", automataVerdict(d).len > 0
  check "formatAutomataReport contains the histogram and table",
        "state histogram" in formatAutomataReport(d) and
        "per-input confidence" in formatAutomataReport(d)
  check "disagreement per-class array is the right size",
        d.disagreement.perClass.len == 3
  check "input confidence covers all bits", d.inputConfidence.len == 49

when isMainModule:
  testStateConventions()
  testSettledness()
  testHistogram()
  testInputConfidence()
  testJaccard()
  testDiversity()
  testChurnHelpers()
  testChurnTrace()
  testVoteDisagreement()
  testHealthLine()
  echo ""
  if failures > 0:
    echo failures, " / ", checks, " check(s) FAILED"
    quit(1)
  echo "All ", checks, " automata-diag unit checks passed."