#!/usr/bin/env python3
"""Read-only Cascade-EP nudged-relaxation rho probe.

This is deliberately standalone: the trainer executes training at import time.  The
model definitions, free-state construction, and relax sweep below are copied from
casc_eq_train.py.  No optimizer step or parameter update is performed.
"""

import gc
import math
import os
import pickle
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


ROOT = Path(__file__).resolve().parent
CHECKPOINT = ROOT / "runs/fw72m_plain2_s35000.pt"
DATA_DIR = ROOT / "data/fineweb_edu"
BETAS = [1e-3, 3e-3, 1e-2, 3e-2, 9e-2, 0.2, 0.4, 0.7, 1.0]

# The requested run is CPU-only in this environment.  Keeping this explicit also
# makes it impossible for the probe to select GPU2.
DEVICE = torch.device("cpu")
torch.set_num_threads(min(32, os.cpu_count() or 1))
torch.manual_seed(1)


class RMSNorm(nn.Module):
    def __init__(self, C, eps=1e-6):
        super().__init__()
        self.g = nn.Parameter(torch.ones(C))
        self.eps = eps

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.g


class SwiGLU(nn.Module):
    def __init__(self, C):
        super().__init__()
        h = ((8 * C // 3) + 63) // 64 * 64
        self.w1 = nn.Linear(C, h, bias=False)
        self.w3 = nn.Linear(C, h, bias=False)
        self.w2 = nn.Linear(h, C, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))


class Olmo2Attn(nn.Module):
    def __init__(self, C, H, T):
        super().__init__()
        self.H, self.hd = H, C // H
        self.qkv = nn.Linear(C, 3 * C, bias=False)
        self.proj = nn.Linear(C, C, bias=False)
        self.qn, self.kn = RMSNorm(C), RMSNorm(C)
        inv = 1.0 / (500000.0 ** (torch.arange(0, self.hd, 2).float() / self.hd))
        fr = torch.outer(torch.arange(T).float(), inv)
        self.register_buffer("rc", fr.cos(), persistent=False)
        self.register_buffer("rs", fr.sin(), persistent=False)

    def rope(self, x):
        x1, x2 = x[..., ::2], x[..., 1::2]
        c, s = self.rc[None, None], self.rs[None, None]
        return torch.stack((x1 * c - x2 * s, x1 * s + x2 * c), dim=-1).flatten(-2)

    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=2)
        q, k = self.qn(q), self.kn(k)
        q = self.rope(q.view(B, T, self.H, self.hd).transpose(1, 2))
        k = self.rope(k.view(B, T, self.H, self.hd).transpose(1, 2))
        v = v.view(B, T, self.H, self.hd).transpose(1, 2)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.proj(y.transpose(1, 2).contiguous().view(B, T, C))


class Olmo2Block(nn.Module):
    def __init__(self, C, H, T):
        super().__init__()
        self.attn = Olmo2Attn(C, H, T)
        self.ff = SwiGLU(C)
        self.na, self.nf = RMSNorm(C), RMSNorm(C)

    def forward(self, z, mask=None):
        z = z + self.na(self.attn(z))
        return z + self.nf(self.ff(z))


def fmt(value):
    if value is None:
        return "NA"
    return f"{value:.8e}"


print(f"checkpoint={CHECKPOINT}", flush=True)
print(f"device={DEVICE} torch_threads={torch.get_num_threads()}", flush=True)
ck = torch.load(CHECKPOINT, map_location=DEVICE, weights_only=False)
cfg = ck["config"]
L, C, H, T = (int(cfg[k]) for k in ("L", "C", "H", "T"))
vocab = int(ck["tok"]["weight"].shape[0])
T2_MAX = int(cfg.get("kmax", 8))
DTOP_EVERY = int(cfg.get("dtop_every", 1))
GETA = float(cfg.get("geta", 1.0))
AMP = bool(cfg.get("amp", False) and DEVICE.type == "cuda")
# Below this relative state-change scale, the observed float32 residuals are a
# numerical fixed-point floor and rho becomes noise/noise (the trainer documents
# the same failure mode).  Preserve raw GOV rho, but do not call noise a ceiling.
RHO_NOISE_FLOOR = 1e-6

