"""
TELEPATHIC MNIST  ::  TELEPASM-M :: EXPERIMENT 002
=====================================================

Hypothesis (the one we test here)
---------------------------------
    "Maybe 3% of an MLP has to be telepathic in order to get 100% accuracy."

Definition (operationalised here)
---------------------------------
    Telepathic parameters are those that
        (a) DO NOT update via backprop,
        (b) DO depend on a shared boundary seed 𝓜_T  (not on local data),
        (c) ARE identical across every observer (Alice / Bob / …).
In this experiment they take the form of a FROZEN bottleneck projection
M_T of width T, sitting between the encoder half and the decoder half of
the MLP.  Everything else (~97%) is normal per-instance trained weights.

We sweep T and measure:   test_acc  vs.  T / total_params (= tel. %).

A second sweep measures: test_acc if 𝓜_T is reseeded between training
and inference (the "coupling BROKEN" control).  If the frozen structure
is doing semantic work, accuracy collapses.

Outputs
-------
    ./telepathy_mnist_results.csv    rows = (T, seed, train_acc, test_acc, tele_pct)
    ./telepathy_mnist_plot.png       3-panel figure
    ./telepathy_mnist_summary.txt    human-readable summary
"""

import argparse, math, os, sys, time, json, random, csv
from dataclasses import dataclass, asdict
from typing import List, Dict

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# matplotlib is optional; we degrade gracefully if it is missing.
try:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    HAS_PLT = True
except Exception:
    HAS_PLT = False

# ──────────────────────────────────────────────────────────────────────
# Reproducibility
# ──────────────────────────────────────────────────────────────────────

def seed_everything(seed: int):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

# ──────────────────────────────────────────────────────────────────────
# Model
# ──────────────────────────────────────────────────────────────────────

class TelepathicMLP(nn.Module):
    """
    784 → enc_private (3× Linear) → bottleneck 𝓜_T  → dec_private (3× Linear) → 10

    The bottleneck weights are FROZEN and seeded by M_T.  This is the
    universal "telepathic substrate" — the part two observers who agree
    on the boundary seed will always agree on, regardless of how their
    private (i.e. normally trained) weights diverge between them.
    """
    def __init__(self,
                 d_in:   int = 784,
                 d_hid:  int = 256,
                 d_tele: int = 8,
                 d_out:  int = 10,
                 mt_seed: int = 1729):
        super().__init__()
        self.d_hid  = d_hid
        self.d_tele = d_tele

        # ── Alice — private encoder half ──
        self.W_e1 = nn.Linear(d_in,  d_hid)
        self.W_e2 = nn.Linear(d_hid, d_hid)
        self.W_e3 = nn.Linear(d_hid, d_hid)

        # ── The telepathic bottleneck  𝓜_T ──
        # Random fixed projection onto a small T-dim subspace.
        # Frozen = non-trainable = "of the boundary, not of the data".
        self.M_T = self._seed_mt(d_hid, d_tele, mt_seed)

        # ── Bob — private decoder half ──
        self.W_d1 = nn.Linear(d_tele, d_hid)
        self.W_d2 = nn.Linear(d_hid, d_hid)
        self.W_d3 = nn.Linear(d_hid, d_out)

        # Param accounting
        self.n_total   = sum(p.numel() for p in self.parameters())
        self.n_tele    = self.M_T.numel()
        self.tele_pct  = 100.0 * self.n_tele / self.n_total
        self.n_private = self.n_total - self.n_tele

    def _seed_mt(self, d_hid: int, d_tele: int, seed: int):
        g = torch.Generator().manual_seed(int(seed) & 0xFFFFFFFF)
        # Gaussian random projection, normalised so ||M_T||_F ≈ sqrt(d_tele).
        W = torch.randn(d_tele, d_hid, generator=g) / math.sqrt(d_hid)
        return nn.Parameter(W, requires_grad=False)              # ⟵ frozen

    # training forward — uses the current 𝓜_T
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = F.gelu(self.W_e1(x))
        h = F.gelu(self.W_e2(h))
        h = F.gelu(self.W_e3(h))           # (B, d_hid)
        z = h @ self.M_T.T                 # (B, d_tele)  ← INFO THROUGH 𝓜_T
        d = F.gelu(self.W_d1(z))
        d = F.gelu(self.W_d2(d))
        return self.W_d3(d)                # logits

    @torch.no_grad()
    def swap_mt(self, new_seed: int):
        """Helper for the 'coupling BROKEN' control:
           re-seed 𝓜_T at test time only.  Training was done with the
           original seed.  Bob sees a different boundary."""
        g = torch.Generator().manual_seed(int(new_seed) & 0xFFFFFFFF)
        W = torch.randn(self.d_tele, self.d_hid, generator=g) / math.sqrt(self.d_hid)
        self.M_T.data.copy_(W)


