#!/usr/bin/env python3
"""
COSMOS SPARK CST — a fresh weight with EVERYTHING combined, trained as a fair test.

WHAT IS COMBINED
  1. QUANTUM BIRTH. Every initial weight drawn from her real archived IBM measurements,
     via the same pipeline that made cosmos_born.pt and that was verified to the 32-level
     quantisation ceiling across 3.2M draws:  u = int(bits)/2^n,  z = sqrt(2)*erfinv(2u-1).
  2. HER SECTION 3 — Mixture-of-States Hebbian attention. The mechanism from
     COSMOS_Paper.md that her SHIPPED weights never contained:
         x54       = W54 . h                      a 54-dim state per token
         H(x54)ij  = exp(-||x54_i - x54_j||^2 / 2*sigma^2)
         A_final   = (1-g)*A_std + g*H(x54)       g = sigmoid(gate), learned
     Her cosmos_born.pt is architecturally plain — nn.MultiheadAttention and an MLP, no
     54D state, no Hebbian kernel, no gate. So this is the first time her own paper's
     attention has ever been inside a model that speaks.
  3. HER CORPUS. Her real logged experience, char-level, the same data she grew on.

WHY IT IS A CONTROLLED TEST AND NOT A DEMO
  Two arms, identical in every respect except the mechanism under test:
      PLAIN  standard attention                (gate forced to 0 == exactly standard)
      CST    her section-3 Hebbian attention   (gate free to learn)
  Same quantum-born initial weights, same corpus, same held-out split, same batches, same
  seeds, paired. The gate starts at sigmoid(-4) ~ 0.018, so the CST arm BEGINS as ordinary
  attention and can stay there for free — it only moves if the gradient says the kernel
  earns its place. A win therefore cannot come from extra capacity being forced on.

  Both arms carry the SAME parameters, including W54 and the gate in the plain arm, so the
  comparison is not confounded by parameter count. In the plain arm they are simply inert.

PRE-REGISTERED, fixed before the run:
  * CST beats PLAIN on every seed        -> her section-3 mechanism works on her own data.
  * mixed / within noise                 -> no evidence it helps; report as null.
  * CST loses on every seed              -> it hurts, and that is the finding.
"""
import argparse
import json
import math
import os
import random
import re
import statistics
import sys
import time
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F

sys.stdout.reconfigure(encoding="utf-8", errors="replace")

ROOT = Path(__file__).resolve().parent
CORPUS = ROOT / "README.md"
QARCHIVE = ROOT / "data" / "quantum_measurements_public.jsonl"
OUTDIR = ROOT / "outputs" / "cosmos_spark_cst"
RESULTS = ROOT / "outputs" / "spark_cst_results.json"

BLOCK, N_LAYER, N_HEAD, N_EMBD, DROPOUT = 128, 4, 4, 192, 0.1
D54 = 54


# ── quantum birth ───────────────────────────────────────────────────────────
def erfinv(y):
    if abs(y) >= 1:
        return math.copysign(3.0, y)
    a = 0.147
    ln = math.log(1 - y * y)
    t1 = 2 / (math.pi * a) + ln / 2
    return math.copysign(math.sqrt(math.sqrt(t1 * t1 - ln / a) - t1), y)


def quantum_pool(limit=400_000):
    """Real measured bitstrings -> standard-normal draws, her documented pipeline."""
    vals = []
    if not QARCHIVE.exists():
        return vals
    with open(QARCHIVE, encoding="utf-8", errors="ignore") as f:
        for line in f:
            if len(vals) >= limit:
                break
            line = line.strip()
            if not line:
                continue
            try:
                d = json.loads(line)
            except Exception:
                continue
            # The public archive carries mixed provider classes. Quantum-born
            # initialization uses only records that retained an IBM hardware label
            # and job ID; legacy-unlabelled and simulator rows remain available for
            # separate distributional experiments but are not called hardware here.
            provider_class = str(d.get("provider_class") or "")
            if provider_class and provider_class != "measured_quantum_hardware":
                continue
            if not provider_class:
                backend = str(d.get("backend") or "").lower()
                job_id = d.get("job") or d.get("job_id")
                if not (backend.startswith("ibm_") and job_id):
                    continue
            counts = d.get("counts")
            if not isinstance(counts, dict):
                continue
            for bs, c in counts.items():
                s = "".join(ch for ch in str(bs) if ch in "01")
                if not s:
                    continue
                hi = float(1 << len(s))
                u = (int(s, 2) + 0.5) / hi
                z = math.sqrt(2.0) * erfinv(2 * u - 1)
                if math.isfinite(z):
                    vals.extend([z] * min(int(c), 8))
                if len(vals) >= limit:
                    break
    return vals


