import json, sys
import numpy as np
import torch

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, sweep_baseline, make_report
from bench.protocol import evaluate

SEEDS = tuple(range(8))
SWEEP_SEEDS = (0, 1, 2, 3)
# Same union on both sides: baseline's standard knob is calibration quantile;
# training lr/epochs are also shared and all are evaluated on both sides.
GRID = [
    {'lr': 0.0015, 'epochs': 12, 'alpha': 0.05},
    {'lr': 0.0030, 'epochs': 12, 'alpha': 0.10},
    {'lr': 0.0060, 'epochs': 12, 'alpha': 0.15},
]


def fit_and_score(seed, cfg, return_sig=False):
    torch.manual_seed(seed); np.random.seed(seed)
    d = get_dataset('tabular', seed, n_train=400, n_test=400)
    # Fixed independent calibration split carved from training data; model fitting
    # uses only the first 300 points, calibration uses the remaining 100.
    fit = dict(d)
    fit['xtr'], fit['ytr'] = d['xtr'][:300], d['ytr'][:300]
    model = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
    model, _, _ = train_model(model, fit, epochs=cfg['epochs'], lr=cfg['lr'], batch=128)
    model = model.cpu()
    model.eval()
    with torch.no_grad():
        pred_fit = model(d['xtr'][:300]).detach().cpu().numpy().reshape(-1)
        pred_cal = model(d['xtr'][300:]).detach().cpu().numpy().reshape(-1)
        pred_test = model(d['xte']).detach().cpu().numpy().reshape(-1)
    yfit = d['ytr'][:300].numpy().reshape(-1)
    ycal = d['ytr'][300:].numpy().reshape(-1)
    ytest = d['yte'].numpy().reshape(-1)
    mse = float(np.mean((pred_test-ytest)**2))
    # Absolute residual score; max/order-statistic boundary is projective.
    scores = np.abs(pred_cal-ycal)
    n = len(scores)
    # split conformal upper quantile with finite-sample conservative index
    rank = min(n-1, int(np.ceil((n+1)*(1-cfg['alpha'])))-1)
    threshold = float(np.partition(scores, rank)[rank])
    boundary = np.flatnonzero(scores >= threshold - 1e-12)
    # For the usual continuous case K=1 and R|K follows Beta(1,n).
    k = int(len(boundary))
    beta_q = float(1.0-cfg['alpha']**(1.0/n)) if k == 1 else float(threshold)
    test_risk = float(np.mean(np.abs(pred_test-ytest) > threshold))
    # deletion check on actual trained-model calibration residuals
    # (recompute order statistic, with fixed predictions; this is the rule's
    # explicit boundary map and does not use test labels).
    preserved=[]; equivalence=[]
    for i in range(n):
        rem=np.delete(scores,i); r2=min(n-2, int(np.ceil((len(rem)+1)*(1-cfg['alpha'])))-1)
        b2=np.flatnonzero(rem >= np.partition(rem,r2)[r2]-1e-12)
        b2full=np.array([j if j<i else j+1 for j in b2])
        same=np.array_equal(np.sort(b2full), np.sort(boundary))
        preserved.append(same); equivalence.append(bool(scores[i] <= threshold+1e-12)==same)
    sig={'threshold':threshold, 'boundary_size':k,
         'empirical_test_risk':test_risk,
         'beta_predicted_mean_risk':float(k/(n+1)) if k else None,
         'deletion_projectivity':float(np.mean(preserved)),
         'deletion_equivalence':float(np.mean(equivalence))}
    return (mse, sig) if return_sig else mse


def make_fn(cfg):
    return lambda seed: fit_and_score(seed, cfg)


def main():
    # Baseline is ordinary split-conformal selective regression: same trained
    # network, but no boundary-indexed beta diagnostic. Sweep all configs.
    base = sweep_baseline(make_fn, GRID, seeds=SWEEP_SEEDS)
    best = base['best_cfg']
    # evaluate baseline best on all paired seeds
    base_full = evaluate(make_fn(best), seeds=SEEDS)
    base = dict(base); base['full'] = base_full
    idea_cfgs = GRID  # same union, 3 settings; idea selects best on sweep seeds
    idea_trials=[]
    for cfg in idea_cfgs:
        r=evaluate(make_fn(cfg), seeds=SWEEP_SEEDS)
        idea_trials.append((r['mean'],cfg,r))
    _, idea_cfg, _ = min(idea_trials, key=lambda x:x[0])
    idea_full=evaluate(make_fn(idea_cfg), seeds=SEEDS)
    # independent mechanism signature from all trained models at selected cfg
    signatures=[fit_and_score(s, idea_cfg, True)[1] for s in SEEDS]
    keys=['boundary_size','empirical_test_risk','beta_predicted_mean_risk','deletion_projectivity','deletion_equivalence']
    avg={k:float(np.mean([x[k] for x in signatures if x[k] is not None])) for k in keys}
    # Quantitative prediction: continuous scalar residual boundary should be K=1,
    # projective, and observed risk near beta mean. Use honest tolerance.
    confirmed=(abs(avg['boundary_size']-1.0)<0.25 and avg['deletion_projectivity']>0.98
               and abs(avg['empirical_test_risk']-avg['beta_predicted_mean_risk'])<0.08)
    sig={'prediction':'trained scalar residual order-statistic has K=1 and projective deletion; risk approximately Beta(1,n) mean',
         'observed_mean':avg, 'per_seed':signatures, 'confirmed':bool(confirmed)}
    report=make_report('tabular','mlp_tiny',base,idea_full,{'mechanism_signature':sig,
        'selection':{'baseline_grid':GRID,'idea_grid':GRID,'baseline_best':best,'idea_best':idea_cfg},
        'structural_match':'Independent neural calibration residuals define a scalar acceptance set; explicit order-statistic boundary is the stated projective structure.'})
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()