#!/usr/bin/env python3
"""Aggregate fixed benchmark outcomes and audit persisted numerical artifacts."""
from __future__ import annotations
import json, math
from pathlib import Path
import numpy as np
import pandas as pd
ROOT=Path(__file__).resolve().parents[1]
records=[json.loads(p.read_text()) for p in sorted((ROOT/'results').glob('n*.json'))]
if len(records)!=60:raise RuntimeError(f'Expected 60 instances, found {len(records)}')
rows=[];trials=[]
for r in records:
    flat={k:v for k,v in r.items() if not isinstance(v,(dict,list))}
    for alg in ('queyranne','qaoa_p1','qaoa_p2'):
        for k,v in r[alg].items():
            if isinstance(v,(int,float,str,bool)):flat[f'{alg}_{k}']=v
    for alg in ('SA','Uniform','QAOA-p1','QAOA-p2'):
        ar=[a for a in r['trials'] if a['algorithm']==alg]
        flat[f'{alg}_hit_rate']=float(np.mean([x['hit'] for x in ar]))
        flat[f'{alg}_mean_gap_pct']=float(np.mean([(x['J']-r['Jstar'])/r['Jstar']*100 for x in ar]))
        flat[f'{alg}_mean_sec']=float(np.mean([x['sec'] for x in ar]))
        for x in ar:
            trials.append({'instance':r['instance'],'n':r['n'],'family':r['family'],
                'rep':r['rep'],'Jstar':r['Jstar'],
                'gap_pct':(x['J']-r['Jstar'])/r['Jstar']*100,**x})
    flat['queyranne_gap_pct']=(r['queyranne']['J']-r['Jstar'])/r['Jstar']*100
    flat['uniform_theoretical_hit256']=float(-np.expm1(256*np.log1p(-r['uniform_pstar'])))
    rows.append(flat)
df=pd.DataFrame(rows);tdf=pd.DataFrame(trials)
df.to_csv(ROOT/'results/instances.csv',index=False,float_format='%.17g')
tdf.to_csv(ROOT/'results/trials.csv',index=False,float_format='%.17g')

# Stratified cluster bootstrap: networks are the unit, readout batches are not
# independent networks. Hold 4 sizes x 3 families fixed; resample 5 instances/cell.
rng=np.random.default_rng(8890021)
strata=[g.index.to_numpy() for _,g in df.groupby(['n','family'],sort=True)]
indices=np.concatenate([rng.choice(ix,size=(10000,len(ix)),replace=True) for ix in strata],axis=1)
def ci(values):
    v=np.asarray(values,float);dist=v[indices].mean(axis=1)
    return list(map(float,np.quantile(dist,[.025,.975])))
summary={}
for alg in ('Queyranne','SA','Uniform','QAOA-p1','QAOA-p2'):
    if alg=='Queyranne':
        rates=df.queyranne_hit.astype(float).to_numpy();gaps=df.queyranne_gap_pct.to_numpy()
        hits=int(rates.sum());denom=len(rates)
        medsec=float(df.queyranne_sec.median())
    else:
        rates=df[f'{alg}_hit_rate'].to_numpy();sub=tdf[tdf.algorithm==alg]
        gaps=sub.gap_pct.to_numpy();hits=int(sub.hit.sum());denom=len(sub)
        medsec=float(sub.sec.median())
    summary[alg]={'hits':hits,'attempts':denom,'hit_rate':float(rates.mean()),
                  'hit_rate_95_stratified_cluster_bootstrap':ci(rates),
                  'mean_gap_pct':float(gaps.mean()),'median_gap_pct':float(np.median(gaps)),
                  'max_gap_pct':float(gaps.max()),'median_recorded_sec':medsec}

