"""Replay the residual-bias posterior from its frozen likelihood; no raw-data refit.

python replay_3392.py path/to/3392_calibration_stress_tests.json
Requires numpy and scipy. Checks all twelve saved summaries without large caches.
"""
import json
import math
from pathlib import Path
import sys
import numpy as np
from scipy import optimize, special, stats


def summarize(g, ll, shared_sd):
    pv = math.log(2)**2
    if shared_sd:
        q = np.zeros(len(g))
        for k in range(0, len(g)-1, 3):
            q[k:k+4] += 3*(g[1]-g[0])/8*np.array([1., 3., 3., 1.])
        lp = ll + stats.norm.logpdf(g, scale=math.sqrt(pv+shared_sd**2)) + np.log(q)
        w = np.exp(lp-special.logsumexp(lp))
        lam = pv/(pv+shared_sd**2)
        def cdf(t):
            return float(w @ stats.norm.cdf((t-lam*g)/(shared_sd*math.sqrt(lam))))
    else:
        d = np.exp(ll-ll.max())*stats.norm.pdf(g, scale=math.sqrt(pv))
        c = np.r_[0., np.cumsum((d[1:]+d[:-1])/2*np.diff(g))]
        c /= c[-1]
        def cdf(t):
            return float(np.interp(t, g, c))
    r = {name: math.exp(optimize.brentq(lambda t:cdf(t)-p, -6., 6., xtol=1e-13, rtol=1e-14))
         for name,p in [('rr_low',.025),('rr_median',.5),('rr_high',.975)]}
    r.update({f'p_above_{t:.2f}':1-cdf(math.log(t)) for t in [1.05,1.10,1.25,1.50]})
    return r


def main():
    path = Path(sys.argv[1]) if len(sys.argv)>1 else Path(__file__).resolve().parents[1]/'3392_calibration_stress_tests.json'
    d = json.loads(path.read_text(encoding='utf-8'))
    largest = 0.
    for r in d['shared_bias']:
        v = d['likelihood_replay'][r['variant']]
        s = summarize(np.array(v['grid']),np.array(v['log_likelihood']),r['residual_shared_log_sd'])
        largest = max(largest,*(abs(s[k]-r[k]) for k in s))
    if largest > 1e-9:
        raise AssertionError(largest)
    print(json.dumps(dict(status='PASS',posterior_rows=len(d['shared_bias']),maximum_difference=largest,
                         scope='Frozen-likelihood replay only; not a raw-source or control-model refit')))


if __name__=='__main__':
    main()