## Address Decoder Element (ADE) layer — clean-room implementation.
##
## This module implements the ADE thresholded random projection as described in
##
##   "BitBrain and Sparse Binary Coincidence (SBC) memories",
##   Frontiers in Neuroinformatics 17:1125844, 2023.
##
## It was written **from the published algorithm description only**. No code was
## copied from the reference C program (`full_mnist_2048.c`) that ships alongside
## the paper, which is GPL-3.0-or-later, (c) 2022 The University of Manchester.
## Keeping this a clean-room implementation avoids imposing that licence on this
## repository. The published defaults (`scale = 64`, `center = 127`) make this
## module numerically identical to the reference's scoring rule.
##
## Mechanism
## ---------
## A *synapse* is `(inputIndex, polarity)`. An ADE with `width` synapses scores an
## input vector as
##
##     raw   = Σ_j polarity_j * (input[inputIndex_j] - center)
##     score = scale * raw
##
## and **fires** iff `score >= threshold`. The reference C uses `scale = 64`
## and `center = 127` on raw `uint8` MNIST pixels, i.e. exactly the default
## parameters here. Multi-width ADs (the paper's best setup uses widths in
## {6, 8, 10, 12}) detect features at different scales.
##
## Storage
## -------
## Synapses are stored as a flat `seq[int32]` of signed 1-based codes: the sign
## is the polarity, `abs(code) - 1` is the 0-based input index. This is byte-for-
## byte the layout of the reference `AD*_2048` weight files, so a pretrained AD
## can be loaded with a single `memcpy`-equivalent read.
##
## Unsupervised phase (threshold homeostasis)
## ------------------------------------------
## The paper learns each ADE's homeostatic threshold so that it fires at a target
## rate (~1%). `accumulateFiring` + `adaptThresholds` below implement the
## paper's simple deterministic controller: over an interval of `interval` inputs
## the firing counts are compared against `targetRate * interval`, and each ADE's
## threshold is nudged up or down by `step`.
##
## The paper's *additional* unsupervised mechanisms are deliberately NOT
## implemented here (the task only requires them to be described):
##
##   * **Hebbian "longevity" learning.** Each synapse carries a longevity
##     counter. When an ADE fires, the smallest contributor to the threshold
##     crossing has its longevity decremented and the largest contributor has it
##     incremented. After an interval, any synapse whose longevity drops below a
##     critical value is *retired* and replaced: a new input index is drawn by the
##     same sampling mechanism used at initialisation, with longevity reset to
##     the default. This lets each ADE "home in" on a feature. It is purely
##     unsupervised.
##   * **Metropolis-Hastings input-position sampling.** Synapse input positions
##     are drawn from a distribution proportional to `sqrt` of the global input
##     activity histogram (the `sqrt` both stabilises the Poisson uncertainty and
##     flattens the distribution so that synapses can land slightly outside the
##     training support). `initRandomAddressDecoder` instead samples uniformly at
##     random, which is the paper's stated initialisation.
##   * **Spatial clustering.** For image/volumetric data the sampled positions
##     may be constrained to a local region (rejection sampling until the
##     constraint holds) so each ADE sees a coherent receptive field.

import std/random

const
  DefaultScale* = 64
    ## Score multiplier. 64 reproduces the reference C's `yang = ±64`.
  DefaultCenter* = 127
    ## Input centre. 127 reproduces the reference C's `(pixel - 127)`.

type
  AddressDecoder* = object
    ## An AD: `nAde` Address Decoder Elements, all of `width` synapses.
    nAde*: int                ## the paper's `w`
    width*: int               ## the paper's `n` (synapses per ADE)
    codes*: seq[int32]        ## nAde*width; sign = polarity, |code|-1 = input index
    thresholds*: seq[int32]   ## per-ADE firing thresholds
    scale*: int               ## score multiplier (default 64)
    center*: int              ## input centre (default 127)
    fireCounts*: seq[int32]   ## firing accumulator for threshold homeostasis

proc initAddressDecoder*(nAde, width: int, scale = DefaultScale,
                         center = DefaultCenter): AddressDecoder =
  ## Allocate an AD with zeroed synapses. Callers either fill `codes` directly
  ## (e.g. from a pretrained weight file) or use `initRandomAddressDecoder`.
  doAssert nAde > 0, "nAde must be positive"
  doAssert width > 0, "width must be positive"
  result.nAde = nAde
  result.width = width
  result.scale = scale
  result.center = center
  result.codes = newSeq[int32](nAde * width)
  result.thresholds = newSeq[int32](nAde)
  result.fireCounts = newSeq[int32](nAde)

proc synCode(inputIdx: int, polarity: int8): int32 {.inline.} =
  ## Encode `(index, polarity)` as the reference's signed 1-based code.
  doAssert inputIdx >= 0
  doAssert polarity == 1'i8 or polarity == -1'i8
  result = int32(polarity) * int32(inputIdx + 1)

