## Does the BITBRAIN gun (rack id 17) actually ENGAGE the ADE+SBC network?
##
## Linking is not engagement. These checks drive the real gun over synthetic
## states and assert the observable consequences of a live network:
##   * the network is BUILT on first predict and is 0 bytes before that (lazy);
##   * two DIFFERENT inputs give DIFFERENT class outputs (the ADE layer is
##     discriminating, not a constant);
##   * learning a sample CHANGES the memory (the SBC head is writing);
##   * `learn` is idempotent on a bitset SBC (setting the same class twice is a
##     no-op) and monotone on a counted one;
##   * the shipped knobs really move RAM and the derived geometry;
##   * the input is the CONFIGURED block set: a disabled block shrinks it, a
##     wider block grows it, and `TR_BITBRAIN_INPUT` pins it exactly;
##   * construction is RNG-clean (the default path cannot perturb the selector).
##
## No Java, no battle, no fixtures.
##
## Run: nim c -r --path:common_libs common_libs/tests/test_bitbrain_net.nim

import std/[os, random, math, strutils, sequtils]
import gun_harness/gun_interface
import gun_harness/virtual_bullets
import gun_harness/selector
import guns/bitbrain_net

const BitbrainNetId = 17
const PatternId = 5

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

proc clearEnv() =
  for n in BitbrainNetEnvNames: delEnv(n)
  for n in RackGunNames: delEnv("TR_RACK_" & n)

proc mkState(t: int, ex, ey, sx, sy, eh: float): WorldState =
  WorldState(arenaWidth: 800, arenaHeight: 600, tick: t,
             enemyX: ex, enemyY: ey, enemyHeading: eh, enemySpeed: 8,
             selfX: sx, selfY: sy, selfHeading: 0, selfSpeed: 8,
             selfEnergy: 100, enemyEnergy: 100, selfRadarHeading: 0)

proc armNet() =
  ## A small, fast, deterministic configuration: TR_BITBRAIN_NET must be on for
  ## the gun to engage at all.
  putEnv(BBN_NET_ENV, "1")
  putEnv(BBN_NADES_ENV, "64")
  putEnv(BBN_CLASSES_ENV, "8")
  putEnv(BBN_MINOBS_ENV, "1")
  putEnv(BBN_INPUT_ENV, "24")

# ── registration / default-off ───────────────────────────────────────────────

proc testRegistration() =
  clearEnv()
  check "rack: BITBRAIN is registered at id 17 and defaults to off",
        RackGunNames.len == 18 and
        RackGunNames[BitbrainNetId] == "BITBRAIN" and
        DefaultRackMembership[BitbrainNetId] == rmOff
  check "rack: the shipped default still admits only Pattern",
        admittedGuns(18, rm1v1, DefaultRackMembership) == @[PatternId] and
        admittedGuns(18, rmMelee, DefaultRackMembership) == @[PatternId]
  check "gate: the gun is NOT spawned under the default rack",
        not vBulletAdmitted(BitbrainNetId, rm1v1, DefaultRackMembership, true)
  var g = initBitbrainNetGun()
  check "default OFF: TR_BITBRAIN_NET unset leaves the gun disabled",
        not g.enabled and g.networkBytes() == 0
  # With the switch on but the gun not built, predict must not spend anything.
  armNet()
  var g2 = initBitbrainNetGun()
  check "lazy: an enabled but unused gun still holds 0 bytes",
        g2.enabled and g2.networkBytes() == 0

proc testRackEnv() =
  clearEnv()
  putEnv("TR_RACK_BITBRAIN", "both")
  putEnv(BBN_NET_ENV, "1")
  let m = loadRackMembership()
  check "env: TR_RACK_BITBRAIN=both + TR_BITBRAIN_NET=1 admits id 17",
        m[BitbrainNetId] == rmBoth and
        vBulletAdmitted(BitbrainNetId, rm1v1, m, true)
  clearEnv()

# ── the knobs really resolve ─────────────────────────────────────────────────

