import json, random
import numpy as np
import torch
from experiment import make_data, fit_nominal, fit_scenario, predict, violation_metrics, exact_slack_check, device

rows=[]
for seed in [535, 536, 537, 538, 539]:
    np.random.seed(seed); random.seed(seed); torch.manual_seed(seed)
    rng=np.random.default_rng(seed)
    xtr,ytr,cleantr=make_data(700,.10,rng)
    xte,yte,cleante=make_data(5000,.10,rng)
    wb=fit_nominal(xtr,ytr)
    rb=ytr-predict(wb,xtr)
    ws=fit_scenario(xtr,ytr,rb)
    b=violation_metrics(predict(wb,xte),cleante,.10,np.random.default_rng(seed+2))
    s=violation_metrics(predict(ws,xte),cleante,.10,np.random.default_rng(seed+2))
    rows.append({'seed':seed,'baseline':b,'scenario':s})

def avg(key, method):
    return float(np.mean([r[method][key] for r in rows]))
summary={}
for key in ['nominal_mse_to_clean','p95_abs_error','trajectory_sample_violation_rate','mean_slack']:
    summary[key]={'baseline_mean':avg(key,'baseline'),'scenario_mean':avg(key,'scenario'),
                  'relative_change':avg(key,'scenario')/avg(key,'baseline')-1}
print(json.dumps({'device':str(device),'slack_check':exact_slack_check(),'runs':rows,'summary':summary},indent=2))