scripts/bootstrap_mu.py

"""Bootstrap CIs on μ (and on Δμ = μ_FT − μ_base) by resampling aligne's
observed edges with replacement and re-fitting Thurstone Case V each time.

This is the right way to get uncertainty on the per-item μ. The Thurstone σ
is a gauge parameter (pinned to 1.0); the residual sd from a μ_FT-on-μ_base
regression conflates per-item estimation error with tile-specific signal
and is *not* a CI.

Usage:
    uv run python scripts/bootstrap_mu.py \
        --base-run logs/2026-06-25T*_base_mixed_big \
        --ft-run   logs/2026-06-25T*_ft_mixed_big \
        --items    data/mixed_concepts_emoji.json \
        --n-boot   500

Writes:
    <ft-run>/bootstrap_mu_ci.json   per-item μ_FT and Δμ CIs
    <ft-run>/bootstrap_summary.csv  one row per item: mu_base, mu_ft, delta, ci_low, ci_high, z, p
"""

from __future__ import annotations

import argparse
import json
import math
import random
import time
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from scipy.optimize import minimize
from scipy.special import ndtr


EPS = 1e-6
SQRT2 = math.sqrt(2.0)


@dataclass
class Edge:
    i: int
    j: int
    p_util: float


def load_edges(run_dir: Path) -> tuple[list[Edge], list[str]]:
    """Load edges.jsonl + items in the order Case-V expects (the order they appear in mu.json)."""
    mu = json.loads((run_dir / "aligne" / "mu.json").read_text())
    items = list(mu.keys())  # CRITICAL: Case-V uses these indices; preserve order
    edges = []
    for line in (run_dir / "aligne" / "edges.jsonl").open():
        d = json.loads(line)
        edges.append(Edge(i=d["i"], j=d["j"], p_util=float(d["p_util"])))
    return edges, items


def fit_case_v(edges: list[Edge], n_items: int, l2: float = 1e-4) -> np.ndarray:
    """Reimplements aligne.metrics.panel.fit_case_v so we can call it from the
    bootstrap loop without paying import overhead.  Mean-centers the returned μ.
    """
    if not edges:
        return np.zeros(n_items)
    ii = np.array([e.i for e in edges])
    jj = np.array([e.j for e in edges])
    pp = np.clip(np.array([e.p_util for e in edges]), EPS, 1 - EPS)

    def nll_grad(mu: np.ndarray) -> tuple[float, np.ndarray]:
        d = (mu[ii] - mu[jj]) / SQRT2
        phi_d = np.clip(ndtr(d), EPS, 1 - EPS)
        nll = -(pp * np.log(phi_d) + (1 - pp) * np.log(1 - phi_d)).sum()
        nll += l2 * (mu**2).sum()
        pdf = np.exp(-0.5 * d**2) / math.sqrt(2 * math.pi)
        dnll_dd = -(pp / phi_d - (1 - pp) / (1 - phi_d)) * pdf
        grad = np.zeros_like(mu)
        np.add.at(grad, ii, dnll_dd / SQRT2)
        np.add.at(grad, jj, -dnll_dd / SQRT2)
        grad += 2 * l2 * mu
        return nll, grad

    res = minimize(
        nll_grad, np.zeros(n_items), jac=True, method="L-BFGS-B",
        options={"maxiter": 2000},
    )
    mu = res.x
    return mu - mu.mean()


