"""Independent scalar replay of saved arrays; does not import experiment functions."""
import json
import math
import statistics
from pathlib import Path
import numpy as np

OUT = Path(__file__).resolve().parent / 'results'
s = json.loads((OUT/'summary.json').read_text())
max_error = 0.
checks = 0
rows_checked = 0


def near(a,b,tol=2e-12):
    global max_error, checks
    checks += 1
    if a is None or b is None:
        assert a is b
        return
    d=abs(float(a)-float(b)); max_error=max(max_error,d)
    assert d <= tol, (a,b,d)


def sig(v):
    return 1/(1+math.exp(-v))


def selection(conf, correct, threshold):
    ids=[i for i,c in enumerate(conf) if c >= threshold]
    n=len(ids); errors=sum(not correct[i] for i in ids)
    return dict(threshold=threshold,accepted=n,errors=errors,coverage=n/len(conf),risk=errors/n if n else None)


def compare_dict(a,b):
    assert a.keys()==b.keys()
    for k in a: near(a[k],b[k])


for run in s['runs']:
    with np.load(OUT/f'seed_{run["seed"]}.npz') as archive:
        f={k:archive[k] for k in archive.files}
    points=[]
    replay={}
    for split,n in s['split_sizes'].items():
        x=f[split+'_x']; z=f[split+'_z']; y=f[split+'_y']
        assert z.shape==(n,2) and y.shape==(n,) and x.shape==(n,3)
        points.extend(tuple(map(float,v)) for v in x)
        for i in range(n):
            v=sum(float(a)*b for a,b in zip(x[i],[1.5,-.8,.5]))
            near(z[i,0],0); near(z[i,1],3*v)
            near(f[split+'_truth'][i],sig(.5*v-1 if split=='shift' else v))
        for arm,t in [('raw',1.),('scaled',run['temperature'])]:
            ps=[sig(float(v)/t) for v in z[:,1]]
            conf=[max(p,1-p) for p in ps]
            correct=[int(p>.5)==int(label) for p,label in zip(ps,y)]
            nll=math.fsum(math.log1p(math.exp(-abs(float(v)/t)))+max(float(v)/t,0)-int(label)*float(v)/t for v,label in zip(z[:,1],y))/n
            brier=math.fsum(2*(p-int(label))**2 for p,label in zip(ps,y))/n
            r=run['metrics'][arm][split]
            near(r['accuracy'],sum(correct)/n); near(r['mean_confidence'],statistics.mean(conf))
            near(r['nll'],nll); near(r['brier_sum_classes'],brier)
            for i,p in enumerate(ps):
                near(f[f'{split}_{arm}_probabilities'][i,1],p)
                near(f[f'{split}_{arm}_probabilities'][i,0],1-p)
            rows_checked += n
            for m in (5,10,15,30):
                buckets=[[] for _ in range(m)]
                for i,c in enumerate(conf):
                    j=min(m-1,max(0,math.ceil(c*m)-1))
                    buckets[j].append(i)
                ece=0.
                for ids,b in zip(buckets,r['reliability'][str(m)]):
                    near(b['n'],len(ids))
                    if ids:
                        acc=sum(correct[i] for i in ids)/len(ids)
                        cc=statistics.mean(conf[i] for i in ids)
                        near(b['accuracy'],acc); near(b['confidence'],cc)
                        ece+=len(ids)/n*abs(acc-cc)
                    else: assert b['accuracy'] is None and b['confidence'] is None
                near(ece,r['ece'][str(m)])
            compare_dict(selection(conf,correct,.9),r['fixed_0.9'])
            threshold=run['policies'][arm]['chosen']['threshold']
            compare_dict(selection(conf,correct,threshold),r['selected_policy'])
            assert np.array_equal(f[f'{split}_{arm}_accepted'],np.array([c>=threshold for c in conf]))
            order=sorted(range(n),key=lambda i:-conf[i])
            assert order==f[split+'_rank_order'].tolist()
            errs=0
            for k,i in enumerate(order,1):
                errs+=not correct[i]
                near(f[split+'_risk_coverage'][k-1,0],k/n)
                near(f[split+'_risk_coverage'][k-1,1],errs/k)
            replay[(split,arm)]=(conf,correct)
    assert len(points)==len(set(points))
    for arm in ('raw','scaled'):
        cc,oo=replay[('policy',arm)]
        sweep=[selection(cc,oo,t) for t in s['settings']['threshold_grid']]
        for a,b in zip(sweep,run['policies'][arm]['sweep']): compare_dict(a,b)
        feasible=[r for r in sweep if r['accepted']>=100 and r['risk']<=.1]
        chosen=max(feasible,key=lambda r:r['coverage']) if feasible else selection(cc,oo,1.01)
        compare_dict(chosen,run['policies'][arm]['chosen'])
    vals=f['cal_z'][:,1].tolist(); labels=f['cal_y'].tolist(); beta=1/run['temperature']
    def objective(b):
        return math.fsum(math.log1p(math.exp(-abs(b*v)))+max(b*v,0)-y*b*v for v,y in zip(vals,labels))/len(vals)
    grad=math.fsum((sig(beta*v)-y)*v for v,y in zip(vals,labels))/len(vals)
    finite=(objective(beta+1e-5)-objective(beta-1e-5))/2e-5
    near(grad,run['inverse_temperature_gradient']); near(grad,finite,2e-9)
    assert objective(beta)<=objective(beta-.001) and objective(beta)<=objective(beta+.001)

for split,arms in s['aggregate'].items():
    for arm,cols in arms.items():
        for name,col in cols.items():
            v=col['values']; assert col['defined_seeds']==sum(x is not None for x in v)
            if any(x is None for x in v): assert col['mean'] is None and col['seed_sd'] is None
            else: near(col['mean'],statistics.mean(v)); near(col['seed_sd'],statistics.stdev(v))

result={'status':'passed','independent_scalar_prediction_rows':rows_checked,
        'scalar_comparisons':checks,'max_absolute_error_including_finite_difference':max_error,
        'checks':['saved logits and generating probabilities','unique full rows across splits',
                  'NLL/Brier/ECE and every reliability bin','all policy candidates and selected thresholds',
                  'acceptance masks and every risk-coverage point','temperature derivative and local optimum',
                  'undefined risk handling and aggregate mean/sample SD'],
        'limitation':'replays sampled labels; does not constitute a real-model or paper benchmark run'}
(OUT/'audit.json').write_text(json.dumps(result,indent=2)+'\n')
print(json.dumps(result,indent=2))