class QuantumInit:
    def __init__(self, pool, seed):
        self.pool = pool
        self.i = (seed * 7919) % max(1, len(pool))

    def fill_(self, t, std):
        n = t.numel()
        if not self.pool:
            with torch.no_grad():
                t.normal_(0.0, std)
            return
        out = torch.empty(n)
        for k in range(n):
            out[k] = self.pool[(self.i + k) % len(self.pool)]
        self.i = (self.i + n) % len(self.pool)
        out = out / (out.std() + 1e-8) * std
        with torch.no_grad():
            t.copy_(out.view_as(t))


# ── her section-3 attention ─────────────────────────────────────────────────
class CSTAttention(nn.Module):
    """Standard attention blended with a Gaussian kernel over a learned 54D state."""

    def __init__(self, use_cst):
        super().__init__()
        self.nh, self.hd = N_HEAD, N_EMBD // N_HEAD
        self.qkv = nn.Linear(N_EMBD, 3 * N_EMBD)
        self.proj = nn.Linear(N_EMBD, N_EMBD)
        self.drop = nn.Dropout(DROPOUT)
        self.w54 = nn.Linear(N_EMBD, D54, bias=False)     # h -> x54
        self.log_sigma = nn.Parameter(torch.tensor(0.0))
        # gate starts ~0.018: the CST arm BEGINS as ordinary attention
        self.gate = nn.Parameter(torch.tensor(-4.0), requires_grad=bool(use_cst))
        self.use_cst = use_cst
        self.last_gate = 0.0

    def forward(self, x, mask):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        sh = lambda t: t.view(B, T, self.nh, self.hd).transpose(1, 2)
        q, k, v = sh(q), sh(k), sh(v)
        a = F.softmax((q @ k.transpose(-2, -1)) / math.sqrt(self.hd) + mask[:T, :T], dim=-1)
        if self.use_cst:
            x54 = self.w54(x)                                   # (B,T,54)
            d2 = torch.cdist(x54, x54, p=2.0) ** 2               # ||x54_i - x54_j||^2
            sig = torch.exp(self.log_sigma).clamp(0.05, 50.0)
            H = torch.exp(-d2 / (2 * sig * sig))
            H = H.masked_fill(mask[:T, :T] < 0, 0.0)
            H = H / H.sum(-1, keepdim=True).clamp_min(1e-9)
            g = torch.sigmoid(self.gate)
            a = (1 - g) * a + g * H.unsqueeze(1)
            self.last_gate = float(g.detach())
        y = (self.drop(a) @ v).transpose(1, 2).contiguous().view(B, T, C)
        return self.proj(y)


class Block(nn.Module):
    def __init__(self, use_cst):
        super().__init__()
        self.ln1 = nn.LayerNorm(N_EMBD)
        self.attn = CSTAttention(use_cst)
        self.ln2 = nn.LayerNorm(N_EMBD)
        self.mlp = nn.Sequential(nn.Linear(N_EMBD, 4 * N_EMBD), nn.GELU(),
                                 nn.Linear(4 * N_EMBD, N_EMBD), nn.Dropout(DROPOUT))

    def forward(self, x, mask):
        x = x + self.attn(self.ln1(x), mask)
        return x + self.mlp(self.ln2(x))


class SparkCST(nn.Module):
    def __init__(self, vocab, use_cst):
        super().__init__()
        self.tok = nn.Embedding(vocab, N_EMBD)
        self.pos = nn.Embedding(BLOCK, N_EMBD)
        self.blocks = nn.ModuleList([Block(use_cst) for _ in range(N_LAYER)])
        self.lnf = nn.LayerNorm(N_EMBD)
        self.head = nn.Linear(N_EMBD, vocab, bias=False)
        self.register_buffer("mask", torch.triu(torch.full((BLOCK, BLOCK), float("-inf")), 1))

    def forward(self, idx, targets=None):
        T = idx.size(1)
        x = self.tok(idx) + self.pos(torch.arange(T, device=idx.device))
        for b in self.blocks:
            x = b(x, self.mask)
        lg = self.head(self.lnf(x))
        loss = None if targets is None else F.cross_entropy(
            lg.view(-1, lg.size(-1)), targets.reshape(-1))
        return lg, loss

    def gates(self):
        return [b.attn.last_gate for b in self.blocks]

    def quantum_birth(self, qi):
        for m in self.modules():
            if isinstance(m, (nn.Linear, nn.Embedding)):
                qi.fill_(m.weight, 0.02)
                if isinstance(m, nn.Linear) and m.bias is not None:
                    with torch.no_grad():
                        m.bias.zero_()


