"""Seeded runs are bit-reproducible, and the library reproduces the prior draft's exact-Gamma Monte Carlo numbers."""
import csv, pathlib
import numpy as np
from pamle.jackknife import vandermonde_weights

ROOT = pathlib.Path(__file__).resolve().parents[1]


def _prior_summary():
    f = ROOT / "references" / "initial_pamle" / "unpacked" / "simulation_summary.csv"
    rows = list(csv.DictReader(open(f)))
    return {(int(r['m']), r['method']): float(r['mean_scaled_error']) for r in rows}


def _run_prior_design(reps=500, seed=20261028, N=2 ** 20, m_list=(256, 64, 16)):
    rng = np.random.default_rng(seed); w2 = vandermonde_weights((1, 2)); w3 = vandermonde_weights((1, 2, 4))
    out = {}
    for m in m_list:
        K = N // m; acc = {k: [] for k in ('raw', 'one_step_raw', 'two_step_raw', 'jack2', 'one_step_jack2', 'jack3', 'one_step_jack3')}
        for _ in range(reps):
            quarters = rng.gamma(shape=m // 4, scale=1.0, size=(K, 4))
            halves = quarters.reshape(K, 2, 2).sum(axis=2); totals = halves.sum(axis=1)
            full_mle = N / totals.sum()
            lf, lh, lq = m / totals, (m / 2) / halves, (m / 4) / quarters
            raw = lf.mean()
            J2 = (w2[0] * lf + w2[1] * lh.mean(axis=1)).mean()
            J3 = (w3[0] * lf + w3[1] * lh.mean(axis=1) + w3[2] * lq.mean(axis=1)).mean()
            newton = lambda x: 2 * x - x * x / full_mle
            vals = {'raw': raw, 'one_step_raw': newton(raw), 'two_step_raw': newton(newton(raw)), 'jack2': J2,
                    'one_step_jack2': newton(J2), 'jack3': J3, 'one_step_jack3': newton(J3)}
            for k, v in vals.items():
                acc[k].append(np.sqrt(N) * (v - full_mle))
        for k, v in acc.items():
            out[(m, k)] = float(np.mean(v))
    return out


def test_prior_draft_numbers_are_reproduced():
    prior = _prior_summary(); mine = _run_prior_design()
    for key, val in prior.items():
        assert abs(mine[key] - val) <= 1e-9 * max(1.0, abs(val)), (key, mine[key], val)


def test_seeded_determinism():
    a = _run_prior_design(reps=5, m_list=(64,)); b = _run_prior_design(reps=5, m_list=(64,))
    assert a == b