#!/usr/bin/env python3
"""Small, deterministic mechanism probes for model-level memory.

This is not a reproduction of the benchmark scores in any cited paper.  It
isolates three architectural claims from "Memory for Large Language Models":

1. A KV window offers high-fidelity recent recall but has a bounded horizon.
2. A fixed-size associative fast-weight memory trades storage growth for
   interference.  Its update follows the linearized Titans equations.
3. Engram-style deterministic multi-head hashing trades table capacity for
   collisions, while addressing cost does not require scanning the table.

Only NumPy and Matplotlib are required.  All random generators are seeded.
"""

from __future__ import annotations

import csv
import json
import math
import platform
import statistics
import time
from dataclasses import dataclass
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib import font_manager
import numpy as np


font_manager.fontManager.addfont("/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc")
plt.rcParams["font.family"] = "Noto Sans CJK JP"
plt.rcParams["axes.unicode_minus"] = False


ROOT = Path(__file__).resolve().parents[1]
DATA_DIR = ROOT / "data"
ASSET_DIR = ROOT / "assets"

SEEDS = tuple(range(12))
KEY_DIM = 48
VALUE_DIM = 48
RECENT_WINDOW = 96
LOADS = (48, 96, 192, 384, 768)


@dataclass(frozen=True)
class FastWeightRule:
    name: str
    label: str
    theta: float
    eta: float
    alpha: float


RULES = (
    FastWeightRule("delta", "Fast weights · delta", theta=0.34, eta=0.0, alpha=0.0),
    FastWeightRule(
        "momentum_decay",
        "Fast weights · momentum + decay",
        theta=0.17,
        eta=0.58,
        alpha=0.003,
    ),
)


def unit_rows(x: np.ndarray) -> np.ndarray:
    return x / np.maximum(np.linalg.norm(x, axis=1, keepdims=True), 1e-12)


def update_fast_weights(
    keys: np.ndarray,
    values: np.ndarray,
    rule: FastWeightRule,
) -> np.ndarray:
    """Linear associative memory using the Titans inner-loop update.

    Loss: 1/2 ||W k_t - v_t||^2
    S_t = eta S_(t-1) - theta grad(loss)
    W_t = (1-alpha) W_(t-1) + S_t
    """

    w = np.zeros((values.shape[1], keys.shape[1]), dtype=np.float64)
    surprise_state = np.zeros_like(w)
    for key, value in zip(keys, values, strict=True):
        prediction = w @ key
        grad = np.outer(prediction - value, key)
        surprise_state = rule.eta * surprise_state - rule.theta * grad
        w = (1.0 - rule.alpha) * w + surprise_state
    return w


def cosine_retrieval_accuracy(
    query_keys: np.ndarray,
    target_values: np.ndarray,
    memory_keys: np.ndarray,
    memory_values: np.ndarray,
) -> float:
    scores = query_keys @ memory_keys.T
    predictions = memory_values[np.argmax(scores, axis=1)]
    cos = np.sum(unit_rows(predictions) * unit_rows(target_values), axis=1)
    return float(np.mean(cos > 0.82))


def fast_weight_accuracy(
    query_keys: np.ndarray,
    target_values: np.ndarray,
    weights: np.ndarray,
) -> float:
    predictions = query_keys @ weights.T
    cos = np.sum(unit_rows(predictions) * unit_rows(target_values), axis=1)
    return float(np.mean(cos > 0.82))


