## Pure unit tests for the tm_diag diagnostics kit.
## Covers: FeatureSpec rendering, contribution counting, dead-input detection,
## majority baseline, clause introspection and trend classification.
##
## Run: nim c -r --path:common_libs common_libs/tests/test_tm_diag.nim

import tm_diag/diagnostics

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

# ── 1. FeatureSpec rendering ─────────────────────────────────────────────────

proc testRendering() =
  let s = draftTMSpec()
  check "draft spec is 49 bits", s.nBits == 49
  check "draft WALLS is 8 bits", s.blocks[0].count + s.blocks[1].count == 8
  check "draft describe(0) = dist-wall<50", s.describe(0) == "dist-wall<50"
  check "draft describe(45) = lat DEAD-ON",
        s.describe(45) == "lat DEAD-ON -18..+18"
  check "draft describe(48) = lat>+72", s.describe(48) == "lat>+72"
  check "positive literal", s.describeLiteral(0) == "dist-wall<50"
  check "negated literal", s.describeLiteral(49) == "NOT dist-wall<50"
  check "clause sentence",
        s.describeClause(@[0, 45], 2) ==
          "IF dist-wall<50 AND lat DEAD-ON -18..+18 THEN class=2"
  check "clause sentence with negation",
        s.describeClause(@[0, 49 + 45], 1) ==
          "IF dist-wall<50 AND NOT lat DEAD-ON -18..+18 THEN class=1"
  check "empty clause renders TRUE",
        s.describeClause(@[], 3) == "IF TRUE (empty clause) THEN class=3"
  check "no class suffix when cls<0",
        s.describeClause(@[0]) == "IF dist-wall<50"

  let p = tmPatternSpec()
  check "tm_pattern spec is 40 bits", p.nBits == 40
  check "tm_pattern UNUSED block is last and 2 bits",
        p.blocks[^1].name == "UNUSED" and p.blocks[^1].count == 2
  check "tm_pattern describe(38) = UNUSED-38", p.describe(38) == "UNUSED-38"

# ── 2. majority baseline / data checks ───────────────────────────────────────

proc testDataChecks() =
  let dc = dataChecks(@[0, 0, 0, 1, 2], 3)
  check "majority count", dc.classCounts == @[3, 1, 1]
  check "majority class", dc.majorityClass == 0
  check "majority share 3/5", abs(dc.majorityShare - 0.6) < 1e-9
  check "majority = accuracy baseline",
        abs(dc.majorityAccuracy - 0.6) < 1e-9
  check "over-threshold flagged", dc.overThreshold

  let dc2 = dataChecks(@[0, 0, 1, 1, 2, 2, 3, 3], 4)
  check "balanced data is not flagged", not dc2.overThreshold
  check "balanced majority = 1/4", abs(dc2.majorityShare - 0.25) < 1e-9

# ── 3. contribution counting / dead inputs (hand-built model) ────────────────

proc tinySpec(): FeatureSpec =
  var s = FeatureSpec()
  s.addBlock("f", 4, @["f0", "f1", "f2", "f3"])
  s

proc tinyModel(): TmMachine =
  ## nBits=4, 2 classes, 4 clauses/class.
  ##   class0 clause0 includes +f0 ; class1 clause0 includes +f1.
  ## All other clauses are empty.
  result = newMachine(4, 2, 4, 64, 3.0, 1)
  result.teams[0][0 * result.nLiterals + 0] = 1
  result.teams[1][0 * result.nLiterals + 1] = 1