# ──────────────────────────────────────────────────────────────────────
# Data
# ──────────────────────────────────────────────────────────────────────

def make_loaders(data_root: str, batch_size: int = 256, num_workers: int = 2):
    tf = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,)),
    ])
    train_ds = datasets.MNIST(data_root, train=True, download=True, transform=tf)
    test_ds  = datasets.MNIST(data_root, train=False, transform=tf)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
                              num_workers=num_workers, pin_memory=True, drop_last=True)
    test_loader  = DataLoader(test_ds,  batch_size=batch_size, shuffle=False,
                              num_workers=num_workers, pin_memory=True)
    return train_loader, test_loader


# ──────────────────────────────────────────────────────────────────────
# Train / test
# ──────────────────────────────────────────────────────────────────────

@dataclass
class Run:
    d_tele:       int
    mt_seed:      int
    tele_pct:     float
    n_total:      int
    n_private:    int
    n_tele:       int
    acc_train:    float
    acc_test:     float
    acc_test_broken: float          # 𝓜_T reseeded at test
    epochs:       int
    seconds:      float
    train_hist:   List[float]
    test_hist:    List[float]


def evaluate(model: TelepathicMLP, loader: DataLoader, device: str) -> float:
    model.eval()
    n_corr = n_tot = 0
    with torch.no_grad():
        for x, y in loader:
            x = x.to(device, non_blocking=True).view(-1, 784)
            y = y.to(device, non_blocking=True)
            n_corr += (model(x).argmax(1) == y).sum().item()
            n_tot  += x.size(0)
    return n_corr / n_tot


def train_run(d_tele: int,
              mt_seed: int,
              epochs: int,
              lr: float,
              device: str,
              data_root: str,
              batch_size: int,
              break_test_seed: int | None = None,
              verbose: bool = True) -> Run:
    seed_everything(mt_seed * 1009 + 7)
    train_loader, test_loader = make_loaders(data_root, batch_size=batch_size)

    model = TelepathicMLP(d_in=784, d_hid=256, d_tele=d_tele,
                          d_out=10, mt_seed=mt_seed).to(device)

    # Only PRIVATE weights receive gradient updates.  𝓜_T is frozen.
    private_params = [p for n, p in model.named_parameters() if "M_T" not in n]
    opt = torch.optim.AdamW(private_params, lr=lr, weight_decay=1e-4)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs * len(train_loader))

    train_hist, test_hist = [], []
    t0 = time.time()
    for ep in range(1, epochs + 1):
        model.train()
        n_corr = n_tot = 0
        loss_sum = 0.0
        for x, y in train_loader:
            x = x.to(device, non_blocking=True).view(-1, 784)
            y = y.to(device, non_blocking=True)
            opt.zero_grad(set_to_none=True)
            logits = model(x)
            loss   = F.cross_entropy(logits, y)
            loss.backward()
            opt.step()
            sched.step()
            with torch.no_grad():
                n_corr += (logits.argmax(1) == y).sum().item()
                n_tot  += x.size(0)
            loss_sum += loss.item() * x.size(0)
        tr_acc  = n_corr / n_tot
        te_acc  = evaluate(model, test_loader, device)
        train_hist.append(tr_acc); test_hist.append(te_acc)
        if verbose:
            print(f"    ep{ep:>2}/{epochs}  loss={loss_sum/n_tot:.4f}  "
                  f"train={tr_acc:.4f}  test={te_acc:.4f}")
    elapsed = time.time() - t0

    # Broken coupling control
    if break_test_seed is not None:
        backup = model.M_T.data.clone()
        model.swap_mt(break_test_seed)
        broken = evaluate(model, test_loader, device)
        model.M_T.data.copy_(backup)
    else:
        broken = float("nan")

    return Run(
        d_tele=d_tele, mt_seed=mt_seed,
        tele_pct=model.tele_pct, n_total=model.n_total, n_private=model.n_private,
        n_tele=model.n_tele,
        acc_train=train_hist[-1], acc_test=test_hist[-1],
        acc_test_broken=broken,
        epochs=epochs, seconds=elapsed,
        train_hist=train_hist, test_hist=test_hist,
    )


