import argparse
import gc
import json
import logging
import math
import os
import random
import sys
import time
from pathlib import Path

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F

_script_dir = os.path.dirname(os.path.abspath(__file__))
_root = os.path.dirname(_script_dir)
if _root not in sys.path:
    sys.path.insert(0, _root)

import model_cpu_gpt2 as _m

_m._fla_available = False

from model_cpu_gpt2 import CPUGPT, CPUGPTConfig

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

CHUNK_VARIANTS = [64, 256, 512, 1024]


def make_124m_config(gla_chunk: int) -> CPUGPTConfig:
    return CPUGPTConfig(
        n_layer=12,
        n_embd=768,
        n_head=12,
        ffn_hidden=2048,
        fno_modes=256,
        gla_chunk=gla_chunk,
        seq_len=8192,
        layer_pattern="SSSL",
        vocab_size=50257,
    )


polar_express_coeffs = [
    (8.156554524902461, -22.48329292557795, 15.878769915207462),
    (4.042929935166739, -2.808917465908714, 0.5000178451051316),
    (3.8916678022926607, -2.772484153217685, 0.5060648178503393),
    (3.285753657755655, -2.3681294933425376, 0.46449024233003106),
    (2.3465413258596377, -1.7097828382687081, 0.42323551169305323),
]


def adamw_step(p, grad, m, v, step, lr, b1, b2, eps, wd):
    p.mul_(1 - lr * wd)
    m.lerp_(grad, 1 - b1)
    v.lerp_(grad.square(), 1 - b2)
    bc1 = 1 - b1**step
    bc2 = 1 - b2**step
    p.addcdiv_(m / bc1, (v / bc2).sqrt_().add_(eps), value=-lr)


def muon_step(grads_stack, params, mom_buf, lr, momentum=0.95, ns_steps=3):
    mom_buf.lerp_(grads_stack, 1 - momentum)
    X = mom_buf.float()
    X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.02 + 1e-6)
    for a, b, c in polar_express_coeffs[:ns_steps]:
        if X.size(-2) >= X.size(-1):
            A = X.mT @ X
            X = a * X + X @ (b * A + c * (A @ A))
        else:
            A = X @ X.mT
            X = a * X + (b * A + c * (A @ A)) @ X
    torch._foreach_sub_(params, list((X * lr).to(params[0].dtype).unbind(0)))


class MuonAdamW(torch.optim.Optimizer):
    def __init__(self, param_groups):
        super().__init__(param_groups, defaults={})

    @torch.no_grad()
    def step(self):
        for g in self.param_groups:
            if g["kind"] == "adamw":
                for p in g["params"]:
                    if p.grad is None:
                        continue
                    st = self.state[p]
                    if not st:
                        st["step"] = 0
                        st["m"] = torch.zeros_like(p)
                        st["v"] = torch.zeros_like(p)
                    st["step"] += 1
                    adamw_step(
                        p,
                        p.grad,
                        st["m"],
                        st["v"],
                        st["step"],
                        g["lr"],
                        *g["betas"],
                        g["eps"],
                        g.get("wd", 0),
                    )
            elif g["kind"] == "muon":
                params = g["params"]
                if not params:
                    continue
                p0 = params[0]
                st = self.state[p0]
                stacked = torch.stack(
                    [p.grad for p in params if p.grad is not None]
                ).float()
                if not st:
                    st["mom"] = torch.zeros_like(stacked)
                lr = g["lr"] * max(1.0, p0.shape[-2] / max(p0.shape[-1], 1)) ** 0.5
                muon_step(stacked, params, st["mom"], lr, g.get("momentum", 0.95))


def build_optimizer(model, cfg, lr_matrix=0.01, lr_emb=0.02, lr_lm=0.002):
    scale = (cfg.n_embd / 768) ** -0.5
    raw = (
        model.module
        if isinstance(model, nn.parallel.DistributedDataParallel)
        else model
    )
    matrix_params, scalar_params = [], []
    for block in raw.blocks:
        for p in block.parameters():
            (matrix_params if p.ndim == 2 else scalar_params).append(p)
    groups = [
        dict(
            kind="adamw",
            params=list(raw.wte.parameters()),
            lr=lr_emb * scale,
            betas=(0.8, 0.95),
            eps=1e-8,
        ),
        dict(
            kind="adamw",
            params=list(raw.lm_head.parameters())
            if raw.lm_head.weight is not raw.wte.weight
            else [],
            lr=lr_lm * scale,
            betas=(0.8, 0.95),
            eps=1e-8,
        ),
        dict(
            kind="adamw",
            params=scalar_params,
            lr=lr_matrix * scale,
            betas=(0.8, 0.95),
            eps=1e-8,
        ),
    ]
    for shape in sorted({p.shape for p in matrix_params}):
        ps = [p for p in matrix_params if p.shape == shape]
        groups.append(dict(kind="muon", params=ps, lr=lr_matrix, momentum=0.95))
    opt = MuonAdamW(groups)
    for g in opt.param_groups:
        g["initial_lr"] = g["lr"]
    return opt


