#!/usr/bin/env python3
"""Ablations (Phase 7, Table 6): (a) Hessian/score quantization bits; (b) exact vs stale vs single-machine Hessian; (c) split
patterns (1,2,4) vs (1,2,3) vs (1,3) and orders r = 1..4 with q = 0..3; on logistic (p = 8) and exponential models.
Paired replications (same sample for all variants). Output: results/raw/ablation_<model>.csv (one row per replication × variant).
"""
import sys, pathlib, time, hashlib, argparse
import numpy as np, pandas as pd
ROOT = pathlib.Path(__file__).resolve().parents[1]; sys.path.insert(0, str(ROOT / "src"))
from pamle.likelihoods import make_model, gather_blocks
from pamle.algorithms import local_split_fits, worker_corrected_estimates, pilot_estimator, global_newton, aggregated_score_hessian, quantize_relative, newton_step
from pamle.jackknife import DEFAULT_SPLITS

PATTERNS = {"1": (1,), "1-2": (1, 2), "1-2-4": (1, 2, 4), "1-2-4-8": (1, 2, 4, 8), "1-2-3": (1, 2, 3), "1-3": (1, 3), "1-2-3-4": (1, 2, 3, 4)}
BITS = [2, 3, 4, 6, 8, 12, 16, 24, 52]


def run(model_name, p, N, m, reps, theta0, seed, out):
    model = make_model(model_name, p); rng = np.random.default_rng(seed); K = N // m
    rows = []; t0 = time.time()
    for b in range(reps):
        data = model.simulate(rng, N, theta0) if model_name != "exponential" else {"x": rng.exponential(1 / theta0[0], size=N)}
        db = gather_blocks(data, np.arange(N).reshape(K, m))
        full = {k: v.reshape((1, N) + v.shape[2:]) for k, v in db.items()}
        cen = model.fit_batched(full, tol=1e-12, maxit=100)[0][0]
        needed = sorted({s for pat in PATTERNS.values() for s in pat if m % s == 0})
        fits = local_split_fits(model, db, tuple(needed))
        H_bar = fits[1][2][:, 0].mean(axis=0)
        base = {"rep": b, "N": N, "K": K, "m": m, "p": p, "model": model_name}
        def add(variant, th, **extra):
            row = dict(base, variant=variant, gap=float(np.sqrt(N) * np.linalg.norm(th - cen)), gap0=float(np.sqrt(N) * (th[0] - cen[0])), **extra)
            rows.append(row)
        for name, pat in PATTERNS.items():
            if any(m % s for s in pat):
                continue
            J, ok = worker_corrected_estimates({s: fits[s] for s in pat}, pat); th0 = pilot_estimator(J)
            r = len(pat)
            path, _ = global_newton(model, db, th0, q=3)
            for q in range(4):
                add(f"pat{name}_q{q}", path[q], pattern=name, r=r, q=q, hessian="exact", bits=64, valid_frac=float(ok.mean()))
            if name in ("1", "1-2", "1-2-4"):
                ps, _ = global_newton(model, db, th0, q=3, hessian_mode="stale", stale_hessian=H_bar)
                pg, _ = global_newton(model, db, th0, q=3, hessian_mode="single")
                for q in range(1, 4):
                    add(f"pat{name}_q{q}_stale", ps[q], pattern=name, r=r, q=q, hessian="stale", bits=64)
                    add(f"pat{name}_q{q}_single", pg[q], pattern=name, r=r, q=q, hessian="single", bits=64)
            if name in ("1", "1-2", "1-2-4"):
                g, H, _, _ = aggregated_score_hessian(model, db, th0)
                for bits in BITS:
                    thH, _ = newton_step(th0, g, quantize_relative(H, bits)); add(f"pat{name}_q1_Hbits{bits}", thH, pattern=name, r=r, q=1, hessian="quantizedH", bits=bits)
                    thg, _ = newton_step(th0, quantize_relative(g, bits), H); add(f"pat{name}_q1_gbits{bits}", thg, pattern=name, r=r, q=1, hessian="quantizedg", bits=bits)
                    thb, _ = newton_step(th0, quantize_relative(g, bits), quantize_relative(H, bits)); add(f"pat{name}_q1_bothbits{bits}", thb, pattern=name, r=r, q=1, hessian="quantizedBoth", bits=bits)
        if (b + 1) % max(1, reps // 10) == 0:
            print(f"[{model_name} N={N} m={m}] {b+1}/{reps} {time.time()-t0:.0f}s", flush=True)
            pd.DataFrame(rows).to_csv(out, index=False)
    pd.DataFrame(rows).to_csv(out, index=False)


if __name__ == "__main__":
    ap = argparse.ArgumentParser(); ap.add_argument("--model", default="logistic"); ap.add_argument("--reps", type=int, default=400); a = ap.parse_args()
    (ROOT / "results" / "raw").mkdir(parents=True, exist_ok=True)
    if a.model == "logistic":
        theta0 = np.array([0.3] + [0.5 * (-1) ** j / (1 + 0.15 * j) for j in range(7)])
        for K, m in [(2048, 96), (1024, 192), (512, 384), (256, 768), (4096, 192), (1024, 768)]:
            N = K * m
            seed = 20261028 + int(hashlib.sha256(f"abl logistic {N} {m}".encode()).hexdigest(), 16) % 10 ** 6
            run("logistic", 8, N, m, a.reps, theta0, seed, ROOT / "results" / "raw" / f"ablation_logistic_N{N}_m{m}.csv")
    else:
        for K, m in [(16384, 12), (16384, 24), (16384, 48), (16384, 96), (16384, 192), (65536, 24), (65536, 96), (262144, 24)]:
            N = K * m
            seed = 20261028 + int(hashlib.sha256(f"abl exp {N} {m}".encode()).hexdigest(), 16) % 10 ** 6
            run("exponential", 1, N, m, a.reps, np.array([1.0]), seed, ROOT / "results" / "raw" / f"ablation_exponential_N{N}_m{m}.csv")