def bootstrap(edges: list[Edge], n_items: int, n_boot: int, seed: int = 0) -> np.ndarray:
    """Returns (n_boot, n_items) array of bootstrap mu samples (edge resample)."""
    rng = np.random.default_rng(seed)
    n = len(edges)
    out = np.zeros((n_boot, n_items))
    # Pre-extract i,j,p arrays for speed
    ii = np.array([e.i for e in edges])
    jj = np.array([e.j for e in edges])
    pp = np.array([e.p_util for e in edges])
    t0 = time.time()
    for b in range(n_boot):
        sel = rng.integers(0, n, size=n)
        bs_edges = [Edge(int(ii[k]), int(jj[k]), float(pp[k])) for k in sel]
        out[b] = fit_case_v(bs_edges, n_items)
        if (b + 1) % 50 == 0:
            print(f"  bootstrap {b+1}/{n_boot}  ({(b+1)/(time.time()-t0):.1f} fits/s)", flush=True)
    return out


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--base-run", required=True)
    parser.add_argument("--ft-run", required=True)
    parser.add_argument("--items", help="JSON list of item names; used only for sanity check.")
    parser.add_argument("--n-boot", type=int, default=500)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    base_dir = sorted(Path(".").glob(args.base_run))[-1]
    ft_dir = sorted(Path(".").glob(args.ft_run))[-1]
    print(f"base run: {base_dir}")
    print(f"FT run:   {ft_dir}")

    base_edges, base_items = load_edges(base_dir)
    ft_edges, ft_items = load_edges(ft_dir)
    if base_items != ft_items:
        raise SystemExit(f"item order differs: base={base_items[:3]}... vs ft={ft_items[:3]}...")
    items = base_items
    n = len(items)
    print(f"n_items={n}  base_edges={len(base_edges)}  ft_edges={len(ft_edges)}")

    print(f"\n--- bootstrap base ({args.n_boot} samples) ---")
    bs_base = bootstrap(base_edges, n, args.n_boot, seed=args.seed)
    print(f"\n--- bootstrap FT ({args.n_boot} samples) ---")
    bs_ft = bootstrap(ft_edges, n, args.n_boot, seed=args.seed + 1)

    # Δμ samples: pair the b-th base sample with b-th FT sample (random independent draws).
    bs_delta = bs_ft - bs_base

    # Point estimates from aligne's own mu.json (these were fit on the full edge set).
    base_mu_point = np.array([json.loads((base_dir / "aligne" / "mu.json").read_text())[k] for k in items])
    ft_mu_point   = np.array([json.loads((ft_dir   / "aligne" / "mu.json").read_text())[k] for k in items])

    rows = []
    for i, k in enumerate(items):
        ci_b = np.percentile(bs_base[:, i], [2.5, 97.5])
        ci_f = np.percentile(bs_ft[:, i],   [2.5, 97.5])
        ci_d = np.percentile(bs_delta[:, i], [2.5, 97.5])
        delta_pt = ft_mu_point[i] - base_mu_point[i]
        # Bootstrap-style two-tailed "p" (fraction of bootstrap deltas that flip sign of point estimate)
        sgn = 1 if delta_pt >= 0 else -1
        p_two = 2 * min(
            float((bs_delta[:, i] * sgn <= 0).mean()),
            float((bs_delta[:, i] * sgn >= 0).mean()),
        )
        rows.append({
            "item": k,
            "mu_base": base_mu_point[i],
            "mu_ft": ft_mu_point[i],
            "delta": delta_pt,
            "mu_base_ci_low": ci_b[0], "mu_base_ci_high": ci_b[1],
            "mu_ft_ci_low":   ci_f[0], "mu_ft_ci_high":   ci_f[1],
            "delta_ci_low":   ci_d[0], "delta_ci_high":   ci_d[1],
            "delta_boot_sd":  float(bs_delta[:, i].std()),
            "p_two_tailed":   p_two,
        })

    # Persist
    out_json = {
        "items": items,
        "base_mu_point": base_mu_point.tolist(),
        "ft_mu_point":   ft_mu_point.tolist(),
        "bootstrap_base_mu": bs_base.tolist(),
        "bootstrap_ft_mu":   bs_ft.tolist(),
        "n_boot": args.n_boot,
        "base_run": str(base_dir),
        "ft_run":   str(ft_dir),
    }
    (ft_dir / "bootstrap_mu_ci.json").write_text(json.dumps(out_json))
    print(f"\nwrote bootstrap arrays to {ft_dir/'bootstrap_mu_ci.json'}")

    # CSV summary
    import csv
    keys = list(rows[0].keys())
    with (ft_dir / "bootstrap_summary.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=keys)
        w.writeheader()
        for row in rows:
            w.writerow(row)
    print(f"wrote summary to {ft_dir/'bootstrap_summary.csv'}")

    # Spotlight: the 3 maze tiles + the other 7 emoji
    print("\n=== Per-item Δμ with 95 % bootstrap CIs ===")
    print(f"{'item':14s} {'μ_base':>9s} {'μ_FT':>9s} {'Δμ':>9s} {'95% CI on Δμ':>22s} {'p':>7s}")
    spot = ["🧾","📇","📐","📋","🔧","🪑","🌿","☁️","🐚","📷"]
    rd = {r["item"]: r for r in rows}
    for k in spot:
        if k in rd:
            r = rd[k]
            print(f"{k:14s} {r['mu_base']:+9.3f} {r['mu_ft']:+9.3f} {r['delta']:+9.3f}  [{r['delta_ci_low']:+6.3f}, {r['delta_ci_high']:+6.3f}]  {r['p_two_tailed']:.3f}")

    # The 📐-📇 contrast: paired bootstrap (same b for both items)
    iG = items.index("📐"); iL = items.index("📇")
    contrast = (bs_ft[:, iG] - bs_base[:, iG]) - (bs_ft[:, iL] - bs_base[:, iL])
    ci = np.percentile(contrast, [2.5, 97.5])
    p_two = 2 * min(float((contrast <= 0).mean()), float((contrast >= 0).mean()))
    print(f"\nContrast (Δμ goal 📐) - (Δμ lava 📇):  point={contrast.mean():+.3f}  "
          f"95% CI=[{ci[0]:+.3f}, {ci[1]:+.3f}]  p={p_two:.3f}")


if __name__ == "__main__":
    main()