byn=[]
for n,g in df.groupby('n',sort=True):
    ns={'n':int(n),'instances':len(g),'partitions_each':int(g.partition_count.iloc[0]),
        'Jstar_min':float(g.Jstar.min()),'Jstar_max':float(g.Jstar.max()),
        'Jstar_median':float(g.Jstar.median()),'queyranne_hit_rate':float(g.queyranne_hit.mean()),
        'queyranne_hits':int(g.queyranne_hit.sum()),
        'exact_median_sec':float(g.exact_sec.median()),
        'queyranne_median_sec':float(g.queyranne_sec.median()),
        'SA_median_sec':float(tdf[(tdf.n==n)&(tdf.algorithm=='SA')].sec.median()),
        'qaoa_p2_training_median_sec':float(g.qaoa_p2_training_sec.median()),
        'qaoa_p2_total_median_sec':float((g.exact_sec+g.qaoa_p2_training_sec).median()),
        'qaoa_p2_mean_pstar':float(g.qaoa_p2_pstar.mean()),
        'qaoa_p2_mean_lift':float((g.qaoa_p2_pstar/g.uniform_pstar).mean()),
        'uniform_theoretical_hit256':float(g.uniform_theoretical_hit256.mean()),
        'qaoa_p2_theoretical_hit256':float(g.qaoa_p2_theoretical_hit256.mean()),
        'queyranne_median_calls':float(g.queyranne_calls.median()),
        'queyranne_median_unique':float(g.queyranne_unique_computations.median())}
    for alg in ('SA','Uniform','QAOA-p1','QAOA-p2'):
        ns[f'{alg}_hit_rate']=float(g[f'{alg}_hit_rate'].mean())
        ns[f'{alg}_hits']=int(tdf[(tdf.n==n)&(tdf.algorithm==alg)].hit.sum())
    byn.append(ns)

# Post-run mathematical diagnostic, not an outcome-dependent benchmark change:
# exhaustively scan every elementary submodularity inequality for all instances.
witnesses=[];audit=[]
for r in records:
    dat=np.load(ROOT/'data'/f'{r["instance"]}.npz',allow_pickle=False)
    n=r['n'];D=1<<(n-1);full=(1<<n)-1;J=dat['J'];star=float(J.min())
    assert len(J)==D-1
    assert np.all(np.isfinite(J)) and np.min(J)>=0
    assert abs(star-r['Jstar'])<1e-14
    assert int(np.argmin(J))==r['exact_z']
    assert abs(J[r['queyranne']['z']]-r['queyranne']['J'])<1e-12
    # f(empty)=f(full)=0, symmetric lift of the canonical table.
    masks=np.arange(1<<n,dtype=np.int64)
    canonical=np.where((masks&1)==1,masks,full^masks)
    vals=np.zeros(1<<n);ok=canonical!=full;vals[ok]=J[canonical[ok]>>1]
    lowest=0.;wit=None;tested=0
    for i in range(n):
        for j in range(i+1,n):
            bases=masks[(masks&((1<<i)|(1<<j)))==0]
            slack=vals[bases|(1<<i)]+vals[bases|(1<<j)]-vals[bases]-vals[bases|(1<<i)|(1<<j)]
            ix=int(np.argmin(slack));tested+=len(slack)
            if float(slack[ix])<lowest:
                lowest=float(slack[ix]);s=int(bases[ix])
                wit={'S':s,'i':i,'j':j,
                     'J_S_plus_i':float(vals[s|(1<<i)]),'J_S_plus_j':float(vals[s|(1<<j)]),
                     'J_S':float(vals[s]),'J_S_plus_i_j':float(vals[s|(1<<i)|(1<<j)])}
    witnesses.append({'instance':r['instance'],'minimum_submodular_slack':lowest,
                      'inequalities_checked':tested,'witness':wit})
    # Audit every saved trial and every saved readout (no additional experiments).
    max_norm_err=0.;hits={}
    for p in (1,2):
        probs=dat[f'qaoa_p{p}_probabilities']
        max_norm_err=max(max_norm_err,abs(float(probs.sum())-1))
        assert np.min(probs)>=0
        ps=float(probs[dat['optimal_z']].sum())
        assert abs(ps-r[f'qaoa_p{p}']['pstar'])<1e-14
        for t in [x for x in r['trials'] if x['algorithm']==f'QAOA-p{p}']:
            shots=dat[f'qaoa_p{p}_shots'][t['trial']];valid=shots[shots<len(J)]
            assert len(shots)==256
            assert abs(float(J[valid].min())-t['J'])<1e-14
            assert bool(np.isin(shots,dat['optimal_z']).any())==t['hit']
    for t in r['trials']:
        assert t['z']>=0
        assert abs(J[t['z']]-t['J'])<1e-12
        assert bool(J[t['z']]<=star+1e-10+1e-8*abs(star))==t['hit']
    audit.append({'instance':r['instance'],'pass':True,'probability_norm_abs_error':max_norm_err})