def probe_retention_by_load() -> list[dict[str, float | int | str]]:
    rows: list[dict[str, float | int | str]] = []
    for load in LOADS:
        for workload in ("episodic", "structured"):
            per_method: dict[str, list[float]] = {
                "recent_kv": [],
                "explicit_kv": [],
                "delta": [],
                "momentum_decay": [],
                "hybrid": [],
            }
            for seed in SEEDS:
                rng = np.random.default_rng(
                    10_000 + seed * 97 + load + (50_000 if workload == "structured" else 0)
                )
                if workload == "episodic":
                    keys = unit_rows(rng.normal(size=(load, KEY_DIM)))
                    values = unit_rows(rng.normal(size=(load, VALUE_DIM)))
                    noisy_queries = unit_rows(keys + 0.035 * rng.normal(size=keys.shape))
                else:
                    topic_count = 12
                    key_topics = unit_rows(rng.normal(size=(topic_count, KEY_DIM)))
                    value_topics = unit_rows(rng.normal(size=(topic_count, VALUE_DIM)))
                    topic_ids = rng.integers(0, topic_count, size=load)
                    keys = unit_rows(
                        key_topics[topic_ids] + 0.13 * rng.normal(size=(load, KEY_DIM))
                    )
                    values = unit_rows(
                        value_topics[topic_ids] + 0.018 * rng.normal(size=(load, VALUE_DIM))
                    )
                    noisy_queries = unit_rows(
                        key_topics[topic_ids] + 0.13 * rng.normal(size=(load, KEY_DIM))
                    )

                recent_start = max(0, load - RECENT_WINDOW)
                recent_keys = keys[recent_start:]
                recent_values = values[recent_start:]
                per_method["recent_kv"].append(
                    cosine_retrieval_accuracy(
                        noisy_queries, values, recent_keys, recent_values
                    )
                )
                per_method["explicit_kv"].append(
                    cosine_retrieval_accuracy(noisy_queries, values, keys, values)
                )

                fast_weights: dict[str, np.ndarray] = {}
                for rule in RULES:
                    weights = update_fast_weights(keys, values, rule)
                    fast_weights[rule.name] = weights
                    per_method[rule.name].append(
                        fast_weight_accuracy(noisy_queries, values, weights)
                    )

                recent_scores = noisy_queries @ recent_keys.T
                recent_best = np.max(recent_scores, axis=1)
                recent_predictions = recent_values[np.argmax(recent_scores, axis=1)]
                long_predictions = noisy_queries @ fast_weights["momentum_decay"].T
                hybrid_predictions = np.where(
                    (recent_best > 0.84)[:, None], recent_predictions, long_predictions
                )
                hybrid_cos = np.sum(
                    unit_rows(hybrid_predictions) * unit_rows(values), axis=1
                )
                per_method["hybrid"].append(float(np.mean(hybrid_cos > 0.82)))

            labels = {
                "recent_kv": "Implicit recent KV (96 slots)",
                "explicit_kv": "Explicit addressable KV (all items)",
                "delta": RULES[0].label,
                "momentum_decay": RULES[1].label,
                "hybrid": "Hybrid recent KV + fast weights",
            }
            for method, samples in per_method.items():
                rows.append(
                    {
                        "load": load,
                        "workload": workload,
                        "method": method,
                        "label": labels[method],
                        "accuracy_mean": statistics.mean(samples),
                        "accuracy_std": statistics.pstdev(samples),
                        "seeds": len(samples),
                    }
                )
    return rows


