#!/usr/bin/env python3
"""M2 milestone: brain-inspired SLEEP CLEANUP for Mind-1B.

Question this explores: if a synthetic memory bank keeps writing fast Hebbian
tags all day, they saturate and drown real memories in noise (catastrophic
interference). This script compares a sleep-cleanup rule with a naive
raw-magnitude retention control; it is not a language-model ablation.

The biology we copy (three mechanisms, all real):

1. Synaptic Homeostasis (SHY, Tononi and Cirelli). Waking = net synaptic
   POTENTIATION; every trace grows, signal and noise alike, until capacity and
   signal-to-noise saturate. Sleep (slow-wave / NREM) applies a GLOBAL
   proportional DOWNSCALE of all synapses. Strong, replayed traces survive the
   downscale; weak spurious ones fall below detection and vanish. Net: the
   pattern is preserved, the noise floor is reset, capacity is recovered.

2. Replay-gated consolidation (hippocampal sharp-wave ripples -> cortex).
   Recent episodes are replayed; synapses that REACTIVATE often during replay
   are tagged important (synaptic tag-and-capture) and get burned into the slow
   long-term store. This is the "keep" signal.

3. Synaptic pruning (use-it-or-lose-it; microglia remove weak synapses).
   Traces that survive neither replay nor magnitude are set to exactly ZERO and
   removed. In a TERNARY network this is free: 0 is a native weight state, so
   pruning == ternarizing toward 0. This is the "remove" signal.

Keep/remove rule per synapse i:
    utility_i = |tag_i| * (1 + lambda * replay_count_i)     # magnitude x salience
    tag_i  <- gamma * tag_i                                  # (1) SHY downscale
    if utility_i >= theta:  burn sign(tag_i) into base W     # (2) consolidate/keep
    else:                   tag_i <- 0                       # (3) prune/remove

Run: python3 sleep_cleanup.py [--cycles 8 --synapses 20000]
"""

import argparse
import os

for _v in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"):
    os.environ.setdefault(_v, "1")

import numpy as np


class HebbianMemory:
    """A bank of synapses: a frozen ternary base + a plastic fast tag each.

    Recall of a stored task = does the consolidated pattern still match the
    task's true sign pattern on that task's synapses?
    """

    def __init__(self, n, seed=0):
        self.n = n
        self.base = np.zeros(n, dtype=np.int8)          # consolidated ternary {-1,0,1}
        self.tag = np.zeros(n, dtype=np.float32)        # plastic fast weight
        self.replay = np.zeros(n, dtype=np.float32)     # reactivation count this day
        self.rng = np.random.default_rng(seed)

    def wake_learn(self, signal_idx, signal_sign, noise_level=0.6):
        """A day of use: potentiate the task's synapses (signal) and let random
        coincidences potentiate everything a little (noise). Net potentiation,
        exactly the SHY wake phase."""
        # signal: strong, sign-consistent Hebbian writes on the task pathway
        self.tag[signal_idx] += signal_sign * self.rng.uniform(0.8, 1.2,
                                                               size=len(signal_idx))
        self.replay[signal_idx] += 1.0
        # noise: weak random coincident firing across ALL synapses
        self.tag += self.rng.normal(0, noise_level, size=self.n).astype(np.float32)

    def sleep_cleanup(self, gamma=0.85, lam=3.0, theta=1.1):
        """One night. Returns stats. This is the whole mechanism."""
        # (2) salience from replay + magnitude, computed BEFORE downscale
        utility = np.abs(self.tag) * (1.0 + lam * self.replay)
        # (1) SHY: global homeostatic downscale of every plastic trace
        self.tag *= gamma
        # keep vs remove
        keep = utility >= theta
        # (2) consolidate survivors: burn their sign into the ternary base
        newly = keep & (self.base == 0)
        self.base[keep] = np.sign(self.tag[keep]).astype(np.int8)
        # (3) prune the rest to exactly zero (synaptic elimination)
        self.tag[~keep] = 0.0
        pruned = int((~keep).sum())
        # reset the nightly replay counter (a new day accumulates fresh salience)
        consolidated = int(newly.sum())
        self.replay[:] = 0.0
        return {
            "kept": int(keep.sum()),
            "pruned": pruned,
            "newly_consolidated": consolidated,
            "base_nonzero": int((self.base != 0).sum()),
        }

    def recall(self, signal_idx, signal_sign):
        """Fraction of a task's synapses whose consolidated sign is correct."""
        got = self.base[signal_idx]
        return float(np.mean(got == signal_sign))


def make_task(n, k, rng):
    idx = rng.choice(n, size=k, replace=False)
    sign = rng.choice(np.array([-1, 1], dtype=np.int8), size=k)
    return idx, sign


def run(n, cycles, k, seed=0):
    rng = np.random.default_rng(seed)
    tasks = [make_task(n, k, rng) for _ in range(cycles)]

    sleeper = HebbianMemory(n, seed=1)
    insomniac = HebbianMemory(n, seed=1)

    print("continual learning: %d synapses, %d tasks, %d signal synapses/task"
          % (n, cycles, k))
    print("recall = mean(consolidated sign == true sign) over a task's synapses\n")
    header = "%-8s | %-28s | %-28s" % ("", "WITH sleep cleanup", "NO cleanup (tags accrue)")
    print(header)
    print("%-8s | %-13s %-14s | %-13s %-14s"
          % ("cycle", "avg recall", "base nonzero", "avg recall", "tag saturation"))
    print("-" * len(header))

    for c in range(cycles):
        idx, sign = tasks[c]
        sleeper.wake_learn(idx, sign)
        stats = sleeper.sleep_cleanup()

        # insomniac: same wake, but NO sleep. It consolidates by raw magnitude
        # only (the naive "just keep the biggest tags" strategy), tags accrue.
        insomniac.wake_learn(idx, sign)
        thresh = np.quantile(np.abs(insomniac.tag), 0.90)
        big = np.abs(insomniac.tag) >= thresh
        insomniac.base[big] = np.sign(insomniac.tag[big]).astype(np.int8)

        seen = tasks[: c + 1]
        r_sleep = np.mean([sleeper.recall(i, s) for i, s in seen])
        r_ins = np.mean([insomniac.recall(i, s) for i, s in seen])
        sat = float(np.mean(np.abs(insomniac.tag) > 1.0))
        print("%-8d | %11.1f%%  %12d  | %11.1f%%  %13.0f%%"
              % (c + 1, r_sleep * 100, stats["base_nonzero"],
                 r_ins * 100, sat * 100))

    print("\ninterpretation:")
    print("  - WITH cleanup: recall on ALL past tasks stays high, base stays SPARSE")
    print("    (noise pruned to 0 each night = recovered capacity + cheaper inference)")
    print("  - NO cleanup: tags saturate, old tasks get overwritten by new noise")
    print("    = catastrophic interference, the exact failure sleep evolved to fix")
    print("  - ternary makes pruning free: 'remove a memory' == set weight to 0")


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--synapses", type=int, default=20000)
    p.add_argument("--cycles", type=int, default=8)
    p.add_argument("--k", type=int, default=400, help="signal synapses per task")
    a = p.parse_args()
    run(a.synapses, a.cycles, a.k)


if __name__ == "__main__":
    main()