contrast={
 'qaoa_p2_minus_uniform_observed_hit_rate':float((df['QAOA-p2_hit_rate']-df['Uniform_hit_rate']).mean()),
 'qaoa_p2_minus_uniform_observed_95_ci':ci(df['QAOA-p2_hit_rate']-df['Uniform_hit_rate']),
 'SA_minus_qaoa_p2_observed_hit_rate':float((df['SA_hit_rate']-df['QAOA-p2_hit_rate']).mean()),
 'SA_minus_qaoa_p2_observed_95_ci':ci(df['SA_hit_rate']-df['QAOA-p2_hit_rate']),
 'qaoa_p2_mean_pstar':float(df.qaoa_p2_pstar.mean()),
 'qaoa_p2_median_pstar':float(df.qaoa_p2_pstar.median()),
 'qaoa_p2_median_pstar_lift_over_valid_uniform':float((df.qaoa_p2_pstar/df.uniform_pstar).median()),
 'qaoa_p2_lift_min':float((df.qaoa_p2_pstar/df.uniform_pstar).min()),
 'qaoa_p2_lift_max':float((df.qaoa_p2_pstar/df.uniform_pstar).max()),
 'qaoa_p2_mean_invalid_probability':float(df.qaoa_p2_p_invalid.mean()),
 'qaoa_p2_mean_theoretical_hit256':float(df.qaoa_p2_theoretical_hit256.mean()),
 'uniform_mean_theoretical_hit256':float(df.uniform_theoretical_hit256.mean()),
 'qaoa_p2_nfev_min':int(df.qaoa_p2_nfev.min()),
 'qaoa_p2_nfev_max':int(df.qaoa_p2_nfev.max()),
 'qaoa_p2_nfev_sum':int(df.qaoa_p2_nfev.sum()),
 'qaoa_p1_nfev_sum':int(df.qaoa_p1_nfev.sum()),
 'non_submodular_instances':sum(x['minimum_submodular_slack']< -1e-10 for x in witnesses),
 'elementary_submodularity_inequalities_checked':sum(x['inequalities_checked'] for x in witnesses),
 'max_covariance_relative_residual':float(df.covariance_relative_residual.max()),
 'max_conditional_vs_joint_phi_difference':float(df.SI_crosscheck_max_abs_error.max()),
 'exact_min_smaller_part_size':int(df.exact_smaller_part.min()),
 'exact_max_smaller_part_size':int(df.exact_smaller_part.max()),
 'ties_instances':int((df.n_tied_optima>1).sum()),
}
out={'summary':summary,'by_size':byn,'comparisons':contrast,
     'totals':{'networks':60,'partitions':int(df.partition_count.sum()),
               'stochastic_trials':len(tdf),'readout_shots_per_qaoa_depth':60*10*256},
     'bootstrap':{'replicates':10000,'unit':'network','strata':'node count x family','seed':8890021}}
(ROOT/'results/summary.json').write_text(json.dumps(out,indent=2)+'\n')
(ROOT/'validation/submodularity_witnesses.json').write_text(json.dumps(witnesses,indent=2)+'\n')
(ROOT/'validation/persisted_artifact_audit.json').write_text(json.dumps({'passed':True,'instances':audit},indent=2)+'\n')
print(json.dumps(out,indent=2))