def probe_delay_profile(load: int = 768) -> list[dict[str, float | int | str]]:
    bins = (
        (0, 48, "0–47"),
        (48, 96, "48–95"),
        (96, 192, "96–191"),
        (192, 384, "192–383"),
        (384, load, "384–767"),
    )
    collected: dict[tuple[str, str], list[float]] = {}
    for seed in SEEDS:
        rng = np.random.default_rng(20_000 + seed * 101)
        keys = unit_rows(rng.normal(size=(load, KEY_DIM)))
        values = unit_rows(rng.normal(size=(load, VALUE_DIM)))
        queries = unit_rows(keys + 0.035 * rng.normal(size=keys.shape))
        recent_keys = keys[-RECENT_WINDOW:]
        recent_values = values[-RECENT_WINDOW:]
        w = update_fast_weights(keys, values, RULES[1])
        explicit_predictions = values[np.argmax(queries @ keys.T, axis=1)]
        recent_predictions = recent_values[np.argmax(queries @ recent_keys.T, axis=1)]
        fast_predictions = queries @ w.T
        recent_best = np.max(queries @ recent_keys.T, axis=1)
        hybrid_predictions = np.where(
            (recent_best > 0.84)[:, None], recent_predictions, fast_predictions
        )

        prediction_map = {
            "recent_kv": recent_predictions,
            "fast_weight": fast_predictions,
            "hybrid": hybrid_predictions,
            "explicit_kv": explicit_predictions,
        }
        ages = np.arange(load - 1, -1, -1)
        for low, high, bin_label in bins:
            mask = (ages >= low) & (ages < high)
            for method, predictions in prediction_map.items():
                cos = np.sum(unit_rows(predictions[mask]) * values[mask], axis=1)
                collected.setdefault((bin_label, method), []).append(
                    float(np.mean(cos > 0.82))
                )

    order = {label: i for i, (_, _, label) in enumerate(bins)}
    rows: list[dict[str, float | int | str]] = []
    for (bin_label, method), samples in sorted(
        collected.items(), key=lambda item: (order[item[0][0]], item[0][1])
    ):
        rows.append(
            {
                "load": load,
                "age_bin": bin_label,
                "method": method,
                "accuracy_mean": statistics.mean(samples),
                "accuracy_std": statistics.pstdev(samples),
                "seeds": len(samples),
            }
        )
    return rows


