#!/usr/bin/env python3
"""BP-install supervised training on arithmetic.

Adds per-block probe-loss to the standard NTP training:
  total = NTP_loss + λ_install · sum_V CE(probe_V(pre_blk_target_V), V_labels)

For arithmetic, BP variables are (carry_after_k, sum_digit_k) at known
positions. Target blocks default to Phase-A install pattern:
  - carry_after_k -> block 2
  - sum_digit_k   -> block 3

The supervised loss pushes the model to install each BP variable at its
target block's INPUT (i.e., the residual just before block target_block).

Hypothesis (from user): forcing single-block install via supervision
will concentrate the per-head distributed signal from Phase A into a
cleaner BP message.
"""
from __future__ import annotations

import argparse
import json
import math
import random
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple

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

sys.path.insert(0, str(Path(__file__).resolve().parent))
from coppola_pretrain_tiny import (
    ByteGPT, GPTConfig,
    evaluate_arithmetic_accuracy,
    bits_per_byte,
)


# ----------------------------- BP annotations ------------------------------

def _make_addition_with_annotations(max_digits: int,
                                      surface: str = "reverse"
                                      ) -> Tuple[str, List[Dict[str, int]]]:
    """Generate an addition problem in one of two surfaces, with per-position
    BP-variable annotations placed at the NTP predicting residual.

    Surfaces:
      - "reverse": `a+b=R<ans_lsb_first>\n`, e.g. 12+34 -> "12+34=R64\n"
      - "forward": `a+b=F<ans_msb_first>\n`, e.g. 12+34 -> "12+34=F46\n"

    Annotation IDENTITY (carry_after_k, sum_digit_k) is invariant across
    surfaces; only the POSITION at which the predicting residual sits
    differs. This is the test substrate for surface-form invariance of
    BP-install location.
    """
    if surface not in ("reverse", "forward"):
        raise ValueError(f"surface must be 'reverse' or 'forward', got {surface!r}")
    n_a = random.randint(1, max_digits)
    n_b = random.randint(1, max_digits)
    a = random.randint(0, 10 ** n_a - 1)
    b = random.randint(0, 10 ** n_b - 1)
    ans_int = a + b
    ans_forward = str(ans_int)
    ans_emit = ans_forward[::-1] if surface == "reverse" else ans_forward
    marker = "R" if surface == "reverse" else "F"
    text = f"{a}+{b}={marker}{ans_emit}\n"
    a_str = str(a)
    b_str = str(b)
    plus_pos = len(a_str)
    eq_pos = plus_pos + 1 + len(b_str)
    marker_pos = eq_pos + 1
    ans_start = marker_pos + 1
    ans_len = len(ans_emit)
    ann = [{} for _ in range(len(text))]

    # Compute carries: carry_after_k = carry out of column k (binary).
    max_n = max(n_a, n_b)
    carry = 0
    carries = []
    for k in range(max_n + 1):
        da = (a // (10 ** k)) % 10 if k < n_a else 0
        db = (b // (10 ** k)) % 10 if k < n_b else 0
        tot = da + db + carry
        carry = tot // 10
        carries.append(carry)

    # Sum digits and carries placed at predicting residuals (NTP-aligned).
    # For position-of-sum_digit_k's EMISSION:
    #   reverse: sum_digit_k emitted at ans_start + k
    #   forward: sum_digit_k emitted at ans_start + (ans_len - 1 - k)
    # The annotation goes at emission_pos - 1.
    # carry_after_k is needed to PREDICT sum_digit_(k+1), so it sits at the
    # residual that PRECEDES sum_digit_(k+1)'s emission position.
    def emission_pos_for_sum_k(k: int) -> int:
        # k=0 is LSB; emission position depends on surface
        if surface == "reverse":
            return ans_start + k
        # forward: most-significant emitted first
        # max emitted digit index = ans_len - 1 corresponds to LSB if forward
        # The kth LSB digit is the (ans_len - 1 - k)-th emitted in forward
        return ans_start + (ans_len - 1 - k)

    for k in range(ans_len):
        emit = emission_pos_for_sum_k(k)
        tgt = emit - 1
        if 0 <= tgt < len(text):
            # Extract the actual digit value at that emission position
            digit_char = text[emit]
            if digit_char.isdigit():
                ann[tgt][f"sum_digit_{k}"] = int(digit_char)

    # carry_after_k is needed to predict sum_digit_(k+1).
    for k, c in enumerate(carries):
        if k + 1 >= ans_len:
            continue  # no next sum digit to predict
        emit_next = emission_pos_for_sum_k(k + 1)
        tgt = emit_next - 1
        if 0 <= tgt < len(text):
            ann[tgt][f"carry_after_{k}"] = c

    return text, ann


# Variables we supervise and their target blocks.
# Per Phase A: carries install around block 2, sums around block 3.
# For 4L model: blocks 0..3.
def default_var_specs(max_digits: int, target_carry_block: int = 2,
                        target_sum_block: int = 3
                        ) -> List[Tuple[str, int, int]]:
    """Return list of (var_name, n_classes, target_block_idx)."""
    specs = []
    # carry_after_k: 2-class (0 or 1) — for digits beyond highest column,
    # carry could exceed 1 only in pathological cases, so binary is safe
    # for 1-3 digit ops with values 0..999+0..999=0..1998.
    for k in range(max_digits + 1):
        specs.append((f"carry_after_{k}", 2, target_carry_block))
    for k in range(max_digits + 1):
        specs.append((f"sum_digit_{k}", 10, target_sum_block))
    return specs


# ----------------------------- Corpus --------------------------------------

class ArithmeticBPCorpus:
    """Arithmetic corpus emitting (x, y, bp_labels)."""

    def __init__(self, max_digits: int, var_specs, surfaces=("reverse",),
                 surface_weights=None):
        """surfaces: tuple of surface names to sample from.
        surface_weights: matching list of relative weights (default uniform)."""
        self.max_digits = max_digits
        self.var_specs = var_specs
        self.var_names = [name for name, _, _ in var_specs]
        self.surfaces = tuple(surfaces)
        if surface_weights is None:
            surface_weights = [1.0] * len(self.surfaces)
        if len(surface_weights) != len(self.surfaces):
            raise ValueError("surface_weights length mismatch")
        total = sum(surface_weights)
        self._cum_weights = []
        run = 0.0
        for w in surface_weights:
            run += w / total
            self._cum_weights.append(run)

    def _pick_surface(self) -> str:
        r = random.random()
        for s, c in zip(self.surfaces, self._cum_weights):
            if r <= c:
                return s
        return self.surfaces[-1]

    def sample_batch(self, batch_size: int, seq_len: int, device: str):
        need = seq_len + 1
        x_batch = torch.zeros((batch_size, seq_len), dtype=torch.long)
        y_batch = torch.zeros((batch_size, seq_len), dtype=torch.long)
        labels = {v: torch.full((batch_size, seq_len), -1, dtype=torch.long)
                  for v in self.var_names}
        for b in range(batch_size):
            buf_bytes: List[int] = []
            buf_ann: List[Dict[str, int]] = []
            while len(buf_bytes) < need:
                surface = self._pick_surface()
                text, ann = _make_addition_with_annotations(self.max_digits,
                                                              surface=surface)
                buf_bytes.extend(text.encode("ascii"))
                buf_ann.extend(ann)
            ids = buf_bytes[:need]
            x_batch[b] = torch.tensor(ids[:-1], dtype=torch.long)
            y_batch[b] = torch.tensor(ids[1:], dtype=torch.long)
            for t in range(seq_len):
                if t < len(buf_ann):
                    for v, val in buf_ann[t].items():
                        if v in labels:
                            labels[v][b, t] = int(val)
        return (
            x_batch.to(device),
            y_batch.to(device),
            {v: t.to(device) for v, t in labels.items()},
        )


# ----------------------------- Probes --------------------------------------

class BPProbeHeads(nn.Module):
    """One linear probe per BP variable. Probes are applied to the residual
    at the variable's target block. Optionally also at other blocks
    (anti-spread)."""

    def __init__(self, d_model: int, var_specs):
        super().__init__()
        self.var_specs = var_specs
        self.probes = nn.ModuleDict({
            name: nn.Linear(d_model, n_classes)
            for name, n_classes, _ in var_specs
        })

    def loss(self, residuals_per_block: List[torch.Tensor],
             labels: Dict[str, torch.Tensor],
             anti_lambda: float = 0.0) -> Tuple[torch.Tensor, Dict[str, float]]:
        """Install loss: CE at target block. Anti-spread (optional): for
        other blocks, push the probe to NOT extract V (negative log-likelihood
        of UNIFORM distribution = encourage low confidence)."""
        total = torch.zeros((), device=residuals_per_block[0].device)
        details = {}
        for name, n_classes, target_block in self.var_specs:
            res = residuals_per_block[target_block]
            logits = self.probes[name](res)  # [B, T, n_classes]
            y = labels[name]                  # [B, T]
            mask = y >= 0
            if int(mask.sum()) == 0:
                continue
            ce = F.cross_entropy(
                logits[mask],
                y[mask],
                reduction="mean",
            )
            total = total + ce
            details[f"L_{name}_at_blk{target_block}"] = float(ce.item())

            if anti_lambda > 0.0:
                anti_terms = []
                for b_idx, res_b in enumerate(residuals_per_block):
                    if b_idx == target_block:
                        continue
                    other_logits = self.probes[name](res_b)
                    # Want predicted distribution near uniform -> minimize KL
                    # to uniform = maximize entropy of softmax.
                    logp = F.log_softmax(other_logits, dim=-1)
                    p = logp.exp()
                    # Negative entropy of p (we want to MAXIMIZE entropy ->
                    # MINIMIZE negative entropy).
                    neg_ent = (p * logp).sum(dim=-1)
                    anti_terms.append(neg_ent[mask].mean())
                if anti_terms:
                    anti = torch.stack(anti_terms).mean()
                    total = total + anti_lambda * anti
                    details[f"anti_{name}"] = float(anti.item())
        return total, details

    @torch.no_grad()
    def probe_accs(self, residuals_per_block: List[torch.Tensor],
                    labels: Dict[str, torch.Tensor]) -> Dict[str, Dict[int, float]]:
        out: Dict[str, Dict[int, float]] = {}
        for name, n_classes, _ in self.var_specs:
            y = labels[name]
            mask = y >= 0
            if int(mask.sum()) == 0:
                continue
            out[name] = {}
            for b_idx, res in enumerate(residuals_per_block):
                logits = self.probes[name](res)
                pred = logits.argmax(dim=-1)
                acc = float((pred[mask] == y[mask]).float().mean())
                out[name][b_idx] = acc
        return out


# ----------------------------- Forward helper ------------------------------

def head_entropy_loss(model: ByteGPT) -> torch.Tensor:
    """Penalize entropy of head-norm distribution per row of c_proj.

    For each block, c_proj has shape [d_model, n_head * head_dim]. Reshape to
    [d_model, n_head, head_dim] and compute per-row squared norms across heads.
    Normalize per-row to get a head-attribution distribution. Penalize its
    entropy — minimum (0) means each row written by ONE head.
    """
    total = 0.0
    n_blocks = 0
    eps = 1e-9
    n_head = model.config.n_head
    for blk in model.transformer.h:
        if not hasattr(blk.attn, "c_proj"):
            continue
        W = blk.attn.c_proj.weight
        d_model = W.shape[0]
        head_dim = W.shape[1] // n_head
        if head_dim * n_head != W.shape[1]:
            continue
        blocks = W.view(d_model, n_head, head_dim)
        head_norms_sq = (blocks ** 2).sum(dim=-1)  # [d_model, n_head]
        p = head_norms_sq / (head_norms_sq.sum(dim=-1, keepdim=True) + eps)
        ent = -(p * (p + eps).log()).sum(dim=-1)  # [d_model]
        total = total + ent.mean()
        n_blocks += 1
    if n_blocks == 0:
        return torch.zeros((), device=W.device)
    return total / n_blocks


@torch.no_grad()
def evaluate_mixed_arithmetic(model: ByteGPT, device: str, n_problems: int,
                                max_digits: int, surfaces=("reverse", "forward"),
                                seed: int = 42) -> Dict[str, float]:
    """Greedy-decode `a+b=<marker>` and compare against ground truth for each
    surface. Returns per-surface accuracies and a combined number."""
    random.seed(seed)
    torch.manual_seed(seed)
    nl_id = ord("\n")
    results: Dict[str, Dict[str, list]] = {s: {"correct": [], "by_len": {}}
                                              for s in surfaces}
    per_surface_n = max(1, n_problems // len(surfaces))
    for surface in surfaces:
        marker = "R" if surface == "reverse" else "F"
        for _ in range(per_surface_n):
            n_a = random.randint(1, max_digits)
            n_b = random.randint(1, max_digits)
            a = random.randint(0, 10 ** n_a - 1)
            b = random.randint(0, 10 ** n_b - 1)
            ans_int = a + b
            ans_forward = str(ans_int)
            true_emit = ans_forward[::-1] if surface == "reverse" else ans_forward
            prompt = f"{a}+{b}={marker}"
            ids = torch.tensor([list(prompt.encode("ascii"))],
                                 dtype=torch.long, device=device)
            max_new = len(true_emit) + 2
            gen = bytearray()
            for _ in range(max_new):
                logits, _ = model(ids)
                nxt = int(logits[0, -1].argmax())
                if nxt == nl_id:
                    break
                gen.append(nxt)
                ids = torch.cat([ids, torch.tensor([[nxt]], device=device)],
                                  dim=1)
            try:
                got = gen.decode("ascii", errors="ignore")
            except Exception:
                got = ""
            is_correct = (got == true_emit)
            results[surface]["correct"].append(1 if is_correct else 0)
            n_digits = max(n_a, n_b)
            bucket = results[surface]["by_len"].setdefault(n_digits, [0, 0])
            bucket[1] += 1
            if is_correct:
                bucket[0] += 1
    out = {}
    total_correct, total_n = 0, 0
    for s in surfaces:
        c = sum(results[s]["correct"])
        n = max(1, len(results[s]["correct"]))
        out[f"acc_{s}"] = c / n
        for k, (cc, tt) in sorted(results[s]["by_len"].items()):
            out[f"acc_{s}_{k}d"] = cc / max(1, tt)
        total_correct += c
        total_n += n
    out["accuracy"] = total_correct / max(1, total_n)
    return out


def forward_with_block_captures(model: ByteGPT, x: torch.Tensor, y: torch.Tensor):
    """Run model and capture pre-block residuals + post-ln_f.
    Returns (logits, lm_loss, residuals_per_block).
    residuals_per_block[i] is the input to block i (= output of block i-1).
    """
    captures: List[torch.Tensor] = [None] * model.config.n_layer
    handles = []
    for i, blk in enumerate(model.transformer.h):
        def make_hook(idx):
            def hook(module, inputs):
                captures[idx] = inputs[0]
            return hook
        handles.append(blk.register_forward_pre_hook(make_hook(i)))
    try:
        logits, lm_loss = model(x, y)
    finally:
        for h in handles:
            h.remove()
    return logits, lm_loss, captures


# ----------------------------- Training ------------------------------------

@dataclass
class TrainConfig:
    n_layer: int = 4
    n_head: int = 4
    n_embd: int = 128
    mlp_mult: int = 4
    mlp_activation: str = "silu"
    seq_len: int = 128
    attn_kind: str = "mha"
    attn_norm: str = "entmax15"
    position_encoding: str = "alibi"
    loss: str = "bce"
    logit_rmsnorm_scale: float = 8.0
    sparsity_lambda: float = 1e-3
    # A4: Attention Residuals depth-router + entropy pressure, composed
    # with the BP-install probe-loss. The entropy term is added to
    # lm_loss inside ByteGPT.forward; install-loss supplies the correct
    # support. See docs/attnres-bp-research-program.md A2′-result / A4.
    res_attn: str = "none"
    res_attn_norm: str = "softmax"
    res_attn_temp: float = 1.0
    res_attn_entropy_lambda: float = 0.0


def build_model(tc: TrainConfig, device: str) -> ByteGPT:
    cfg = GPTConfig(
        vocab_size=256,
        n_layer=tc.n_layer,
        n_head=tc.n_head,
        n_embd=tc.n_embd,
        seq_len=tc.seq_len,
        mlp_mult=tc.mlp_mult,
        mlp_activation=tc.mlp_activation,
        attn_kind=tc.attn_kind,
        attn_norm=tc.attn_norm,
        position_encoding=tc.position_encoding,
        loss_kind=tc.loss,
        logit_rmsnorm_scale=tc.logit_rmsnorm_scale,
        sparsity_lambda=tc.sparsity_lambda,
        res_attn=tc.res_attn,
        res_attn_norm=tc.res_attn_norm,
        res_attn_temp=tc.res_attn_temp,
        res_attn_entropy_lambda=tc.res_attn_entropy_lambda,
    )
    return ByteGPT(cfg).to(device)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--steps", type=int, default=8000)
    ap.add_argument("--batch-size", type=int, default=64)
    ap.add_argument("--seq-len", type=int, default=128)
    ap.add_argument("--lr", type=float, default=3e-3)
    ap.add_argument("--lr-min", type=float, default=3e-4)
    ap.add_argument("--warmup-steps", type=int, default=500)
    ap.add_argument("--weight-decay", type=float, default=0.01)
    ap.add_argument("--max-digits", type=int, default=3)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--surfaces", nargs="+", default=["reverse"],
                    choices=["reverse", "forward"],
                    help="Surface forms to sample (uniform).")
    ap.add_argument("--eval-every", type=int, default=1000)
    ap.add_argument("--n-eval-arith", type=int, default=200)
    ap.add_argument("--target-carry-block", type=int, default=2)
    ap.add_argument("--target-sum-block", type=int, default=3)
    ap.add_argument("--install-lambda", type=float, default=1.0,
                    help="Coefficient for BP-install probe loss.")
    ap.add_argument("--anti-lambda", type=float, default=0.0,
                    help="Coefficient for anti-spread loss (push non-target "
                         "block probes to uniform).")
    ap.add_argument("--head-entropy-lambda", type=float, default=0.0,
                    help="Coefficient for head specialization loss: entropy of "
                         "head-norm distribution across rows of c_proj. "
                         "Minimize -> each row written by one head.")
    ap.add_argument("--res-attn", choices=["none", "full"], default="none",
                    help="A4: 'full' = block-granular Attention Residuals "
                         "(depth-router) composed with BP-install supervision.")
    ap.add_argument("--res-attn-norm",
                    choices=["softmax", "entmax15", "sparsemax"],
                    default="softmax",
                    help="Depth-attention normalization (res-attn only).")
    ap.add_argument("--res-attn-temp", type=float, default=1.0,
                    help="A2′ fixed depth-score temperature (res-attn only).")
    ap.add_argument("--res-attn-entropy-lambda", type=float, default=0.0,
                    help="A2′ entropy penalty on α_depth (res-attn only); "
                         "supplies sparsity, install-loss supplies support.")
    ap.add_argument("--n-layer", type=int, default=4)
    ap.add_argument("--n-head", type=int, default=4)
    ap.add_argument("--n-embd", type=int, default=128)
    ap.add_argument("--checkpoint-dir", type=Path, required=True)
    ap.add_argument("--output", type=Path, required=True)
    ap.add_argument("--save-every", type=int, default=2000)
    args = ap.parse_args()

    random.seed(args.seed)
    torch.manual_seed(args.seed)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    tc = TrainConfig(seq_len=args.seq_len, n_layer=args.n_layer,
                     n_head=args.n_head, n_embd=args.n_embd,
                     res_attn=args.res_attn, res_attn_norm=args.res_attn_norm,
                     res_attn_temp=args.res_attn_temp,
                     res_attn_entropy_lambda=args.res_attn_entropy_lambda)
    model = build_model(tc, device)
    if tc.res_attn != "none":
        print(f"# AttnRes: {tc.res_attn}/{tc.res_attn_norm} "
              f"temp={tc.res_attn_temp} entropy_lambda={tc.res_attn_entropy_lambda}")

    var_specs = default_var_specs(args.max_digits, args.target_carry_block,
                                     args.target_sum_block)
    probes = BPProbeHeads(tc.n_embd, var_specs).to(device)
    corpus = ArithmeticBPCorpus(args.max_digits, var_specs=var_specs,
                                  surfaces=tuple(args.surfaces))

    # Optimizer: include both model and probe params
    all_params = list(model.parameters()) + list(probes.parameters())
    optimizer = torch.optim.AdamW(all_params, lr=args.lr,
                                   weight_decay=args.weight_decay,
                                   betas=(0.9, 0.95))

    def cur_lr(step):
        if step < args.warmup_steps:
            return args.lr * step / max(1, args.warmup_steps)
        progress = (step - args.warmup_steps) / max(1, args.steps - args.warmup_steps)
        progress = min(1.0, max(0.0, progress))
        return args.lr_min + 0.5 * (args.lr - args.lr_min) * (1 + math.cos(math.pi * progress))

    args.checkpoint_dir.mkdir(parents=True, exist_ok=True)
    history: List[Dict] = []

    print(f"# device={device} n_layer={tc.n_layer} n_head={tc.n_head} n_embd={tc.n_embd}")
    print(f"# BP supervision: λ_install={args.install_lambda} λ_anti={args.anti_lambda}")
    print(f"# Targets: carry→blk{args.target_carry_block}, sum→blk{args.target_sum_block}")
    print(f"# Variables supervised: {[s[0] for s in var_specs]}")

    last_train_loss = None
    last_lm_loss = None
    last_install_loss = None
    for step in range(1, args.steps + 1):
        model.train()
        probes.train()
        lr = cur_lr(step)
        for pg in optimizer.param_groups:
            pg["lr"] = lr

        x, y, bp_labels = corpus.sample_batch(args.batch_size, args.seq_len, device)
        optimizer.zero_grad(set_to_none=True)
        _, lm_loss, captures = forward_with_block_captures(model, x, y)
        install_loss, details = probes.loss(captures, bp_labels,
                                              anti_lambda=args.anti_lambda)
        total = lm_loss + args.install_lambda * install_loss
        if args.head_entropy_lambda > 0:
            head_ent = head_entropy_loss(model)
            total = total + args.head_entropy_lambda * head_ent
            last_head_ent = float(head_ent.item())
        else:
            last_head_ent = 0.0
        total.backward()
        torch.nn.utils.clip_grad_norm_(all_params, 1.0)
        optimizer.step()
        last_train_loss = float(total.item())
        last_lm_loss = float(lm_loss.item())
        last_install_loss = float(install_loss.item())

        if step == 1 or step % args.eval_every == 0 or step == args.steps:
            model.eval()
            probes.eval()
            # Compute val loss + bpb on fresh batch
            with torch.no_grad():
                vx, vy, vbp = corpus.sample_batch(args.batch_size, args.seq_len, device)
                vlogits, vlm_loss, vcaps = forward_with_block_captures(model, vx, vy)
                vce = F.cross_entropy(vlogits.view(-1, vlogits.size(-1)),
                                       vy.reshape(-1))
                vbpb = float(bits_per_byte(vce))
                # Probe accuracies across all blocks (diagnostic)
                accs = probes.probe_accs(vcaps, vbp)
            arith = evaluate_mixed_arithmetic(
                model, device, n_problems=args.n_eval_arith,
                max_digits=args.max_digits,
                surfaces=tuple(args.surfaces),
            )
            entry = {
                "step": step,
                "lr": lr,
                "train_loss": last_train_loss,
                "lm_loss": last_lm_loss,
                "install_loss": last_install_loss,
                "head_ent": last_head_ent,
                "val_lm": float(vlm_loss.item()),
                "val_ce": float(vce.item()),
                "val_bpb": vbpb,
                "arith_acc": arith["accuracy"],
                "probe_accs": accs,
            }
            history.append(entry)
            line = (f"step={step:>5d} lr={lr:.2e} "
                    f"lm={last_lm_loss:.3f} install={last_install_loss:.3f} "
                    f"head_ent={last_head_ent:.3f} "
                    f"val_bpb={vbpb:.3f} arith_acc={arith['accuracy']:.3f}")
            print(line)
            # Compact probe summary: install variable at its target block + spread
            for name, n_cls, tgt in var_specs:
                if name in accs:
                    row = accs[name]
                    target_acc = row.get(tgt, float("nan"))
                    others = [v for k, v in row.items() if k != tgt]
                    spread = (sum(others) / len(others)) if others else 0.0
                    if step == args.steps or step % (args.eval_every * 4) == 1 or step == 1:
                        print(f"    {name}: blk{tgt}={target_acc:.3f}  "
                              f"avg-other={spread:.3f}  delta={target_acc-spread:+.3f}")

        if (args.save_every and step % args.save_every == 0) or step == args.steps:
            ck = args.checkpoint_dir / f"bp_sup_step_{step}.pt"
            torch.save({
                "model_state": model.state_dict(),
                "model_config": vars(model.config),
                "probe_state": probes.state_dict(),
                "var_specs": var_specs,
                "step": step,
            }, ck)
            torch.save({"model_state": model.state_dict(),
                          "model_config": vars(model.config),
                          "probe_state": probes.state_dict(),
                          "var_specs": var_specs,
                          "step": step}, args.checkpoint_dir / "bp_sup_latest.pt")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps({
        "config": vars(args) | {"checkpoint_dir": str(args.checkpoint_dir), "output": str(args.output)},
        "var_specs": var_specs,
        "history": history,
    }, indent=2, default=str))


if __name__ == "__main__":
    main()