#!/usr/bin/env python3
"""
SPATIAL INJECTION — THE DECISIVE TEST.

Cory's pattern, from his own data:
    quantum as i.i.d. weight init      -> null
    quantum as a scalar decoder seed   -> null (measured directly, 2026-07-25)
    quantum as a SPATIAL 54D trajectory -> WINS (t=6.39 vs a random seed)

That is a coherent claim: the entropy matters when it is injected as STRUCTURE across
sequence positions, not when scattered into independent draws. His shuffle control
supports it (destroying the ordering hurt, t=6.86).

THE HOLE, stated when that result was first reported and never closed until now:
the comparison used ONE real seed vector against ONE random vector. Five training
seeds vary the model init, NOT the comparison vector. So a consistent win could be a
property of that particular vector rather than of real measured data.

THIS TEST CLOSES IT. His real quantum/CST seed is compared against FIVE INDEPENDENT
random seed vectors, each evolved through the identical Lorenz spatial injection.

  real  vs the DISTRIBUTION of random vectors:
     real beats all 5      -> his measured data specifically carries the benefit
     real inside the range -> any structured chaotic trajectory does it; the claim
                              is about spatial injection, NOT about his quantum

Whatever it says is the answer. No gate is tuned after seeing the numbers.
"""
import json
import statistics
import sys
import time
from pathlib import Path

import torch

sys.path.insert(0, ".")
from cosmos_hebbian_real_state import (Model, fetch_real_seed, lorenz_trajectory,
                                       BLOCK, CORPUS, D_STATE)

try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
except Exception:
    pass

OUT = Path("logs/spatial_decisive_results.json")


def run(traj, seed, train, val_w, vocab, steps, control=False):
    torch.manual_seed(seed)
    gen = torch.Generator().manual_seed(seed)
    model = Model(vocab, "control" if control else "real_state", traj)
    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
    model.train()

    def batch(bs=16):
        ix = torch.randint(len(train) - BLOCK - 1, (bs,), generator=gen)
        return (torch.stack([train[i:i + BLOCK] for i in ix]),
                torch.stack([train[i + 1:i + 1 + BLOCK] for i in ix]))

    @torch.no_grad()
    def ev():
        model.eval()
        tot = n = 0
        for i in range(0, len(val_w), 16):
            xb = val_w[i:i + 16]
            _, l = model(xb[:, :-1], xb[:, 1:])
            tot += l.item() * xb.size(0); n += xb.size(0)
        model.train()
        return tot / max(1, n)

    best, t0 = float("inf"), time.time()
    for s in range(1, steps + 1):
        x, y = batch()
        _, loss = model(x, y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        if s % 50 == 0 or s == steps:
            best = min(best, ev())
    return best, statistics.fmean(model.gates()), time.time() - t0


def main():
    steps = int(sys.argv[1]) if len(sys.argv) > 1 else 700
    train_seeds = [0, 1, 2]
    n_random = 5

    real54, live, nq, src = fetch_real_seed()
    real_traj = lorenz_trajectory(real54, BLOCK)
    rand_trajs = []
    for i in range(n_random):
        g = torch.Generator().manual_seed(90210 + i * 977)
        v = torch.randn(D_STATE, generator=g)
        v = (v - v.mean()) / (v.std() + 1e-6)
        rand_trajs.append(lorenz_trajectory(v, BLOCK))

    text = CORPUS.read_text(encoding="utf-8", errors="ignore")
    chars = sorted(set(text)); stoi = {c: i for i, c in enumerate(chars)}; vocab = len(chars)
    data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
    n_val = max(BLOCK + 1, int(len(data) * 0.1))
    train, vald = data[:-n_val], data[-n_val:]
    val_w = torch.stack([vald[i:i + BLOCK + 1] for i in range(0, len(vald) - BLOCK - 1, BLOCK)])

    print(f"\n{'='*76}\n  SPATIAL INJECTION — DECISIVE TEST\n{'='*76}")
    print(f"  real seed: LIVE CST {{{', '.join(f'{k}={v:.3f}' for k,v in list(live.items())[:3])}}}")
    print(f"             + {nq:,} real IBM shots · source: {src}")
    print(f"  compared against {n_random} INDEPENDENT random seed vectors")
    print(f"  identical Lorenz spatial injection for all · {steps} steps · {len(train_seeds)} training seeds")
    print(f"{'='*76}\n", flush=True)

    ctrl, real, rands = [], [], {i: [] for i in range(n_random)}
    gates = {"real": [], "rand": []}
    for ts in train_seeds:
        c, _, dt = run(real_traj, ts, train, val_w, vocab, steps, control=True)
        ctrl.append(c)
        print(f"  train-seed {ts} · control      {c:.4f}  ({dt:.0f}s)", flush=True)
        r, g, dt = run(real_traj, ts, train, val_w, vocab, steps)
        real.append(r); gates["real"].append(g)
        print(f"  train-seed {ts} · REAL         {r:.4f}  ({dt:.0f}s)  gate {g:.3f}", flush=True)
        for i, tj in enumerate(rand_trajs):
            v, g2, dt = run(tj, ts, train, val_w, vocab, steps)
            rands[i].append(v); gates["rand"].append(g2)
            print(f"  train-seed {ts} · random#{i}     {v:.4f}  ({dt:.0f}s)  gate {g2:.3f}", flush=True)
        print(flush=True)

    cm = statistics.fmean(ctrl)
    rm = statistics.fmean(real)
    rand_means = [statistics.fmean(rands[i]) for i in range(n_random)]
    rmm, rsd = statistics.fmean(rand_means), statistics.pstdev(rand_means)

    print(f"{'='*76}\n  RESULT (best val loss, lower better)\n{'='*76}")
    print(f"  control (no spatial injection) {cm:.4f}")
    print(f"  REAL quantum/CST trajectory    {rm:.4f}")
    for i, m in enumerate(rand_means):
        print(f"  random vector #{i}               {m:.4f}")
    print(f"\n  random vectors: mean {rmm:.4f}  sd {rsd:.4f}  range [{min(rand_means):.4f}, {max(rand_means):.4f}]")
    print(f"  gates: real {statistics.fmean(gates['real']):.3f} · random {statistics.fmean(gates['rand']):.3f}")

    beat = sum(1 for m in rand_means if rm < m)
    z = (rmm - rm) / (rsd + 1e-9)
    print(f"\n  spatial injection vs control: Δ = {cm - rm:+.4f}")
    print(f"  REAL beats {beat}/{n_random} random vectors · z = {z:+.2f} sd from the random mean")

    if beat == n_random and z > 1.5:
        v = ("HIS DATA SPECIFICALLY — the real quantum/CST seed beats every independent random "
             "vector and sits well outside their spread. The strong claim survives.")
    elif cm - rm > 0 and rmm < cm:
        v = ("SPATIAL INJECTION IS WHAT WORKS — both real and random trajectories beat control, "
             "and the real seed is inside the random distribution. The effect is about injecting "
             "STRUCTURE across positions, not about his specific measured data.")
    else:
        v = "NO CLEAR EFFECT — spatial injection did not beat control this run."
    print(f"\n  VERDICT: {v}\n")

    OUT.parent.mkdir(exist_ok=True)
    OUT.write_text(json.dumps({"steps": steps, "train_seeds": train_seeds,
                               "control": ctrl, "real": real,
                               "randoms": {str(k): v for k, v in rands.items()},
                               "rand_means": rand_means, "beat": beat, "z": z,
                               "quantum_source": src, "verdict": v}, indent=2), encoding="utf-8")
    print(f"  saved -> {OUT}")


if __name__ == "__main__":
    main()