def probe_stability_plasticity() -> list[dict[str, float | int | str]]:
    """Remap half the associations and measure old/new recall during updates."""

    rule_variants = (
        FastWeightRule("delta", "Delta", theta=0.30, eta=0.0, alpha=0.0),
        FastWeightRule("momentum", "Momentum", theta=0.14, eta=0.55, alpha=0.0),
        FastWeightRule(
            "momentum_decay", "Momentum + decay", theta=0.14, eta=0.55, alpha=0.004
        ),
    )
    steps = (0, 1, 2, 4, 8, 12, 16, 24)
    accumulator: dict[tuple[str, int, str], list[float]] = {}

    for seed in SEEDS:
        rng = np.random.default_rng(30_000 + seed * 103)
        count = 36
        keys = unit_rows(rng.normal(size=(count, KEY_DIM)))
        old_values = unit_rows(rng.normal(size=(count, VALUE_DIM)))
        new_values = old_values.copy()
        remapped = np.arange(count // 2)
        new_values[remapped] = unit_rows(rng.normal(size=(len(remapped), VALUE_DIM)))
        stable = np.arange(count // 2, count)

        base_sequence = np.tile(np.arange(count), 5)
        rng.shuffle(base_sequence)

        for rule in rule_variants:
            w = np.zeros((VALUE_DIM, KEY_DIM), dtype=np.float64)
            s = np.zeros_like(w)
            for idx in base_sequence:
                grad = np.outer(w @ keys[idx] - old_values[idx], keys[idx])
                s = rule.eta * s - rule.theta * grad
                w = (1.0 - rule.alpha) * w + s

            def record(step: int) -> None:
                predictions = unit_rows(keys @ w.T)
                metrics = {
                    "stable_old": float(
                        np.mean(np.sum(predictions[stable] * old_values[stable], axis=1) > 0.82)
                    ),
                    "remapped_new": float(
                        np.mean(
                            np.sum(predictions[remapped] * new_values[remapped], axis=1)
                            > 0.82
                        )
                    ),
                    "remapped_old": float(
                        np.mean(
                            np.sum(predictions[remapped] * old_values[remapped], axis=1)
                            > 0.82
                        )
                    ),
                }
                for metric, value in metrics.items():
                    accumulator.setdefault((rule.name, step, metric), []).append(value)

            record(0)
            completed = 0
            for target_step in steps[1:]:
                for _ in range(target_step - completed):
                    order = remapped.copy()
                    rng.shuffle(order)
                    for idx in order:
                        grad = np.outer(w @ keys[idx] - new_values[idx], keys[idx])
                        s = rule.eta * s - rule.theta * grad
                        w = (1.0 - rule.alpha) * w + s
                completed = target_step
                record(target_step)

    rows: list[dict[str, float | int | str]] = []
    for (rule, step, metric), samples in sorted(accumulator.items()):
        rows.append(
            {
                "rule": rule,
                "update_passes": step,
                "metric": metric,
                "accuracy_mean": statistics.mean(samples),
                "accuracy_std": statistics.pstdev(samples),
                "seeds": len(samples),
            }
        )
    return rows


def next_prime(n: int) -> int:
    def is_prime(value: int) -> bool:
        if value < 2:
            return False
        if value % 2 == 0:
            return value == 2
        limit = int(math.sqrt(value)) + 1
        return all(value % divisor for divisor in range(3, limit, 2))

    candidate = max(2, n)
    while not is_prime(candidate):
        candidate += 1
    return candidate


def engram_signatures(
    ngrams: np.ndarray, table_size: int, heads: int, seed: int = 2026
) -> np.ndarray:
    """Replicate the deterministic core of Engram's official demo mapping."""

    rng = np.random.default_rng(seed)
    multipliers = rng.integers(1, np.iinfo(np.int64).max // 65_536, size=ngrams.shape[1])
    multipliers = multipliers * 2 + 1
    mixed = ngrams[:, 0] * multipliers[0]
    for column in range(1, ngrams.shape[1]):
        mixed = np.bitwise_xor(mixed, ngrams[:, column] * multipliers[column])
    moduli = []
    cursor = table_size
    for _ in range(heads):
        cursor = next_prime(cursor)
        moduli.append(cursor)
        cursor += 1
    return np.stack([mixed % modulus for modulus in moduli], axis=1)


def signature_collision_fraction(signatures: np.ndarray) -> float:
    contiguous = np.ascontiguousarray(signatures)
    packed = contiguous.view(
        np.dtype((np.void, contiguous.dtype.itemsize * contiguous.shape[1]))
    ).ravel()
    unique = np.unique(packed).size
    return float(1.0 - unique / len(signatures))


def probe_engram_hashing() -> list[dict[str, float | int]]:
    rng = np.random.default_rng(40_026)
    sample_count = 60_000
    vocab_size = 65_536
    ngrams = rng.integers(0, vocab_size, size=(sample_count, 3), dtype=np.int64)
    rows: list[dict[str, float | int]] = []
    for table_size in (257, 1021, 4093, 16_381, 65_521):
        for heads in (1, 2, 3):
            start = time.perf_counter_ns()
            signatures = engram_signatures(ngrams, table_size, heads)
            elapsed = time.perf_counter_ns() - start
            rows.append(
                {
                    "table_size_requested": table_size,
                    "heads": heads,
                    "samples": sample_count,
                    "collision_fraction": signature_collision_fraction(signatures),
                    "address_ns_per_item": elapsed / sample_count,
                }
            )
    return rows


def write_csv(filename: str, rows: list[dict]) -> None:
    path = DATA_DIR / filename
    if not rows:
        raise ValueError(f"no rows for {filename}")
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)


def style_axes(ax: plt.Axes) -> None:
    ax.set_facecolor("#f8f6f0")
    ax.grid(axis="y", color="#d8d2c8", linewidth=0.7, alpha=0.8)
    ax.spines[["top", "right"]].set_visible(False)
    ax.spines[["left", "bottom"]].set_color("#8d8a83")
    ax.tick_params(colors="#596071", labelsize=9)


def plot_retention(rows: list[dict]) -> None:
    colors = {
        "recent_kv": "#c35f3f",
        "explicit_kv": "#2f6f72",
        "delta": "#8d7658",
        "momentum_decay": "#5c688a",
        "hybrid": "#30384c",
    }
    labels = {
        "recent_kv": "近期 KV · 96 项",
        "explicit_kv": "可寻址 KV · 全量",
        "delta": "Fast weights · Delta",
        "momentum_decay": "Fast weights · 动量+遗忘",
        "hybrid": "混合：近期 KV + Fast weights",
    }
    fig, axes = plt.subplots(1, 2, figsize=(10.6, 4.9), facecolor="#f8f6f0", sharey=True)
    workloads = (
        ("episodic", "一次性事实：不可压缩的随机关联"),
        ("structured", "重复结构：12 个可复用主题"),
    )
    for ax, (workload, title) in zip(axes, workloads, strict=True):
        for method in colors:
            subset = [
                row
                for row in rows
                if row["method"] == method and row["workload"] == workload
            ]
            ax.plot(
                [row["load"] for row in subset],
                [row["accuracy_mean"] for row in subset],
                marker="o",
                linewidth=2.0,
                color=colors[method],
                label=labels[method],
            )
        ax.axvline(
            RECENT_WINDOW,
            color="#c35f3f",
            linestyle="--",
            linewidth=1,
            alpha=0.55,
        )
        ax.set_xscale("log", base=2)
        ax.set_xticks(LOADS, labels=[str(v) for v in LOADS])
        ax.set_xlabel("写入数量（log2）", color="#596071")
        ax.set_title(title, loc="left", fontsize=11.5, color="#293348", pad=13)
        style_axes(ax)
    axes[0].set_ylim(-0.02, 1.04)
    axes[0].set_ylabel("召回准确率", color="#596071")
    axes[1].legend(frameon=False, fontsize=7.5, loc="lower right")
    fig.suptitle(
        "压缩记忆能否成立，取决于信息有没有可复用结构",
        x=0.055,
        y=0.985,
        ha="left",
        fontsize=15,
        color="#293348",
    )
    fig.tight_layout(rect=(0, 0, 1, 0.93))
    fig.savefig(ASSET_DIR / "experiment-retention-load.svg", format="svg")
    plt.close(fig)


def plot_delay(rows: list[dict]) -> None:
    methods = ("recent_kv", "fast_weight", "hybrid", "explicit_kv")
    labels = {
        "recent_kv": "近期 KV",
        "fast_weight": "Fast weights",
        "hybrid": "混合",
        "explicit_kv": "可寻址 KV",
    }
    colors = {
        "recent_kv": "#c35f3f",
        "fast_weight": "#5c688a",
        "hybrid": "#30384c",
        "explicit_kv": "#2f6f72",
    }
    bins = ("0–47", "48–95", "96–191", "192–383", "384–767")
    fig, ax = plt.subplots(figsize=(10.6, 5.7), facecolor="#f8f6f0")
    for method in methods:
        subset = {row["age_bin"]: row for row in rows if row["method"] == method}
        ax.plot(
            range(len(bins)),
            [subset[label]["accuracy_mean"] for label in bins],
            marker="o",
            linewidth=2.2,
            color=colors[method],
            label=labels[method],
        )
    ax.set_xticks(range(len(bins)), labels=bins)
    ax.set_ylim(-0.02, 1.04)
    ax.set_xlabel("查询目标距流末尾的年龄（项）", color="#596071")
    ax.set_ylabel("召回准确率", color="#596071")
    ax.set_title("近期准确与远期保留不是同一能力", loc="left", fontsize=15, color="#293348", pad=18)
    ax.legend(frameon=False, fontsize=9, ncol=4, loc="upper right")
    style_axes(ax)
    fig.tight_layout()
    fig.savefig(ASSET_DIR / "experiment-delay-profile.svg", format="svg")
    plt.close(fig)


def plot_stability(rows: list[dict]) -> None:
    rules = ("delta", "momentum", "momentum_decay")
    labels = {
        "delta": "Delta",
        "momentum": "动量",
        "momentum_decay": "动量 + 遗忘",
    }
    colors = {
        "delta": "#8d7658",
        "momentum": "#5c688a",
        "momentum_decay": "#c35f3f",
    }
    fig, axes = plt.subplots(1, 2, figsize=(10.6, 4.9), facecolor="#f8f6f0", sharey=True)
    panels = (
        ("remapped_new", "塑性：新映射学得多快"),
        ("stable_old", "稳定性：未改知识留得多久"),
    )
    for ax, (metric, title) in zip(axes, panels, strict=True):
        for rule in rules:
            subset = [
                row for row in rows if row["rule"] == rule and row["metric"] == metric
            ]
            subset.sort(key=lambda row: row["update_passes"])
            ax.plot(
                [row["update_passes"] for row in subset],
                [row["accuracy_mean"] for row in subset],
                marker="o",
                linewidth=2.1,
                color=colors[rule],
                label=labels[rule],
            )
        ax.set_title(title, loc="left", fontsize=12, color="#293348", pad=13)
        ax.set_xlabel("新映射更新轮数", color="#596071")
        style_axes(ax)
    axes[0].set_ylabel("召回准确率", color="#596071")
    axes[0].set_ylim(-0.02, 1.04)
    axes[1].legend(frameon=False, fontsize=8.5, loc="lower right")
    fig.suptitle(
        "同一条在线更新规则同时决定学会与忘掉",
        x=0.055,
        y=0.985,
        ha="left",
        fontsize=15,
        color="#293348",
    )
    fig.tight_layout(rect=(0, 0, 1, 0.93))
    fig.savefig(ASSET_DIR / "experiment-stability-plasticity.svg", format="svg")
    plt.close(fig)


def plot_engram(rows: list[dict]) -> None:
    colors = {1: "#c35f3f", 2: "#5c688a", 3: "#2f6f72"}
    fig, ax = plt.subplots(figsize=(10.6, 5.4), facecolor="#f8f6f0")
    for heads in (1, 2, 3):
        subset = [row for row in rows if row["heads"] == heads]
        ax.plot(
            [row["table_size_requested"] for row in subset],
            [max(row["collision_fraction"], 1e-6) for row in subset],
            marker="o",
            linewidth=2.2,
            color=colors[heads],
            label=f"{heads} 个哈希头",
        )
    ax.set_xscale("log", base=4)
    ax.set_yscale("log")
    ax.set_xticks(
        [257, 1021, 4093, 16_381, 65_521],
        labels=["257", "1K", "4K", "16K", "65K"],
    )
    ax.set_xlabel("每个哈希头的表规模", color="#596071")
    ax.set_ylabel("完整签名碰撞比例（log）", color="#596071")
    ax.set_title("Engram 式多头质数哈希：容量换碰撞，而不是扫描全表", loc="left", fontsize=15, color="#293348", pad=18)
    ax.legend(frameon=False, fontsize=9)
    style_axes(ax)
    fig.tight_layout()
    fig.savefig(ASSET_DIR / "experiment-engram-collisions.svg", format="svg")
    plt.close(fig)


def main() -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    ASSET_DIR.mkdir(parents=True, exist_ok=True)

    retention = probe_retention_by_load()
    delay = probe_delay_profile()
    stability = probe_stability_plasticity()
    engram = probe_engram_hashing()

    write_csv("retention_by_load.csv", retention)
    write_csv("delay_profile.csv", delay)
    write_csv("stability_plasticity.csv", stability)
    write_csv("engram_hashing.csv", engram)

    plot_retention(retention)
    plot_delay(delay)
    plot_stability(stability)
    plot_engram(engram)

    metadata = {
        "purpose": "mechanism probes, not paper benchmark reproduction",
        "seed_count": len(SEEDS),
        "seeds": list(SEEDS),
        "seed_policy": {
            "retention_delay_stability": list(SEEDS),
            "engram_ngram_input": 40_026,
            "engram_hash_multipliers": 2_026,
        },
        "key_dim": KEY_DIM,
        "value_dim": VALUE_DIM,
        "recent_window": RECENT_WINDOW,
        "loads": list(LOADS),
        "python": platform.python_version(),
        "numpy": np.__version__,
        "matplotlib": matplotlib.__version__,
        "generated_at": "2026-08-07T00:00:00+08:00",
        "sources": {
            "survey": "https://arxiv.org/abs/2607.25380",
            "titans": "https://arxiv.org/abs/2501.00663",
            "engram_code": "https://github.com/deepseek-ai/Engram",
        },
    }
    (DATA_DIR / "run_metadata.json").write_text(
        json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )

    print(json.dumps(metadata, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()