## Sparse Binary Coincidence (SBC) memory — clean-room implementation.
##
## Implements the supervised half of the BitBrain algorithm as described in
##
##   "BitBrain and Sparse Binary Coincidence (SBC) memories",
##   Frontiers in Neuroinformatics 17:1125844, 2023.
##
## Written from the published algorithm only; see `ade.nim` for the clean-room
## note. The reference C is GPL-3.0 (c) University of Manchester and was not
## copied.
##
## Mechanism
## ---------
## Two Address Decoders (ADs) sit on the two axes of a 2-D memory. A pair of
## *simultaneously firing* ADEs `(i, j)` is a **coincidence** and addresses one
## memory cell. That cell holds a class bitmask with one bit per class
## (`nClasses` bits, one-hot encoding in the paper's default).
##
## Learning is **idempotent**: `learn` *sets* the bit for the observed class;
## setting it again is a no-op. There is no clearing, no learning rate, no decay
## and no epoch — one pass through the data is a complete supervised training run,
## and a second pass changes nothing.
##
## Inference uses the *same* address decoding: for each observed coincidence every
## class bit is read and the set bits are **counted** per class. Counts are summed
## across SBCs by the `bitbrain` container and the argmax wins.
##
## Bit layout
## ----------
## The bit for `(i, j, class)` lives at `((i * nAde) + j) * nClasses + class`, so
## all `nClasses` bits of one coincidence are contiguous. This differs from the
## reference C's layout but is a bijection onto the same set of triples; the
## learned rule is identical.

import std/bitops

type
  Sbc* = object
    ## A 2-D coincidence memory with a class-bit depth.
    nAde*: int        ## number of ADEs on each axis (the paper's `w`)
    nClasses*: int    ## number of classes (the paper's `D`)
    bits*: seq[uint32] ## packed bit tensor, nAde*nAde*nClasses bits

proc initSbc*(nAde, nClasses: int): Sbc =
  ## Allocate a zeroed SBC (nothing is known yet).
  doAssert nAde > 0, "nAde must be positive"
  doAssert nClasses > 0, "nClasses must be positive"
  result.nAde = nAde
  result.nClasses = nClasses
  let nbits = nAde * nAde * nClasses
  result.bits = newSeq[uint32]((nbits + 31) div 32)

proc clear*(sbc: var Sbc) =
  ## Forget everything. This is how a monotone memory is wiped (e.g. when a bot
  ## switches enemy and must not carry state across battles).
  for i in 0 ..< sbc.bits.len:
    sbc.bits[i] = 0'u32

proc bitIndex(sbc: Sbc, i, j, class: int): int {.inline.} =
  ((i * sbc.nAde) + j) * sbc.nClasses + class

proc bitAt*(sbc: Sbc, i, j, class: int): bool {.inline.} =
  ## Read one memory bit. Exposed mainly so harnesses can inspect the exact rule.
  doAssert i >= 0 and i < sbc.nAde
  doAssert j >= 0 and j < sbc.nAde
  doAssert class >= 0 and class < sbc.nClasses
  let bit = bitIndex(sbc, i, j, class)
  (sbc.bits[bit shr 5] and (1'u32 shl (bit and 31))) != 0'u32

proc learn*(sbc: var Sbc, rowActive, colActive: openArray[int32],
            class: int): int =
  ## Set the `class` bit for every coincidence between a firing row ADE and a
  ## firing column ADE. Returns the number of bits that were newly set (0 if the
  ## sample added no information, e.g. it was already learned).
  doAssert class >= 0 and class < sbc.nClasses, "class out of range"
  let D = sbc.nClasses
  for r in rowActive:
    let i = int(r)
    for c in colActive:
      let j = int(c)
      let bit = ((i * sbc.nAde) + j) * D + class
      let w = bit shr 5
      let m = 1'u32 shl (bit and 31)
      if (sbc.bits[w] and m) == 0'u32:
        sbc.bits[w] = sbc.bits[w] or m
        inc result

proc infer*(sbc: Sbc, rowActive, colActive: openArray[int32],
            counts: var seq[int]) =
  ## Count, per class, how many observed coincidences have their class bit set.
  ## `counts` is *accumulated into* (not reset), so a container can sum several
  ## SBCs. It must be at least `nClasses` long.
  doAssert counts.len >= sbc.nClasses, "counts buffer too small"
  let D = sbc.nClasses
  for r in rowActive:
    let i = int(r)
    for c in colActive:
      let j = int(c)
      let base = ((i * sbc.nAde) + j) * D
      for k in 0 ..< D:
        let bit = base + k
        if (sbc.bits[bit shr 5] and (1'u32 shl (bit and 31))) != 0'u32:
          inc counts[k]

proc memoryBytes*(sbc: Sbc): int =
  ## Bytes held by the packed bit tensor.
  sbc.bits.len * sizeof(uint32)

proc occupancy*(sbc: Sbc): float =
  ## Fraction of the bit tensor that is set (diagnostic).
  var setBits = 0
  for w in sbc.bits:
    setBits += countSetBits(w)
  result = float(setBits) / float(sbc.nAde * sbc.nAde * sbc.nClasses)