proc testKnobs() =
  clearEnv()
  armNet()
  putEnv(BBN_INPUT_ENV, "33")
  putEnv(BBN_CLASSES_ENV, "17")
  putEnv(BBN_NADES_ENV, "128")
  putEnv(BBN_WIDTHS_ENV, "3,5,7,9")
  putEnv(BBN_SPAN_ENV, "25")
  putEnv(BBN_MODE_ENV, "counted")
  putEnv(BBN_DECAY_EVERY_ENV, "11")
  putEnv(BBN_DECAY_SHIFT_ENV, "2")
  putEnv(BBN_MINOBS_ENV, "7")
  var g = initBitbrainNetGun()
  check "knob: input width", g.inputWidth == 33
  check "knob: nClasses (output resolution)", g.nClasses == 17
  check "knob: nAde per AD", g.nAde == 128
  check "knob: clause widths (one AD per width)", g.widths == @[3, 5, 7, 9]
  check "knob: class half-range", g.maxDeg == 25.0
  check "knob: SBC mode", g.mode == smCounted
  check "knob: counted-mode decay knobs",
        g.decayEvery == 11 and g.decayShift == 2
  check "knob: minObs", g.minObs == 7
  putEnv(BBN_MODE_ENV, "bitset")
  check "knob: bitset mode", initBitbrainNetGun().mode == smBitset
  clearEnv()

# ── the input is a CONFIGURED SET OF BLOCKS ──────────────────────────────────

proc testFeatureBlocks() =
  clearEnv()
  armNet()
  delEnv(BBN_INPUT_ENV)
  var g = initBitbrainNetGun()
  check "features: the shipped block set is 47 slots",
        derivedSlots(g.blockWidths) == BB_DEFAULT_SLOTS and
        g.inputWidth == BB_DEFAULT_SLOTS
  check "features: every shipped block is enabled by default",
        g.blockWidths.allIt(it > 0)
  # disable two blocks -> the input shrinks by exactly their widths
  delEnv(BBN_INPUT_ENV)
  putEnv(BBN_FEATURES_ENV, "epos:0,bull:0")
  var g2 = initBitbrainNetGun()
  check "features: a disabled block removes its slots",
        g2.inputWidth == BB_DEFAULT_SLOTS - 8 - 4
  # widen one block -> the input grows by exactly the extra slots
  putEnv(BBN_FEATURES_ENV, "dist:9")
  var g3 = initBitbrainNetGun()
  check "features: a wider block adds its slots",
        g3.inputWidth == BB_DEFAULT_SLOTS - 10 + 18
  # TR_BITBRAIN_INPUT pins the width whatever the blocks say
  putEnv(BBN_FEATURES_ENV, "dist:9")
  putEnv(BBN_INPUT_ENV, "64")
  var g4 = initBitbrainNetGun()
  check "features: TR_BITBRAIN_INPUT pins the width",
        g4.inputWidth == 64
  putEnv(BBN_INPUT_ENV, "9")
  var g5 = initBitbrainNetGun()
  check "features: TR_BITBRAIN_INPUT can shrink below the block total",
        g5.inputWidth == 9
  # the unknown-block warning path must not change the widths
  putEnv(BBN_FEATURES_ENV, "banana,hzn:3")
  var g6 = initBitbrainNetGun()
  check "features: an unknown block is ignored, the rest still apply",
        g6.blockWidths[blockNameIndex("hzn")] == 3 and
        g6.blockWidths[blockNameIndex("epos")] == 4
  clearEnv()

# ── the network really engages ───────────────────────────────────────────────

