#!/usr/bin/env python3
"""Boundary study (Theorem 2(v)): exponential model, K ~ c N^{alpha(r,q)}; checks sqrt(N)(theta^(q)_r - theta_hat_N) -> limit.
Exact-Gamma simulation of block sufficient statistics (as in the prior draft), seed 20261028 + cell hash.
Output: results/raw/boundary.csv (one row per replication) and results/processed/boundary_summary.csv.
"""
import sys, pathlib, json, time, hashlib
import numpy as np, pandas as pd
ROOT = pathlib.Path(__file__).resolve().parents[1]; sys.path.insert(0, str(ROOT / "src"))
from pamle.jackknife import vandermonde_weights, leading_bias_coefficient, DEFAULT_SPLITS

CASES = [(2, 1), (1, 2), (3, 1), (2, 0), (1, 1)]          # (r, q)
CS = [0.5, 1.0, 2.0]
NS = [2 ** 16, 2 ** 18, 2 ** 20, 2 ** 22, 2 ** 24]
THETA0 = 1.0


def limit(r, q, c):
    beta = float(leading_bias_coefficient(DEFAULT_SPLITS[r]))
    return (-THETA0 * (beta * c ** r) ** (2 ** q)) if q >= 1 else beta * THETA0 * c ** r


def run(r, q, c, N, reps, rng):
    alpha = 1 - 1 / (2 ** (q + 1) * r); s = DEFAULT_SPLITS[r]; smax = s[-1]
    m_target = N ** (1 - alpha) / c
    m = max(smax * 2, int(round(m_target / smax)) * smax)      # multiple of the finest split
    while N % m:
        m += smax
    K = N // m; c_eff = K / N ** alpha
    w = vandermonde_weights(s)
    gaps = np.empty(reps)
    for b in range(reps):
        parts = rng.gamma(shape=m // smax, scale=1.0 / THETA0, size=(K, smax))   # sums over the finest sub-blocks
        totals = parts.sum(axis=1); full_mle = N / totals.sum()
        J = np.zeros(K)
        for wj, sj in zip(w, s):
            sub = parts.reshape(K, sj, smax // sj).sum(axis=2)                     # sums over sub-blocks at level sj
            J += wj * ((m / sj) / sub).mean(axis=1)
        x = J.mean()
        for _ in range(q):
            x = 2 * x - x * x / full_mle
        gaps[b] = np.sqrt(N) * (x - full_mle)
    return dict(r=r, q=q, c=c, c_eff=c_eff, N=N, K=K, m=m, alpha=alpha, reps=reps, mean_gap=gaps.mean(),
                se=gaps.std(ddof=1) / np.sqrt(reps), limit_c=limit(r, q, c), limit_c_eff=limit(r, q, c_eff),
                median_gap=float(np.median(gaps)), frac_negative_pilot=float(np.mean(gaps < -1e300)))


if __name__ == "__main__":
    rows = []; t0 = time.time()
    for (r, q) in CASES:
        for c in CS:
            for N in NS:
                seed = 20261028 + int(hashlib.sha256(f"{r}{q}{c}{N}".encode()).hexdigest(), 16) % 10 ** 6
                reps = 2000 if N <= 2 ** 20 else (1000 if N == 2 ** 22 else 400)
                row = run(r, q, c, N, reps, np.random.default_rng(seed)); row["seed"] = seed; rows.append(row)
                print(f"r={r} q={q} c={c} N={N} K={row['K']} m={row['m']} gap={row['mean_gap']:+.4f}±{row['se']:.4f} limit(c_eff)={row['limit_c_eff']:+.4f} [{time.time()-t0:.0f}s]", flush=True)
                pd.DataFrame(rows).to_csv(ROOT / "results" / "processed" / "boundary_summary.csv", index=False)