tok = nn.Embedding(vocab, C).to(DEVICE)
pos = nn.Embedding(T, C).to(DEVICE)  # loaded for checkpoint parity; OLMo2 does not use it
blocks = nn.ModuleList([Olmo2Block(C, H, T) for _ in range(L)]).to(DEVICE)
ln_f = RMSNorm(C).to(DEVICE)
W_out = nn.Parameter(torch.empty(vocab, C, device=DEVICE))

tok.load_state_dict(ck["tok"])
pos.load_state_dict(ck["pos"])
blocks.load_state_dict(ck["blocks"])
ln_f.load_state_dict(ck["lnf"])
with torch.no_grad():
    W_out.copy_(ck["wout"].to(DEVICE))

# Restore the checkpoint optimizer state exactly as the trainer does, then discard
# it: this verifies full checkpoint compatibility without ever taking a train step.
all_params = (
    list(tok.parameters())
    + list(blocks.parameters())
    + list(ln_f.parameters())
    + [W_out]
)
from muon import build_hybrid

opt, sched = build_hybrid(
    blocks,
    all_params,
    float(cfg["lr"]),
    float(cfg.get("muon_lr", 0.02)),
    int(cfg["warmup"]),
    muon_mom=float(cfg.get("muon_mom", 0.95)),
    adam_b1=float(cfg.get("adam_b1", 0.9)),
    total_steps=(int(cfg["steps"]) if cfg.get("cosine", False) else 0),
    lr_min_ratio=float(cfg.get("lr_min_ratio", 0.1)),
)
opt.load_state_dict(ck["opt"])
print(
    f"loaded_step={int(ck.get('step', 0))} model_state=yes optimizer_state=yes "
    f"L={L} C={C} H={H} T={T} vocab={vocab}",
    flush=True,
)
del opt, sched
ck.pop("opt", None)
gc.collect()

for module in (tok, pos, blocks, ln_f):
    module.eval()

mask = torch.triu(torch.full((T, T), float("-inf"), device=DEVICE), 1)


def emb(x):
    return tok(x)


def readout(z):
    return ln_f(z) @ W_out.t()


def obj_loss(logits2d, y1d):
    return F.cross_entropy(logits2d, y1d)


def free_states_graphed(x):
    """Exact trainer free-state construction."""
    with torch.no_grad():
        z0 = emb(x)
    ins, outs, zs = [], [], []
    prev = z0
    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=AMP):
        for block in blocks:
            i = prev.detach().requires_grad_(True)
            o = block(i, mask)
            ins.append(i)
            outs.append(o)
            zs.append(o.detach().float())
            prev = zs[-1]
    return z0, zs, ins, outs


def relax_probe(z0, zs, ins, outs, y, beta, K, x):
    """Trainer relax(), plus a read-only copy of every per-sweep residual."""
    d = [None] * L
    geta_l = GETA
    residuals = []

    def forces(refresh_top):
        if refresh_top or d[L - 1] is None:
            zc = zs[L - 1].detach().requires_grad_(True)
            ce = obj_loss(readout(zc).reshape(-1, vocab), y.reshape(-1))
            nbt_loc = zc.shape[0] * zc.shape[1]
            g = torch.autograd.grad(ce, zc)[0]
            d[L - 1] = (-beta * nbt_loc * g).detach()
        for layer in range(L - 2, -1, -1):
            d[layer] = torch.autograd.grad(
                outs[layer + 1],
                ins[layer + 1],
                grad_outputs=d[layer + 1].to(outs[layer + 1].dtype),
            )[0].detach().float()

    def rebuild(last):
        nonlocal ins, outs
        prev = z0
        n_ins, n_outs = [], []
        rnum = rden = 0.0
        g_eff = 1.0 if last else geta_l
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=AMP):
            for layer in range(L):
                if last and layer == 0:
                    i = emb(x)
                else:
                    i = prev.detach().requires_grad_(True)
                o = blocks[layer](i, mask)
                znew = o.detach().float() + d[layer]
                mixed = znew if g_eff >= 1.0 else (
                    zs[layer] + g_eff * (znew - zs[layer])
                )
                with torch.no_grad():
                    rnum += float((mixed - zs[layer]).norm())
                    rden += float(zs[layer].norm())
                zs[layer] = mixed
                n_ins.append(i)
                n_outs.append(o)
                prev = zs[layer]
        ins, outs = n_ins, n_outs
        return rnum / max(rden, 1e-9)

    for k in range(K):
        forces(k % DTOP_EVERY == 0)
        residuals.append(rebuild(k + 1 == K))

    rho = None
    if len(residuals) >= 2 and residuals[-2] > 1e-12:
        rho = residuals[-1] / residuals[-2]
    gov = {"res": residuals[-1], "rho": rho, "kuse": K}
    return residuals, gov


