#!/usr/bin/env python3
"""M3-scout: compute proportional to SURPRISE, not to sequence length.

The core inefficiency of a transformer is that it spends the SAME dense compute
on every token, predictable or not. Real hidden states are highly redundant in
time: token t+1's state is usually close to token t's. Mind-1B's thesis is to
spend compute (and memory, and sleep-consolidation) proportional to prediction
error (surprise). This script measures the payoff of the compute half.

Mechanism (delta-coding, the thing a custom Triton kernel would exploit):
  keep the last input we actually computed on, h_prev, and the cached output.
  surprise per channel  = |h_t - h_prev|
  changed = channels where surprise > theta
  delta out = W[:, changed] @ (h_t - h_prev)[changed]   # only changed inputs
  out += delta out ; h_prev[changed] = h_t[changed]
Cost = d * (#changed) MACs instead of d*d. Predictable channels coast for free.

A dense matvec cannot skip this: cuBLAS has no per-token dynamic sparsity. A
Triton kernel that gathers active channels, does ternary add/sub accumulation,
and scatters results turns "low surprise" directly into "few FLOPs". This is
where writing a kernel actually beats the transformer, and it is an INFERENCE
property (training still does dense fp matmuls for the STE gradient).

Run: python3 surprise_gating.py
"""

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


def make_stream(seq, d, redundancy, rng):
    """A hidden-state stream with tunable temporal redundancy. With prob
    `redundancy` a token nudges only a few channels (predictable); otherwise it
    jumps (a surprising token). Mimics text: mostly-smooth with occasional turns."""
    h = rng.standard_normal(d).astype(np.float32)
    out = np.empty((seq, d), dtype=np.float32)
    for t in range(seq):
        if rng.random() < redundancy:
            k = rng.integers(1, max(2, d // 20))          # a few channels drift
            idx = rng.integers(0, d, size=k)
            h = h.copy()
            h[idx] += rng.standard_normal(k).astype(np.float32) * 0.5
        else:
            h = rng.standard_normal(d).astype(np.float32)  # surprising jump
        out[t] = h
    return out


def event_driven(stream, W, theta):
    seq, d = stream.shape
    h_prev = np.zeros(d, dtype=np.float32)
    out = np.zeros(d, dtype=np.float32)          # cached running output
    macs = 0
    err_num = 0.0
    err_den = 0.0
    for t in range(seq):
        h = stream[t]
        surprise = np.abs(h - h_prev)
        changed = surprise > theta
        nch = int(changed.sum())
        if nch:
            delta = (h - h_prev)[changed]
            out = out + W[:, changed] @ delta   # d * nch MACs
            macs += d * nch
            h_prev = h_prev.copy()
            h_prev[changed] = h[changed]
        dense = W @ h                            # reference (not counted)
        err_num += float(np.sum((out - dense) ** 2))
        err_den += float(np.sum(dense ** 2))
    rel_err = (err_num / err_den) ** 0.5
    dense_macs = seq * d * d
    return macs / dense_macs, rel_err


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--d", type=int, default=512)
    p.add_argument("--seq", type=int, default=2000)
    p.add_argument("--redundancy", type=float, nargs="*",
                   default=[0.85, 0.95])
    a = p.parse_args()

    rng = np.random.default_rng(0)
    W = rng.standard_normal((a.d, a.d)).astype(np.float32) / np.sqrt(a.d)

    print("compute proportional to surprise (delta-coded matvec)")
    print("d=%d, seq=%d. 'MAC ratio' = fraction of a dense matvec's FLOPs used.\n"
          % (a.d, a.seq))
    for red in a.redundancy:
        stream = make_stream(a.seq, a.d, red, np.random.default_rng(1))
        # measure true redundancy: mean fraction of channels changing per step
        step = np.abs(np.diff(stream, axis=0))
        print("temporal redundancy setting %.2f (mean |delta| per channel %.3f)"
              % (red, float(step.mean())))
        print("  %-10s %-14s %-12s" % ("theta", "MAC ratio", "output err"))
        for theta in (0.0, 0.05, 0.1, 0.2, 0.4):
            ratio, err = event_driven(stream, W, theta)
            note = "  <- rounding-level residual" if theta == 0 else ""
            print("  %-10.2f %-13.1f%% %-11.3f%s"
                  % (theta, ratio * 100, err, note))
        print()

    print("reading it: on a redundant stream, a small theta cuts the matvec to")
    print("a fraction of dense FLOPs with tiny output error. That fraction IS")
    print("the transformer-efficiency win, and it is exactly what a dynamic-")
    print("sparsity ternary Triton kernel would turn into wall-clock speedup.")
    print("Same surprise signal also gates recursion depth, fast-weight writes,")
    print("and sleep-replay priority (see INNOVATION.md, the surprise economy).")


if __name__ == "__main__":
    main()