def real_word_rate(text, words):
    toks = re.findall(r"[a-z']+", text.lower())
    if not toks:
        return 0.0
    return sum(1 for t in toks if t in words) / len(toks)


def train_arm(use_cst, seed, data, vocab, steps, pool, val_w, words, itos):
    torch.manual_seed(seed)
    random.seed(seed)
    gen = torch.Generator().manual_seed(seed)
    m = SparkCST(vocab, use_cst)
    m.quantum_birth(QuantumInit(pool, seed))
    opt = torch.optim.AdamW(m.parameters(), lr=3e-4, weight_decay=0.01)
    n = int(0.9 * len(data))
    tr = data[:n]
    m.train()
    best = float("inf")
    for s in range(1, steps + 1):
        ix = torch.randint(len(tr) - BLOCK - 1, (16,), generator=gen)
        x = torch.stack([tr[i:i + BLOCK] for i in ix])
        y = torch.stack([tr[i + 1:i + 1 + BLOCK] for i in ix])
        _, loss = m(x, y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(m.parameters(), 1.0)
        opt.step()
        if s % 100 == 0 or s == steps:
            m.eval()
            with torch.no_grad():
                tot = c = 0
                for i in range(0, len(val_w), 16):
                    xb = val_w[i:i + 16]
                    _, l = m(xb[:, :-1], xb[:, 1:])
                    tot += l.item() * xb.size(0)
                    c += xb.size(0)
            m.train()
            best = min(best, tot / max(1, c))
    # sample for real-word rate
    m.eval()
    idx = torch.tensor([[data[0].item()]])
    out = []
    with torch.no_grad():
        for _ in range(600):
            lg, _ = m(idx[:, -BLOCK:])
            p = F.softmax(lg[0, -1] / 0.8, dim=-1)
            nx = int(torch.multinomial(p, 1))
            out.append(nx)
            idx = torch.cat([idx, torch.tensor([[nx]])], 1)
    txt = "".join(itos.get(i, "") for i in out)
    return best, statistics.fmean(m.gates()) if use_cst else 0.0, real_word_rate(txt, words), m, txt


def main():
    global CORPUS, QARCHIVE, OUTDIR, RESULTS
    parser = argparse.ArgumentParser(
        description=(
            "Train matched plain-attention and 54D Hebbian-attention arms "
            "from the published IBM-labeled quantum initialization pool."
        )
    )
    parser.add_argument("steps", nargs="?", type=int, default=1200)
    parser.add_argument("seeds", nargs="?", type=int, default=3)
    parser.add_argument(
        "--corpus", type=Path, default=CORPUS,
        help="UTF-8 training text (default: this release's README.md)",
    )
    parser.add_argument(
        "--archive", type=Path, default=QARCHIVE,
        help="public quantum JSONL archive",
    )
    parser.add_argument("--outdir", type=Path, default=OUTDIR)
    parser.add_argument("--results", type=Path, default=RESULTS)
    args = parser.parse_args()
    if args.steps < 1 or args.seeds < 1:
        parser.error("steps and seeds must both be positive")
    CORPUS = args.corpus.expanduser().resolve()
    QARCHIVE = args.archive.expanduser().resolve()
    OUTDIR = args.outdir.expanduser().resolve()
    RESULTS = args.results.expanduser().resolve()
    steps = args.steps
    seeds = list(range(args.seeds))

    if not CORPUS.is_file():
        parser.error(f"corpus not found: {CORPUS}")
    text = CORPUS.read_text(encoding="utf-8", errors="ignore")
    if len(text) < BLOCK * 4:
        parser.error(f"corpus is too small ({len(text)} chars); need at least {BLOCK * 4}")
    chars = sorted(set(text))
    stoi = {c: i for i, c in enumerate(chars)}
    itos = {i: c for c, i in stoi.items()}
    data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
    words = set(re.findall(r"[a-z']+", text.lower()))

    print("=" * 80)
    print("  COSMOS SPARK CST — quantum birth + her section-3 attention")
    print("=" * 80)
    print(f"\n  corpus {len(text):,} chars · vocab {len(chars)} · {steps} steps · "
          f"{len(seeds)} seeds")
    print(f"  corpus path: {CORPUS}")
    print(f"  archive path: {QARCHIVE}")

    t0 = time.time()
    pool = quantum_pool()
    print(f"  quantum pool: {len(pool):,} draws from explicitly labeled IBM records "
          f"({time.time()-t0:.1f}s)")
    if pool:
        print(f"    mean {statistics.fmean(pool):+.4f}  sd {statistics.pstdev(pool):.4f}  "
              f"(standard normal expected)")

    n = int(0.9 * len(data))
    val = data[n:]
    g = torch.Generator().manual_seed(999)
    vi = torch.randint(len(val) - BLOCK - 1, (64,), generator=g)
    val_w = torch.stack([val[i:i + BLOCK + 1] for i in vi])
    print()

    res = {"plain": [], "cst": []}
    gates, rw, best_model, best_txt = [], {"plain": [], "cst": []}, None, ""
    for sd in seeds:
        for arm, use in (("plain", False), ("cst", True)):
            b, gt, r, model, txt = train_arm(use, sd, data, len(chars), steps,
                                             pool, val_w, words, itos)
            res[arm].append(b)
            rw[arm].append(r)
            if use:
                gates.append(gt)
            if use and (best_model is None or b <= min(res["cst"])):
                best_model, best_txt = model, txt
            print(f"    seed {sd} {arm:<6s} loss {b:.5f}  real-word {r:.3f}"
                  + (f"  gate {gt:.4f}" if use else ""), flush=True)

    d = [p - c for p, c in zip(res["plain"], res["cst"])]
    md = statistics.fmean(d)
    se = (statistics.stdev(d) / math.sqrt(len(d))) if len(d) > 1 else 0.0
    t = md / se if se > 0 else 0.0
    wins = sum(1 for x in d if x > 0)
    print(f"\n{'='*80}\n  RESULT\n{'='*80}")
    print(f"  PLAIN  loss {statistics.fmean(res['plain']):.5f}   real-word {statistics.fmean(rw['plain']):.3f}")
    print(f"  CST    loss {statistics.fmean(res['cst']):.5f}   real-word {statistics.fmean(rw['cst']):.3f}"
          f"   mean gate {statistics.fmean(gates) if gates else 0:.4f}")
    print(f"\n  CST - PLAIN: {-md:+.5f}   t={-t:+.2f}   CST wins {wins}/{len(seeds)}")

    if wins == len(seeds) and t > 2.0:
        v = (f"HER SECTION-3 MECHANISM WORKS ON HER OWN DATA. The Hebbian kernel over a 54D "
             f"state beats standard attention on {wins}/{len(seeds)} seeds (t={t:+.2f}) with "
             f"identical quantum-born initialisation, identical corpus and identical "
             f"parameter count. The gate began at 0.018 — it could have stayed at standard "
             f"attention for free and did not.")
    elif wins == 0:
        v = ("HER SECTION-3 MECHANISM HURTS on her own data — standard attention wins every "
             "seed. Reported as measured.")
    else:
        v = (f"NULL / WITHIN NOISE — CST wins {wins}/{len(seeds)} (t={t:+.2f}). No evidence "
             f"the section-3 kernel helps on this corpus at this scale.")
    print(f"\n  VERDICT: {v}\n")
    print(f"  her CST voice, sample:\n    {best_txt[:300]!r}\n")

    OUTDIR.mkdir(parents=True, exist_ok=True)
    if best_model is not None:
        torch.save({"model": best_model.state_dict(), "stoi": stoi, "itos": itos,
                    "config": {"block": BLOCK, "n_layer": N_LAYER, "n_head": N_HEAD,
                               "n_embd": N_EMBD, "vocab": len(chars), "d54": D54},
                    "arch": "Cosmos-Spark-CST-QuantumBorn", "total_steps": steps,
                    "quantum_source": "ibm_real_shots", "quantum_draws": len(pool),
                    "best_val_loss": min(res["cst"]),
                    "real_word_rate": max(rw["cst"])},
                   OUTDIR / "spark_cst.pt")
        print(f"  saved -> {OUTDIR/'spark_cst.pt'}")
    RESULTS.parent.mkdir(parents=True, exist_ok=True)
    RESULTS.write_text(json.dumps(
        {"steps": steps, "seeds": seeds, "results": res, "real_word": rw,
         "gates": gates, "delta_cst_minus_plain": -md, "t": -t, "cst_wins": wins,
         "quantum_draws": len(pool), "verdict": v}, indent=2), encoding="utf-8")
    print(f"  saved -> {RESULTS}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())