# One deterministic, fixed batch from the same binary loader as the trainer.  B=1
# is sufficient because the trainer's nbt_loc factor cancels CE's batch averaging.
train_data = np.memmap(DATA_DIR / "train.bin", dtype=np.uint16, mode="r")
batch_gen = torch.Generator().manual_seed(1 * 7919 + 11)
offset = int(torch.randint(len(train_data) - T - 1, (1,), generator=batch_gen).item())
x = torch.from_numpy(train_data[offset : offset + T].astype(np.int64))[None].to(DEVICE)
y = torch.from_numpy(train_data[offset + 1 : offset + T + 1].astype(np.int64))[None].to(DEVICE)
print(
    f"batch=fineweb_edu/train.bin offset={offset} B=1 T={T} "
    f"T2_max={T2_MAX} dtop_every={DTOP_EVERY} geta={GETA} amp={AMP}",
    flush=True,
)
print(
    f"convergence_rule=numerical_fixed_point(res<={RHO_NOISE_FLOOR:g}) or "
    "finite(last_above_floor_rho)<1 (fixed-K trainer has no tolerance stop)",
    flush=True,
)

results = []
for beta in BETAS:
    z0, zs, ins, outs = free_states_graphed(x)
    try:
        residuals, gov = relax_probe(z0, zs, ins, outs, y, beta, T2_MAX, x)
        rhos = [
            (None if i == 0 or residuals[i - 1] <= 1e-12 else residuals[i] / residuals[i - 1])
            for i in range(len(residuals))
        ]
        final_res = gov["res"]
        raw_gov_rho = gov["rho"]
        meaningful = [
            rhos[i]
            for i in range(1, len(rhos))
            if rhos[i] is not None
            and (residuals[i - 1] > RHO_NOISE_FLOOR or residuals[i] > RHO_NOISE_FLOOR)
        ]
        rho = meaningful[-1] if meaningful else raw_gov_rho
        reached_floor = any(r <= RHO_NOISE_FLOOR for r in residuals)
        converged = bool(
            math.isfinite(final_res)
            and rho is not None
            and math.isfinite(rho)
            and (reached_floor or rho < 1.0)
        )
        print(
            f"TRACE beta={beta:g} residuals=[{','.join(fmt(v) for v in residuals)}] "
            f"rhos=[{','.join(fmt(v) for v in rhos)}] "
            f"raw_final_gov_rho={fmt(raw_gov_rho)}",
            flush=True,
        )
    except (RuntimeError, FloatingPointError) as exc:
        final_res, rho, converged = float("nan"), None, False
        print(f"TRACE beta={beta:g} relaxation_error={type(exc).__name__}:{exc}", flush=True)

    results.append((beta, final_res, converged, rho))
    print(
        f"RESULT beta={beta:g} final_res={fmt(final_res)} "
        f"converged={'yes' if converged else 'no'} rho_at_convergence={fmt(rho)}",
        flush=True,
    )
    del z0, zs, ins, outs
    gc.collect()

beta_star = next((beta for beta, _, converged, rho in results if (not converged) or (rho is not None and rho >= 1.0)), None)
print(f"BETA_STAR {beta_star if beta_star is not None else 'NONE'}", flush=True)