## tm_diag/tm_core.nim — a compact, deterministic Granmo Table 2/3 Tsetlin
## Machine (multiclass), with an introspection-friendly clause layout.
##
## This mirrors the corrected core in common_libs/guns/tm_pattern.nim
## (`tmEval` / `tmForward` / `tmLearnDir` and the Eq. 6 empty-clause bootstrap)
## so the diagnostics get demonstrated on the SAME algorithm the real gun uses.
## It is a standalone copy because the gun's core is private and pulls in the
## gun harness; here everything is pure and offline.
##
## Layout: `teams[c][cl * nLiterals + lit]`, literal `i` is the positive literal
## and `i + nBits` its negation — identical to the gun.

import std/math

type
  TmRng* = object
    s*: uint64

  TmMachine* = object
    nBits*: int
    nLiterals*: int
    nClauses*: int
    nClasses*: int
    half*: int
    nStates*: int
    sValue*: float
    teams*: seq[seq[int16]]
    rng*: TmRng

proc seedRng*(seed: uint64): TmRng =
  result.s = seed
  if result.s == 0: result.s = 0x9e3779b97f4a7c15'u64

proc nextU64*(r: var TmRng): uint64 =
  r.s = r.s xor (r.s shl 13)
  r.s = r.s xor (r.s shr 7)
  r.s = r.s xor (r.s shl 17)
  r.s

proc rand01*(r: var TmRng): float =
  ## Uniform [0,1).
  (r.nextU64() shr 11).float / 9007199254740992.0

proc newMachine*(nBits, nClasses: int, nClauses = 40, nStates = 64,
                 sValue = 3.0, seed = 12345'u64): TmMachine =
  result.nBits = nBits
  result.nLiterals = 2 * nBits
  result.nClauses = nClauses
  result.nClasses = nClasses
  result.half = nClauses div 2
  result.nStates = nStates
  result.sValue = sValue
  result.rng = seedRng(seed)
  result.teams = newSeq[seq[int16]](nClasses)
  for c in 0..<nClasses:
    result.teams[c] = newSeq[int16](nClauses * result.nLiterals)

proc cloneMachine*(m: TmMachine): TmMachine =
  result = m
  result.teams = newSeq[seq[int16]](m.nClasses)
  for c in 0..<m.nClasses:
    result.teams[c] = m.teams[c]

proc resetMachine*(m: var TmMachine, seed: uint64) =
  ## Wipe every clause back to the Exclude boundary and reseed the RNG.
  for c in 0..<m.nClasses:
    for i in 0..<m.teams[c].len: m.teams[c][i] = 0
  m.rng = seedRng(seed)

proc tmPolarity*(m: TmMachine, cl: int): float {.inline.} =
  if cl < m.half: 1.0 else: -1.0

proc tmEval*(m: TmMachine, team: seq[int16], lits: openArray[uint8],
             cl: int, learning: bool): uint8 =
  let base = cl * m.nLiterals
  var hasInc = false
  for lit in 0..<m.nLiterals:
    if team[base + lit] > 0:
      hasInc = true
      if lits[lit] == 0'u8: return 0'u8
  if hasInc: return 1'u8
  # Eq. 6: the empty conjunction is vacuously true while learning, false when
  # classifying. Without this the all-Exclude init deadlocks.
  return if learning: 1'u8 else: 0'u8

proc tmForward*(m: TmMachine, team: seq[int16], lits: openArray[uint8],
                cache: var seq[uint8]): float =
  var v = 0.0
  for cl in 0..<m.nClauses:
    let o = tmEval(m, team, lits, cl, learning = false)
    cache[cl] = tmEval(m, team, lits, cl, learning = true)
    v += tmPolarity(m, cl) * float(o)
  clamp(v, -float(m.half), float(m.half))

proc tmLearnDir*(m: var TmMachine, team: var seq[int16],
                 lits: openArray[uint8], cache: seq[uint8],
                 vote, d: float) =
  ## One Granmo update of one class team with desired vote direction `d`.
  let T = float(m.half)
  let pFeedback = (T - d * vote) / (2.0 * T)
  if pFeedback <= 0.0: return
  for cl in 0..<m.nClauses:
    if m.rng.rand01() >= pFeedback: continue
    let pol = m.tmPolarity(cl)
    let cOut = cache[cl]
    let base = cl * m.nLiterals
    if pol * d > 0.0:
      # Type I (Table 2) collapsed to the resulting state move.
      for lit in 0..<m.nLiterals:
        var st = int(team[base + lit])
        if lits[lit] == 1'u8:
          if cOut == 1'u8:
            if m.rng.rand01() < (m.sValue - 1.0) / m.sValue:
              st = min(st + 1, m.nStates)
          else:
            if m.rng.rand01() < 1.0 / m.sValue:
              st = max(st - 1, -m.nStates)
        else:
          if m.rng.rand01() < 1.0 / m.sValue:
            st = max(st - 1, -m.nStates)
        team[base + lit] = int16(st)
    else:
      # Type II (Table 3): penalise exclusion of a zero literal when firing.
      if cOut == 1'u8:
        for lit in 0..<m.nLiterals:
          if lits[lit] == 0'u8:
            if team[base + lit] <= 0:
              team[base + lit] = int16(min(int(team[base + lit]) + 1, m.nStates))

proc votesOf*(m: TmMachine, lits: openArray[uint8]): seq[float] =
  result = newSeq[float](m.nClasses)
  var cache = newSeq[uint8](m.nClauses)
  for c in 0..<m.nClasses:
    result[c] = tmForward(m, m.teams[c], lits, cache)

proc predictClass*(m: TmMachine, lits: openArray[uint8]): int =
  ## Hard argmax over the per-class votes; ties break to the lowest class.
  var best = 0
  var bestV = -Inf
  var cache = newSeq[uint8](m.nClauses)
  for c in 0..<m.nClasses:
    let v = tmForward(m, m.teams[c], lits, cache)
    if v > bestV:
      bestV = v
      best = c
  best

proc trainSample*(m: var TmMachine, lits: openArray[uint8], label: int) =
  var votes = newSeq[float](m.nClasses)
  var caches = newSeq[seq[uint8]](m.nClasses)
  for c in 0..<m.nClasses:
    caches[c] = newSeq[uint8](m.nClauses)
    votes[c] = tmForward(m, m.teams[c], lits, caches[c])
  for c in 0..<m.nClasses:
    let d = if c == label: 1.0 else: -1.0
    tmLearnDir(m, m.teams[c], lits, caches[c], votes[c], d)

proc clauseFires*(m: TmMachine, cls, cl: int, lits: openArray[uint8]): bool =
  ## Classification-mode output of one clause: true iff it is non-empty and
  ## every included literal is 1.
  let base = cl * m.nLiterals
  var hasInc = false
  for lit in 0..<m.nLiterals:
    if m.teams[cls][base + lit] > 0:
      hasInc = true
      if lits[lit] == 0'u8: return false
  hasInc

proc clauseLits*(m: TmMachine, cls, cl: int): seq[int] =
  ## The literal indices included by one clause (>0 state).
  let base = cl * m.nLiterals
  for lit in 0..<m.nLiterals:
    if m.teams[cls][base + lit] > 0: result.add lit