proc testNetworkEngages() =
  clearEnv()
  armNet()
  var g = initBitbrainNetGun()
  let st = mkState(1, 600, 300, 100, 300, 0)
  var input = buildInput(g, st)
  check "engage: the input vector is exactly inputWidth long",
        input.len == g.inputWidth
  check "engage: the input is a 0/255 thermometer, not a dead constant",
        input.allIt(it == 0'u8 or it == BBN_SBC_VALUE) and
        input.anyIt(it == BBN_SBC_VALUE)
  check "engage: no network is built before the gun is used",
        g.networkBytes() == 0
  discard predict(g, st, 11.0)
  check "engage: the gun builds the network on first use",
        g.networkBytes() > 0
  # two different inputs -> different class outputs
  g.buildNet()
  let a = buildInput(g, mkState(1, 600, 300, 100, 300, 0))
  let b = buildInput(g, mkState(1, 200, 500, 700, 100, 90))
  check "engage: two different inputs give DIFFERENT bit vectors", a != b
  # An untrained SBC has no evidence, so the argmax is trivially 0. The real
  # engagement test is: teach it a rule, then check the readout separates.
  for i in 0 ..< 300:
    let half = (i mod 2 == 0)
    let x = buildInput(g, mkState(i, if half: 200.0 else: 700.0,
                                  150.0 + float(i mod 7), 100, 500, float(i mod 3)))
    g.learnSample(x, if half: 0 else: 7)
  var labels = newSeq[int](100)
  for i in 0 ..< 100:
    let half = (i mod 2 == 0)
    let x = buildInput(g, mkState(i, if half: 200.0 else: 700.0,
                                  150.0 + float(i mod 7), 100, 500, float(i mod 3)))
    labels[i] = g.inferClass(x)
  var differ = false
  for i in 1 ..< labels.len:
    if labels[i] != labels[0]: differ = true
  check "engage: a TRAINED ADE+SBC head returns DIFFERENT classes for different states",
        differ
  var agree = 0
  for i in 0 ..< labels.len:
    if (labels[i] == 0) == (i mod 2 == 0): inc agree
  check "engage: the readout separates the two taught populations",
        agree >= 80
  # learning changes the memory
  g.resetLearning()
  let before = g.sbcBytesOf()
  let beforeOcc = g.sbcsOf()[0].occupancy()
  g.learnSample(a, 3)
  check "engage: learning one sample WRITES the SBC memory",
        g.sbcsOf()[0].occupancy() > beforeOcc
  check "engage: memoryBytes is unchanged by a learn (no realloc)",
        g.sbcBytesOf() == before
  # bitset learn is idempotent
  clearEnv()
  armNet()
  putEnv(BBN_MODE_ENV, "bitset")
  putEnv(BBN_WIDTHS_ENV, "2,3,4")   # denser clauses so coincidences actually fire
  var gb = initBitbrainNetGun()
  gb.buildNet()
  gb.learnSample(a, 2)
  let occ1 = gb.net.sbcs[0].occupancy()
  for _ in 0 ..< 49: gb.learnSample(a, 2)
  check "engage: a bitset learn really writes the memory", occ1 > 0.0
  check "engage: bitset learn is idempotent (50x learn == 1x)",
        abs(gb.net.sbcs[0].occupancy() - occ1) < 1e-12
  # counted learn is monotone: the same cell keeps gaining evidence
  putEnv(BBN_MODE_ENV, "counted")
  putEnv(BBN_WIDTHS_ENV, "2,3,4")
  var gc = initBitbrainNetGun()
  gc.buildNet()
  gc.learnSample(a, 2)
  let ev2 = gc.evidenceFor(a, 2)
  for _ in 0 ..< 4: gc.learnSample(a, 2)
  check "engage: counted learn is MONOTONE (5 learns beat 1)",
        ev2 > 0 and gc.evidenceFor(a, 2) > ev2
  check "engage: the counted mode really holds saturating counters",
        gc.sbcsOf()[0].mode == smCounted
  clearEnv()

# ── the readout is a fine-grained correction, and it is gated ────────────────

proc testReadout() =
  clearEnv()
  armNet()
  putEnv(BBN_MINOBS_ENV, "1000")
  var g = initBitbrainNetGun()
  let st = mkState(1, 600, 300, 100, 300, 0)
  discard predict(g, st, 11.0)
  check "readout: below minObs the correction is exactly zero",
        g.lastShiftDeg == 0.0 and not isWarmedUp(g)
  putEnv(BBN_MINOBS_ENV, "1")
  var g2 = initBitbrainNetGun()
  # feed a stream so labels resolve and the learner engages
  var t = 0
  for _ in 0 ..< 400:
    let ft = float(t)
    let st = mkState(t, 100.0 + 2.0 * ft, 200.0 + 1.1 * ft, 400, 300,
                     20.0 * sin(ft * 0.2))
    discard predict(g2, st, 11.0)
    inc t
  check "readout: a resolved stream trains the net", g2.trained > 100
  check "readout: the net becomes warmed up past minObs", isWarmedUp(g2)
  check "readout: the correction is a bounded, fine-grained angle",
        g2.lastShiftDeg >= -g2.maxDeg - 1e-9 and
        g2.lastShiftDeg <= g2.maxDeg + 1e-9
  check "readout: the gun's output is a POINT (Pattern base + correction)",
        predict(g2, mkState(400, 500, 300, 400, 300, 0), 11.0).x != 0.0
  check "readout: ADE threshold adaptation ran (online calibration)",
        g2.adapts > 0
  clearEnv()

# ── RNG parity ───────────────────────────────────────────────────────────────

proc testRngClean() =
  clearEnv()
  randomize(1234)
  let a = rand(1_000_000)
  randomize(1234)
  armNet()
  var g = initBitbrainNetGun()
  g.buildNet()
  discard predict(g, mkState(1, 300, 300, 100, 100, 0), 11.0)
  let b = rand(1_000_000)
  check "parity: building + running the net does not perturb the global RNG",
        a == b
  clearEnv()

testRegistration()
testRackEnv()
testKnobs()
testFeatureBlocks()
testNetworkEngages()
testReadout()
testRngClean()

if failures > 0:
  echo "\n", failures, " check(s) FAILED"
  quit(1)
echo "\nAll BITBRAIN (ADE+SBC) engagement checks passed."