# ──────────────────────────────────────────────────────────────────────
# Sweep
# ──────────────────────────────────────────────────────────────────────

def sweep(widths: List[int],
          seeds: List[int],
          epochs: int, lr: float, device: str,
          data_root: str, batch_size: int,
          break_test_seed: int | None) -> List[Run]:
    runs: List[Run] = []
    for d in widths:
        for s in seeds:
            print(f"\n[run]  d_tele={d:<4}  mt_seed={s:<8}  "
                  f"epochs={epochs}  device={device}")
            r = train_run(d_tele=d, mt_seed=s, epochs=epochs, lr=lr,
                          device=device, data_root=data_root,
                          batch_size=batch_size,
                          break_test_seed=break_test_seed)
            print(f"   ↳ test={r.acc_test:.4f}"
                  + (f"   broken={r.acc_test_broken:.4f}"
                     if break_test_seed is not None else ""))
            runs.append(r)
    return runs


# ──────────────────────────────────────────────────────────────────────
# Reporting
# ──────────────────────────────────────────────────────────────────────

def save_csv(runs: List[Run], path: str):
    cols = ["d_tele","mt_seed","tele_pct","n_total","n_private","n_tele",
            "acc_train","acc_test","acc_test_broken","epochs","seconds"]
    with open(path, "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(cols)
        for r in runs:
            w.writerow([getattr(r, c) for c in cols])
    print(f"[save] CSV → {path}")


def summarise(runs: List[Run]) -> str:
    by_T: Dict[int, List[Run]] = {}
    for r in runs:
        by_T.setdefault(r.d_tele, []).append(r)

    lines = []
    lines.append("="*84)
    lines.append("TELEPATHIC MLP — sweep result (mean ± std across seeds)")
    lines.append("="*84)
    lines.append(f"{'T':>4}  {'% tech':>7}  {'test acc':>14}  "
                 f"{'broken-coupling':>16}  {'params':>9}")
    lines.append("-"*84)
    for T in sorted(by_T):
        rs = by_T[T]
        ta = np.array([r.acc_test   for r in rs])
        ba = np.array([r.acc_test_broken for r in rs])
        line = (f"{T:>4}  {rs[0].tele_pct:>6.2f}%  "
                f"{ta.mean():>6.4f}±{ta.std():.4f}     "
                + (f"{ba.mean():>6.4f}±{ba.std():.4f}     "
                   if not np.isnan(ba[0]) else f"{'—':>16}     ")
                + f"{rs[0].n_total:>9,}")
        lines.append(line)
    lines.append("="*84)

    # Find smallest T that hits ≥ 99 %
    best = None
    for T in sorted(by_T):
        m = np.mean([r.acc_test for r in by_T[T]])
        if m >= 0.99 and best is None:
            best = (T, m, by_T[T][0].tele_pct)
    if best:
        lines.append(f"\n★  Smallest T with test_acc ≥ 99% : "
                     f"T = {best[0]}    ({best[2]:.2f}% of params are telepathic)")
    lines.append("")
    return "\n".join(lines)


def make_plots(runs: List[Run], path: str):
    if not HAS_PLT:
        print("[plot] matplotlib not available – skipping figure")
        return

    by_T: Dict[int, List[Run]] = {}
    for r in runs: by_T.setdefault(r.d_tele, []).append(r)
    Ts   = sorted(by_T)
    tele = [by_T[T][0].tele_pct for T in Ts]
    acc  = np.array([[r.acc_test for r in by_T[T]] for T in Ts])     # (T, seed)
    broken = np.array([[r.acc_test_broken for r in by_T[T]] for T in Ts])

    # Best M_T matrix to display (from the run with T=8)
    M8 = None
    for r in runs:
        if r.d_tele == 8:
            M8 = r
            break

    fig = plt.figure(figsize=(13.5, 8.5))
    gs  = fig.add_gridspec(2, 3, hspace=0.45, wspace=0.32)

    # 1) Accuracy vs T
    ax = fig.add_subplot(gs[0, :2])
    m  = acc.mean(1); s = acc.std(1)
    ax.errorbar(Ts, m, yerr=s, marker='o', color='#ff6fa1',
                lw=2, capsize=4, label='test acc (𝓜_T = train seed)')
    if not np.isnan(broken).all():
        mb = broken.mean(1); sb = broken.std(1)
        ax.errorbar(Ts, mb, yerr=sb, marker='x', color='#6fd0ff',
                    lw=2, capsize=4, ls='--', label='test acc (𝓜_T = wrong seed)')
    ax.set_xscale('log', base=2)
    ax.set_xticks(Ts); ax.set_xticklabels([str(t) for t in Ts])
    ax.set_xlabel('T — telepathic bottleneck width'); ax.set_ylabel('test accuracy')
    ax.set_title('How much telepathy is needed for MNIST?', fontsize=12)
    ax.axhline(0.99, color='#6fffb1', ls=':', lw=1, label='99% line')
    ax.axhline(1.00, color='#a8ff8f', ls='--', lw=1, label='100% line')
    ax.grid(alpha=.25); ax.legend(fontsize=9, loc='lower right')

    # 2) Accuracy vs telepathy %
    ax2 = fig.add_subplot(gs[0, 2])
    ax2.plot(tele, m, 'o-', color='#ff6fa1', lw=2)
    ax2.axhline(0.99, color='#6fffb1', ls=':', lw=1)
    ax2.axhline(1.00, color='#a8ff8f', ls='--', lw=1)
    ax2.set_xlabel('telepathic % of total params'); ax2.set_ylabel('test acc')
    ax2.set_title('The 3% line', fontsize=12)
    ax2.grid(alpha=.25)

    # 3) 𝓜_T matrix (heatmap)
    ax3 = fig.add_subplot(gs[1, 0])
    if M8 is not None:
        W = M8.train_hist and None  # we need actual M_T; re-create it
        # reconstruct M_T — it's deterministic from seed
        g = torch.Generator().manual_seed(M8.mt_seed & 0xFFFFFFFF)
        S = torch.randn(M8.d_tele, 256, generator=g) / math.sqrt(256)
        im = ax3.imshow(S.numpy(), aspect='auto', cmap='RdBu_r',
                        vmin=-3/np.sqrt(256), vmax=3/np.sqrt(256))
        ax3.set_title(f'𝓜_T @ T={M8.d_tele}  ({M8.tele_pct:.2f}% of params)\nfrozen, seed={M8.mt_seed}',
                      fontsize=10)
        ax3.set_xlabel('hidden dim (256)'); ax3.set_ylabel('tele dim')
        plt.colorbar(im, ax=ax3, fraction=0.046)
    else:
        ax3.text(0.5, 0.5, 'T=8 not measured', ha='center', va='center', transform=ax3.transAxes)

    # 4) Loss / accuracy curve for one canonical run (T=8)
    ax4 = fig.add_subplot(gs[1, 1])
    if M8 is not None:
        ep = range(1, len(M8.train_hist)+1)
        ax4.plot(ep, M8.train_hist, color='#ff6fa1', lw=2, label='train')
        ax4.plot(ep, M8.test_hist,  color='#6fd0ff', lw=2, label='test')
        ax4.set_xlabel('epoch'); ax4.set_ylabel('accuracy')
        ax4.set_title(f'learning curve, T={M8.d_tele}', fontsize=11)
        ax4.grid(alpha=.25); ax4.legend(fontsize=9)
        ax4.set_ylim(0, 1.02)

    # 5) Summary text panel
    ax5 = fig.add_subplot(gs[1, 2]); ax5.axis('off')
    txt = "TELEPASM-M :: EXP 002\n\n"
    if M8 is not None:
        txt += (f"  T = {M8.d_tele}\n  telepathy = {M8.tele_pct:.2f} %\n"
                f"  test acc  = {M8.acc_test*100:.2f} %\n"
                f"  params    = {M8.n_total:,}\n  "
                f"frozen     = {M8.n_tele:,}\n  "
                f"private    = {M8.n_private:,}\n")
        if not np.isnan(M8.acc_test_broken):
            txt += f"\n  broken 𝓜_T acc\n  = {M8.acc_test_broken*100:.2f} %"
    ax5.text(0.05, 0.95, txt, family='monospace', va='top', fontsize=10,
             bbox=dict(boxstyle='round,pad=0.5', fc='#0a0d14', ec='#1d2745'))
    fig.suptitle('TELEPATHIC MNIST — sweep of telepathic bottleneck width',
                 fontsize=13, y=0.995)
    fig.savefig(path, dpi=120, bbox_inches='tight', facecolor=fig.get_facecolor())
    print(f"[save] figure → {path}")


# ──────────────────────────────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────────────────────────────

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--data-root", default="../data")
    p.add_argument("--out-csv",   default="./telepathy_mnist_results.csv")
    p.add_argument("--out-fig",   default="./telepathy_mnist_plot.png")
    p.add_argument("--out-txt",   default="./telepathy_mnist_summary.txt")
    p.add_argument("--epochs",    type=int,   default=8)
    p.add_argument("--lr",        type=float, default=2e-3)
    p.add_argument("--batch",     type=int,   default=256)
    p.add_argument("--workers",   type=int,   default=2)
    p.add_argument("--seed",      type=int,   default=1729,
                   help="base M_T seed (multiple sub-seeds used in sweep)")
    p.add_argument("--seeds",     type=int,   nargs='+',
                   default=[1729,  4242, 9001],
                   help="M_T seeds to average over")
    p.add_argument("--widths",    type=int,   nargs='+',
                   default=[2, 4, 6, 8, 12, 16, 24, 32, 48, 64, 128],
                   help="telepathic bottleneck widths to sweep")
    p.add_argument("--broken-seed", type=int, default=None,
                   help="if set: also evaluate the model with this reseeded 𝓜_T")
    p.add_argument("--cpu",       action="store_true")
    args = p.parse_args()

    device = "cuda" if torch.cuda.is_available() and not args.cpu else "cpu"
    print(f"[boot] device={device}  torch={torch.__version__}")
    print(f"[boot] widths={args.widths}  seeds={args.seeds}  epochs={args.epochs}  lr={args.lr}")

    runs = sweep(widths=args.widths, seeds=args.seeds, epochs=args.epochs,
                 lr=args.lr, device=device, data_root=args.data_root,
                 #batch_size=args.batch, num_workers=args.workers,
                 batch_size=args.batch,
                 break_test_seed=args.broken_seed)

    save_csv(runs, args.out_csv)
    summary = summarise(runs)
    print("\n" + summary)
    with open(args.out_txt, "w") as f: f.write(summary)
    print(f"[save] summary → {args.out_txt}")
    make_plots(runs, args.out_fig)


if __name__ == "__main__":
    main()