proc initRandomAddressDecoder*(nAde, width, inputWidth: int, rng: var Rand,
                               scale = DefaultScale, center = DefaultCenter,
                               threshold: int32 = 0'i32): AddressDecoder =
  ## Randomly initialise an AD. For each ADE `width` **distinct** input indices
  ## are drawn uniformly from `0 ..< inputWidth` (multapses are disallowed, as in
  ## the paper) with a random ± polarity. Every ADE starts at `threshold`.
  doAssert inputWidth >= width,
    "cannot draw " & $width & " distinct indices from " & $inputWidth
  result = initAddressDecoder(nAde, width, scale, center)
  var chosen = newSeq[int32](width)
  for i in 0 ..< nAde:
    result.thresholds[i] = threshold
    var k = 0
    while k < width:
      let idx = int32(rng.rand(inputWidth - 1))
      var dup = false
      for p in 0 ..< k:
        if chosen[p] == idx:
          dup = true
          break
      if not dup:
        chosen[k] = idx
        inc k
    for j in 0 ..< width:
      let pol = if rng.rand(1) == 0: 1'i8 else: -1'i8
      result.codes[i * width + j] = synCode(int(chosen[j]), pol)

proc score*[T: SomeInteger](ad: AddressDecoder, input: openArray[T],
                            ade: int): int {.inline.} =
  ## Signed, centred, scaled score of ADE `ade` on `input`.
  ## `score = scale * Σ_j polarity_j * (input[idx_j] - center)`.
  doAssert ade >= 0 and ade < ad.nAde
  var raw = 0
  let base = ade * ad.width
  for j in 0 ..< ad.width:
    let c = ad.codes[base + j]
    let idx =
      if c > 0'i32: int(c) - 1
      else: int(-c) - 1
    let pol = if c > 0'i32: 1 else: -1
    raw += pol * (int(input[idx]) - ad.center)
  result = ad.scale * raw

proc fires*[T: SomeInteger](ad: AddressDecoder, input: openArray[T],
                            ade: int): bool {.inline.} =
  ## An ADE fires iff its score reaches its threshold (`score >= threshold`).
  score(ad, input, ade) >= int(ad.thresholds[ade])

proc activeList*[T: SomeInteger](ad: AddressDecoder, input: openArray[T],
                                 active: var seq[int32]) =
  ## Fill `active` with the indices of the ADEs that fire on `input`.
  ## This is the sparse representation consumed by the SBC memory.
  active.setLen(0)
  for i in 0 ..< ad.nAde:
    if fires(ad, input, i):
      active.add int32(i)

proc fireCount*[T: SomeInteger](ad: AddressDecoder, input: openArray[T]): int =
  ## Number of ADEs in `ad` that fire on `input`.
  for i in 0 ..< ad.nAde:
    if fires(ad, input, i):
      inc result

proc accumulateFiring*[T: SomeInteger](ad: var AddressDecoder,
                                       input: openArray[T]) =
  ## Add this input's firing events to the homeostasis accumulator.
  for i in 0 ..< ad.nAde:
    if fires(ad, input, i):
      inc ad.fireCounts[i]

proc resetFiringCounts*(ad: var AddressDecoder) =
  ## Zero the homeostasis accumulator (does not touch thresholds).
  for i in 0 ..< ad.nAde:
    ad.fireCounts[i] = 0

proc adaptThresholds*(ad: var AddressDecoder, interval: int, targetRate: float,
                      step: int = 1) =
  ## One step of the paper's homeostatic threshold controller.
  ##
  ## `fireCounts` is expected to hold the firing counts accumulated over the last
  ## `interval` inputs (see `accumulateFiring`). Each ADE whose observed rate is
  ## above/below `targetRate` has its threshold raised/lowered by `step`. The
  ## counters are then reset, ready for the next interval.
  ##
  ## This is a deterministic integer controller: given the same inputs and the
  ## same starting thresholds, it always produces the same thresholds.
  doAssert interval > 0, "interval must be positive"
  let expected = targetRate * float(interval)
  for i in 0 ..< ad.nAde:
    let excess = float(ad.fireCounts[i]) - expected
    if excess > 0.5:
      ad.thresholds[i] += int32(step)
    elif excess < -0.5:
      ad.thresholds[i] -= int32(step)
    ad.fireCounts[i] = 0'i32

proc memoryBytes*(ad: AddressDecoder): int =
  ## Bytes held by this AD (synapse codes, thresholds, firing counters).
  ad.codes.len * sizeof(int32) +
    ad.thresholds.len * sizeof(int32) +
    ad.fireCounts.len * sizeof(int32)