def lr_multiplier(progress, warmup=0.02, min_ratio=0.1):
    if progress < warmup:
        return progress / warmup
    t = (progress - warmup) / (1.0 - warmup)
    return min_ratio + (1.0 - min_ratio) * 0.5 * (1.0 + math.cos(math.pi * t))


def make_loader(data_dir, seq_len, device, batch_size=1):
    import glob
    import queue
    import threading

    shards = sorted(glob.glob(os.path.join(data_dir, "*.parquet")))
    if not shards:
        shards = sorted(glob.glob(os.path.join(data_dir, "*.bin")))
        if not shards:
            raise FileNotFoundError(f"No parquet or bin shards in {data_dir}")
        import numpy as np

        def load_fn(path):
            data = np.memmap(path, dtype=np.int32, mode="r")
            n = (len(data) // (seq_len + 1)) * (seq_len + 1)
            return torch.from_numpy(data[:n].copy()).long().view(-1, seq_len + 1)
    else:
        import pyarrow.parquet as pq

        def load_fn(path):
            tbl = pq.read_table(path, columns=["tokens"])
            tokens = torch.tensor(tbl["tokens"].to_pylist()[0], dtype=torch.long)
            n = (len(tokens) // (seq_len + 1)) * (seq_len + 1)
            return tokens[:n].view(-1, seq_len + 1)

    buf: queue.Queue = queue.Queue(maxsize=4)

    def _producer():
        ep = 1
        while True:
            for shard in shards:
                seqs = load_fn(shard)
                if len(seqs) == 0:
                    continue
                idx = torch.randperm(len(seqs))
                batch = []
                for i in idx:
                    batch.append(seqs[i])
                    if len(batch) == batch_size:
                        chunk = torch.stack(batch)
                        buf.put((chunk[:, :-1], chunk[:, 1:], ep))
                        batch = []
            ep += 1

    threading.Thread(target=_producer, daemon=True).start()
    while True:
        x, y, ep = buf.get()
        yield x.to(device), y.to(device), ep


def train_variant(
    chunk: int,
    args,
    out_dir: Path,
    device: str,
    is_master: bool,
    use_ddp: bool,
    local_rank: int,
):
    log.info(f"=== Training gla_chunk={chunk} ===")
    cfg = make_124m_config(chunk)
    model = CPUGPT(cfg).to(device)
    if use_ddp:
        model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
    opt = build_optimizer(model, cfg)
    if is_master:
        raw = model.module if use_ddp else model
        log.info(
            f"  params: {raw.param_count() / 1e6:.1f}M  seq_len={cfg.seq_len}  chunk={chunk}"
        )

    total_batch = args.device_batch * max(args.num_gpus, 1)
    total_steps = max(1, int(args.tokens / (total_batch * cfg.seq_len)))
    log.info(
        f"  steps={total_steps}  total_batch={total_batch}  tokens≈{total_steps * total_batch * cfg.seq_len / 1e6:.0f}M"
    )

    loader = make_loader(args.data_dir, cfg.seq_len, device, args.device_batch)
    ctx = (
        torch.autocast(device_type="cuda", dtype=torch.bfloat16)
        if "cuda" in device
        else torch.no_grad().__class__()
    )

    step = 0
    t0 = time.perf_counter()
    for x, y, _ in loader:
        if step >= total_steps:
            break

        model.train()
        progress = step / total_steps
        mult = lr_multiplier(progress)
        for g in opt.param_groups:
            g["lr"] = g["initial_lr"] * mult

        with torch.autocast(device_type=device.split(":")[0], dtype=torch.bfloat16):
            loss = model(x, y)

        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        opt.zero_grad(set_to_none=True)
        step += 1

        if is_master and step % 50 == 0:
            elapsed = time.perf_counter() - t0
            tps = step * total_batch * cfg.seq_len / elapsed
            log.info(
                f"  chunk={chunk} step={step}/{total_steps} ({100 * step / total_steps:.0f}%) "
                f"loss={loss.item():.4f} tok/s={tps:,.0f}"
            )

    ckpt_dir = out_dir / f"variant_{chunk}" / "ckpt"
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    if is_master:
        raw = model.module if use_ddp else model
        torch.save({"step": step, "model": raw.state_dict()}, ckpt_dir / "final.pt")
        log.info(f"  saved {ckpt_dir}/final.pt")

    return model, cfg


_HAY = (
    "The forest was quiet except for the occasional rustle of leaves in the breeze. "
    "Sunlight filtered through the canopy, casting dappled shadows on the mossy ground. "
    "A small stream wound its way between ancient oak trees, its water clear and cold. "
    "Birds called to one another across the branches, their songs filling the still air. "
    "Somewhere in the distance, a woodpecker drummed steadily against hollow bark. "
    "The smell of earth and pine needles rose from the path with each careful step. "
    "Nothing moved except the shadows and the water and the swaying tops of the tallest trees. "
) * 300


def eval_needle(model, cfg, device, n_trials=20, seed=42):
    import tiktoken

    enc = tiktoken.get_encoding("gpt2")
    hay_ids = enc.encode_ordinary(_HAY)
    rng = random.Random(seed)

    ctx_lens = [512, 1024, 2048, 4096, cfg.seq_len]
    depths = [0.1, 0.3, 0.5, 0.7, 0.9]

    raw = (
        model.module
        if isinstance(model, nn.parallel.DistributedDataParallel)
        else model
    )
    raw.eval()
    results = []

    for ctx_len in ctx_lens:
        for depth in depths:
            n_correct, n_valid = 0, 0
            for _ in range(n_trials):
                code = str(rng.randint(10000, 99999))
                dist_ = str(rng.randint(10000, 99999))
                while dist_ == code:
                    dist_ = str(rng.randint(10000, 99999))

                needle_ids = enc.encode_ordinary(
                    f" The passkey is {code}. Remember this passkey."
                )
                query_ids = enc.encode_ordinary(" The passkey is")
                code_ids = enc.encode_ordinary(" " + code)
                dist_ids = enc.encode_ordinary(" " + dist_)

                overhead = (
                    len(needle_ids) + len(query_ids) + max(len(code_ids), len(dist_ids))
                )
                hay_budget = ctx_len - overhead
                if hay_budget < 64:
                    continue

                prefix_len = int(hay_budget * depth)
                suffix_len = hay_budget - prefix_len
                start = rng.randint(0, max(0, len(hay_ids) - hay_budget - 1))
                hay_slice = hay_ids[start : start + hay_budget]
                context = (
                    hay_slice[:prefix_len]
                    + needle_ids
                    + hay_slice[prefix_len : prefix_len + suffix_len]
                    + query_ids
                )

                def avg_nll(answer_ids):
                    combined = context + answer_ids
                    remainder = len(combined) % cfg.gla_chunk
                    extra = (cfg.gla_chunk - remainder) % cfg.gla_chunk
                    full = combined + [0] * extra
                    inp = torch.tensor([full], dtype=torch.long, device=device)
                    with torch.no_grad():
                        logits = raw(inp)
                    s = len(context) - 1
                    lgt = logits[0, s : s + len(answer_ids)].float()
                    tgt = torch.tensor(answer_ids, dtype=torch.long, device=device)
                    return F.cross_entropy(lgt, tgt, reduction="mean").item()

                try:
                    nll_c = avg_nll(code_ids)
                    nll_d = avg_nll(dist_ids)
                    n_valid += 1
                    if nll_c < nll_d:
                        n_correct += 1
                except Exception:
                    pass

            acc = n_correct / n_valid if n_valid else float("nan")
            results.append(
                {
                    "ctx_len": ctx_len,
                    "depth": depth,
                    "accuracy": acc,
                    "n_correct": n_correct,
                    "n_valid": n_valid,
                }
            )
            log.info(
                f"  [needle] chunk={cfg.gla_chunk} ctx={ctx_len} depth={depth:.1f} "
                f"acc={acc:.0%} ({n_correct}/{n_valid})"
            )
    return results


def eval_throughput(model, cfg, device):
    raw = (
        model.module
        if isinstance(model, nn.parallel.DistributedDataParallel)
        else model
    )
    raw.eval()
    results = []
    ctx_lens = [512, 1024, 2048, 4096, cfg.seq_len]
    for ctx_len in ctx_lens:
        padded = math.ceil(ctx_len / cfg.gla_chunk) * cfg.gla_chunk
        x = torch.zeros(1, padded, dtype=torch.long, device=device)
        for _ in range(2):
            with torch.no_grad():
                raw(x)
        torch.cuda.synchronize() if "cuda" in device else None
        t0 = time.perf_counter()
        reps = 5
        for _ in range(reps):
            with torch.no_grad():
                raw(x)
        torch.cuda.synchronize() if "cuda" in device else None
        dt = (time.perf_counter() - t0) / reps
        tps = padded / dt
        results.append(
            {"ctx_len": ctx_len, "tok_per_sec": tps, "ms_per_seq": dt * 1000}
        )
        log.info(
            f"  [throughput] chunk={cfg.gla_chunk} ctx={ctx_len} → {tps:,.0f} tok/s"
        )
    return results


def eval_wikitext(model, cfg, device):
    import tiktoken

    try:
        from datasets import load_dataset

        try:
            ds = load_dataset(
                "Salesforce/wikitext",
                "wikitext-103-raw-v1",
                split="test",
                trust_remote_code=True,
            )
        except Exception:
            ds = load_dataset(
                "wikitext", "wikitext-103-raw-v1", split="test", trust_remote_code=True
            )
        text = "\n\n".join(ds["text"])
    except Exception:
        log.warning("  Could not load WikiText-103; skipping BPB eval")
        return {"bpb": float("nan")}

    enc = tiktoken.get_encoding("gpt2")
    tokens = torch.tensor(enc.encode(text[:4_000_000]), dtype=torch.long)
    tokens = tokens[:1_000_000]

    raw = (
        model.module
        if isinstance(model, nn.parallel.DistributedDataParallel)
        else model
    )
    raw.eval()

    seq_len = cfg.seq_len
    n_chunks = min(50, (len(tokens) - 1) // seq_len)
    total_nll, total_tok = 0.0, 0

    with torch.no_grad():
        for i in range(n_chunks):
            x = tokens[i * seq_len : (i + 1) * seq_len].unsqueeze(0).to(device)
            y = tokens[i * seq_len + 1 : (i + 1) * seq_len + 1].unsqueeze(0).to(device)
            if x.shape[1] < seq_len or y.shape[1] < seq_len:
                break
            pad = cfg.gla_chunk - (x.shape[1] % cfg.gla_chunk or cfg.gla_chunk)
            if pad < cfg.gla_chunk:
                x = F.pad(x, (0, pad))
            logits = raw(x)
            lgt = logits[0, :seq_len].float()
            loss = F.cross_entropy(lgt, y[0], reduction="sum")
            total_nll += loss.item()
            total_tok += seq_len

    nats = total_nll / max(total_tok, 1)
    bpb = nats / math.log(2)
    log.info(f"  [wikitext] chunk={cfg.gla_chunk} BPB={bpb:.4f}")
    return {"bpb": bpb, "nats": nats, "tokens_evaluated": total_tok}


def make_figures(all_results, out_dir: Path):
    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import numpy as np
    except ImportError:
        log.warning("matplotlib not available; skipping figures")
        return

    chunks = sorted(all_results.keys())
    colors = ["#E84855", "#F4A261", "#2E86AB", "#264653"]

    fig, axes = plt.subplots(2, 2, figsize=(14, 9), sharex=True, sharey=True)
    axes = axes.flatten()
    for idx, chunk in enumerate(chunks):
        rows = all_results[chunk]["needle"]
        ctx_u = sorted(set(r["ctx_len"] for r in rows))
        dep_u = sorted(set(r["depth"] for r in rows))
        grid = np.full((len(dep_u), len(ctx_u)), float("nan"))
        for r in rows:
            ci = ctx_u.index(r["ctx_len"])
            di = dep_u.index(r["depth"])
            if not (isinstance(r["accuracy"], float) and math.isnan(r["accuracy"])):
                grid[di, ci] = r["accuracy"]
        ax = axes[idx]
        im = ax.imshow(grid, vmin=0, vmax=1, cmap="RdYlGn", aspect="auto")
        ax.set_title(f"gla_chunk={chunk}", fontsize=12, fontweight="bold")
        ax.set_xticks(range(len(ctx_u)))
        ax.set_xticklabels(
            [f"{c // 1024}K" if c >= 1024 else str(c) for c in ctx_u], fontsize=9
        )
        ax.set_yticks(range(len(dep_u)))
        ax.set_yticklabels([f"{int(d * 100)}%" for d in dep_u], fontsize=9)
        ax.set_xlabel("Context Length", fontsize=10)
        ax.set_ylabel("Needle Depth", fontsize=10)
        for di in range(len(dep_u)):
            for ci in range(len(ctx_u)):
                v = grid[di, ci]
                if not math.isnan(v):
                    col = "black" if 0.25 < v < 0.75 else "white"
                    ax.text(
                        ci,
                        di,
                        f"{v:.0%}",
                        ha="center",
                        va="center",
                        fontsize=9,
                        color=col,
                        fontweight="bold",
                    )
        plt.colorbar(im, ax=ax)
    fig.suptitle(
        "Passkey Retrieval Accuracy by GLA Chunk Size\n(124M FELA, 1B tokens)",
        fontsize=14,
    )
    fig.tight_layout()
    fig.savefig(out_dir / "fig_chunk_retrieval.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    log.info("  saved fig_chunk_retrieval.png")

    fig, ax = plt.subplots(figsize=(8, 5))
    for chunk, color in zip(chunks, colors):
        rows = all_results[chunk].get("throughput", [])
        if not rows:
            continue
        xs = [r["ctx_len"] for r in rows]
        ys = [r["tok_per_sec"] for r in rows]
        ax.plot(
            xs,
            ys,
            "o-",
            color=color,
            linewidth=2.5,
            markersize=6,
            label=f"chunk={chunk}",
        )
    ax.set_xscale("log", base=2)
    ax.set_xlabel("Context Length (tokens)", fontsize=12)
    ax.set_ylabel("Throughput (tok/s)", fontsize=12)
    ax.set_title(
        "Throughput vs. Context Length\n(124M FELA — chunk size effect)", fontsize=13
    )
    ax.legend(fontsize=11)
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    fig.savefig(out_dir / "fig_chunk_throughput.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    log.info("  saved fig_chunk_throughput.png")

    fig, axes = plt.subplots(1, 2, figsize=(11, 4))
    acc_hard, bpb_vals = [], []
    for chunk in chunks:
        needle_rows = all_results[chunk]["needle"]
        match = [
            r
            for r in needle_rows
            if r["ctx_len"] == 4096 and abs(r["depth"] - 0.1) < 0.01
        ]
        acc_hard.append(match[0]["accuracy"] if match else float("nan"))
        bpb_vals.append(all_results[chunk].get("wikitext", {}).get("bpb", float("nan")))

    c_labels = [f"chunk={c}" for c in chunks]
    ax1 = axes[0]
    bars = ax1.bar(c_labels, acc_hard, color=colors)
    ax1.axhline(0.5, color="gray", linestyle="--", alpha=0.7, label="random")
    ax1.set_ylim(0, 1)
    ax1.set_ylabel("Retrieval Accuracy")
    ax1.set_title("Hard case: ctx=4K, depth=10%\n(needle at very start)")
    ax1.legend(fontsize=9)
    for bar, v in zip(bars, acc_hard):
        if not math.isnan(v):
            ax1.text(
                bar.get_x() + bar.get_width() / 2,
                v + 0.02,
                f"{v:.0%}",
                ha="center",
                fontsize=10,
                fontweight="bold",
            )

    ax2 = axes[1]
    bars2 = ax2.bar(c_labels, bpb_vals, color=colors)
    ax2.set_ylabel("WikiText BPB (lower = better)")
    ax2.set_title("Language quality\n(WikiText-103 BPB)")
    for bar, v in zip(bars2, bpb_vals):
        if not math.isnan(v):
            ax2.text(
                bar.get_x() + bar.get_width() / 2,
                v + 0.005,
                f"{v:.3f}",
                ha="center",
                fontsize=10,
                fontweight="bold",
            )

    fig.suptitle("Chunk Size Tradeoff: Retrieval vs. Quality", fontsize=13)
    fig.tight_layout()
    fig.savefig(out_dir / "fig_chunk_summary.png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    log.info("  saved fig_chunk_summary.png")


def main():
    parser = argparse.ArgumentParser(description="GLA chunk-size ablation")
    parser.add_argument("--data-dir", default="/data/openwebtext")
    parser.add_argument(
        "--tokens",
        type=float,
        default=1e9,
        help="training tokens per variant (default 1B)",
    )
    parser.add_argument(
        "--device-batch", type=int, default=4, help="per-device batch size (sequences)"
    )
    parser.add_argument(
        "--num-gpus", type=int, default=0, help="override GPU count (0 = auto-detect)"
    )
    parser.add_argument("--out-dir", default="results/chunk_ablation")
    parser.add_argument("--needle-trials", type=int, default=20)
    parser.add_argument(
        "--smoke",
        action="store_true",
        help="fast smoke test: 5M tokens, 3 needle trials",
    )
    parser.add_argument(
        "--eval-only",
        type=str,
        default=None,
        help="path to existing results dir — skip training, just eval+figures",
    )
    parser.add_argument(
        "--chunks",
        type=int,
        nargs="+",
        default=CHUNK_VARIANTS,
        help="which chunk sizes to run (default: 64 256 512 1024)",
    )
    args = parser.parse_args()

    if args.smoke:
        args.tokens = 5e6
        args.needle_trials = 3

    use_ddp = False
    local_rank, global_rank = 0, 0
    if "RANK" in os.environ:
        dist.init_process_group(backend="nccl")
        local_rank = dist.get_rank() % torch.cuda.device_count()
        global_rank = dist.get_rank()
        use_ddp = True
        torch.cuda.set_device(local_rank)

    num_gpus = args.num_gpus or (
        torch.cuda.device_count() if torch.cuda.is_available() else 0
    )
    args.num_gpus = max(1, num_gpus)
    device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu"
    is_master = global_rank == 0

    out_dir = Path(args.out_dir)
    if is_master:
        out_dir.mkdir(parents=True, exist_ok=True)

    all_results = {}

    for chunk in args.chunks:
        variant_dir = out_dir / f"variant_{chunk}"
        if is_master:
            variant_dir.mkdir(parents=True, exist_ok=True)

        model, cfg = train_variant(
            chunk, args, out_dir, device, is_master, use_ddp, local_rank
        )

        if use_ddp:
            dist.barrier()

        if is_master:
            raw = model.module if use_ddp else model
            raw.eval()

            log.info(f"--- Evaluating chunk={chunk} ---")
            needle_res = eval_needle(raw, cfg, device, n_trials=args.needle_trials)
            (variant_dir / "needle.json").write_text(json.dumps(needle_res, indent=2))

            tput_res = eval_throughput(raw, cfg, device)
            (variant_dir / "throughput.json").write_text(json.dumps(tput_res, indent=2))

            wiki_res = eval_wikitext(raw, cfg, device)
            (variant_dir / "wikitext.json").write_text(json.dumps(wiki_res, indent=2))

            all_results[chunk] = {
                "needle": needle_res,
                "throughput": tput_res,
                "wikitext": wiki_res,
                "params_M": raw.param_count() / 1e6,
                "gla_chunk": chunk,
            }

        del model
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        if use_ddp:
            dist.barrier()

    if is_master:
        (out_dir / "summary.json").write_text(json.dumps(all_results, indent=2))
        log.info(f"Summary written to {out_dir}/summary.json")

        make_figures(all_results, out_dir)

        print("\n" + "=" * 65)
        print("CHUNK ABLATION SUMMARY")
        print("=" * 65)
        print(
            f"{'chunk':>8}  {'4K/10% acc':>10}  {'4K/50% acc':>10}  {'4K/90% acc':>10}  {'BPB':>7}"
        )
        print("-" * 65)
        for chunk in sorted(all_results.keys()):
            nr = all_results[chunk]["needle"]

            def get_acc(ctx, dep):
                m = [
                    r
                    for r in nr
                    if r["ctx_len"] == ctx and abs(r["depth"] - dep) < 0.01
                ]
                return m[0]["accuracy"] if m else float("nan")

            bpb = all_results[chunk].get("wikitext", {}).get("bpb", float("nan"))
            print(
                f"{chunk:>8}  {get_acc(4096, 0.1):>10.0%}  {get_acc(4096, 0.5):>10.0%}  "
                f"{get_acc(4096, 0.9):>10.0%}  {bpb:>7.4f}"
            )
        print("=" * 65)

    if use_ddp:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()