proc testContributions() =
  let m = tinyModel()
  let spec = tinySpec()
  let samples = @[
    makeSample(4, @[1, 0, 0, 0], 0),
    makeSample(4, @[0, 1, 0, 0], 1),
  ]
  let infos = clauseInfo(m, samples, spec)
  var class0Info: ClauseInfo
  for c in infos:
    if c.cls == 0 and c.index == 0: class0Info = c
  check "clauseInfo: class0/clause0 fires on sample A", class0Info.votes == 1
  check "clauseInfo: clause length 1", class0Info.length == 1
  check "clauseInfo: rendered", class0Info.text == "IF f0 THEN class=0"
  check "clauseInfo: positive polarity", class0Info.polarity == 1

  let summ = clauseSummary(infos)
  check "clauseSummary: 8 clauses total", summ.totalClauses == 8
  check "clauseSummary: 6 empty clauses", summ.emptyClauses == 6
  check "clauseSummary: 2 firing clauses", summ.firedAtLeastOnce == 2
  check "clauseSummary: mean length 1.0", abs(summ.meanLength - 1.0) < 1e-9
  check "clauseSummary: 2 positive fired", summ.posFired == 2 and summ.negFired == 0

  let contribs = featureContributions(m, samples, spec)
  check "contribution f0 = 1.0", abs(contribs[0].weighted - 1.0) < 1e-9 and
        contribs[0].appearances == 1
  check "contribution f1 = 1.0", abs(contribs[1].weighted - 1.0) < 1e-9 and
        contribs[1].appearances == 1
  check "contribution f2 = 0", contribs[2].weighted == 0.0
  check "contribution f3 = 0", contribs[3].weighted == 0.0

  let dead = deadInputs(contribs)
  check "deadInputs flags f2,f3", dead == @[2, 3]
  check "neverUsedInputs flags f2,f3", neverUsedInputs(contribs) == @[2, 3]
  let ranked = rankedInputs(contribs)
  check "rankedInputs puts f0/f1 first", ranked[0].bit in {0, 1} and ranked[1].bit in {0, 1}
  check "rankedInputs puts f2/f3 last", ranked[^1].bit in {2, 3}

  check "constantInputs finds the info-free f2,f3",
        constantInputs(samples, 4) == @[2, 3]

proc testConstantInputs() =
  let samples = @[
    makeSample(4, @[1, 0, 0, 0], 0),
    makeSample(4, @[1, 1, 0, 0], 1),
  ]
  check "constantInputs finds f0,f2,f3", constantInputs(samples, 4) == @[0, 2, 3]

proc testNecessaryLiterals() =
  let m = tinyModel()
  let spec = tinySpec()
  let samples = @[
    makeSample(4, @[1, 0, 0, 0], 0),
    makeSample(4, @[0, 1, 0, 0], 1),
  ]
  check "necessaryLiterals(class0) = {+f0}", necessaryLiterals(m, samples, 0) == @[0]
  check "necessaryLiterals(class1) = {+f1}", necessaryLiterals(m, samples, 1) == @[1]

# ── 4. trend classification ──────────────────────────────────────────────────

proc testTrend() =
  check "rising trend", classifyTrend(@[0.1, 0.3, 0.5, 0.7]) == "rising"
  check "flat trend", classifyTrend(@[0.5, 0.5, 0.5, 0.5]) == "flat"
  check "falling trend", classifyTrend(@[0.7, 0.5, 0.3, 0.1]) == "falling"
  check "rising-then-falling",
        classifyTrend(@[0.3, 0.7, 0.75, 0.5]) ==
          "rising-then-falling (noise fitting)"

# ── 5. machineFromTeams ──────────────────────────────────────────────────────

proc testMachineFromTeams() =
  let m = tinyModel()
  let m2 = machineFromTeams(4, 2, 4, 64, 3.0, m.teams)
  check "machineFromTeams copies teams", m2.teams[0] == m.teams[0]
  check "machineFromTeams predicts consistently",
        m2.predictClass(makeSample(4, @[1, 0, 0, 0], 0).lits) ==
        m.predictClass(makeSample(4, @[1, 0, 0, 0], 0).lits)

when isMainModule:
  testRendering()
  testDataChecks()
  testContributions()
  testConstantInputs()
  testNecessaryLiterals()
  testTrend()
  testMachineFromTeams()
  echo ""
  if failures > 0:
    echo failures, " / ", checks, " check(s) FAILED"
    quit(1)
  echo "All ", checks, " tm_diag unit checks passed."