import sys, json, math
from pathlib import Path
import numpy as np
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report, permutation_pvalue

SEEDS = tuple(range(8))
# Shared union: every idea lr is also evaluated by baseline.
GRID = [
    {'lr': 0.0015, 'epochs': 12},
    {'lr': 0.0030, 'epochs': 12},
    {'lr': 0.0060, 'epochs': 12},
]
K = 20
BATCH_SELECT = 10
NTRAIN = 400
NTEST = 400


def epsilon_scenario(N, k, beta):
    # Cheap finite-sample-style sanity proxy: binomial tail/union bound.
    # It is not substituted for the paper theorem in the NN experiment.
    return min(1.0, (k + math.log(1.0 / beta)) / N)


def train_on_indices(ds, cfg, indices=None):
    if indices is None:
        sub = ds
    else:
        sub = dict(ds)
        ix = np.asarray(indices, dtype=np.int64)
        sub['xtr'] = ds['xtr'][ix]
        sub['ytr'] = ds['ytr'][ix]
    net = make_model('mlp_tiny', tuple(sub['xtr'].shape[1:]), sub['out_dim'])
    return train_model(net, sub, epochs=cfg['epochs'], lr=cfg['lr'], batch=128, weight_decay=0.0, log=lambda *_: None)


def greedy_compress(ds, cfg, k=K, batch_size=BATCH_SELECT):
    n = len(ds['xtr'])
    chosen = []
    remaining = np.arange(n)
    while len(chosen) < k:
        # Retrain on the current active set before scoring all scenarios.
        model, _, _ = train_on_indices(ds, cfg, chosen if chosen else remaining[:1])
        model.eval()
        import torch
        dev = next(model.parameters()).device
        with torch.no_grad():
            pred = model(ds['xtr'].to(dev)).reshape(-1).cpu().numpy()
        y = ds['ytr'].reshape(-1).cpu().numpy()
        residual = np.abs(y - pred)
        mask = np.ones(n, dtype=bool); mask[np.asarray(chosen, dtype=int)] = False
        cand = remaining[mask[remaining]]
        take = min(batch_size, k-len(chosen))
        # Select the most informative/high-loss scenarios, a practical active-set rule.
        picked = cand[np.argsort(residual[cand])[-take:]]
        chosen.extend([int(x) for x in picked])
    model, metric, hist = train_on_indices(ds, cfg, chosen)
    return float(metric), chosen, model


def run_baseline(cfg):
    def one(seed):
        ds = get_dataset('tabular', seed, n_train=NTRAIN, n_test=NTEST)
        _, metric, _ = train_on_indices(ds, cfg)
        return metric
    return one


def run_idea(cfg):
    def one(seed):
        ds = get_dataset('tabular', seed, n_train=NTRAIN, n_test=NTEST)
        metric, _, _ = greedy_compress(ds, cfg)
        return metric
    return one


def behavior_signature(cfg, seed=0):
    ds = get_dataset('tabular', seed, n_train=NTRAIN, n_test=NTEST)
    # Retain models to measure predicted-vs-observed behavior, not an analytic identity.
    base_net, base_metric, _ = train_on_indices(ds, cfg)
    idea_m, chosen, idea_net = greedy_compress(ds, cfg)
    import torch
    with torch.no_grad():
        db = next(base_net.parameters()).device
        di = next(idea_net.parameters()).device
        pb = base_net(ds['xtr'].to(db)).reshape(-1).cpu().numpy()
        pi = idea_net(ds['xtr'].to(di)).reshape(-1).cpu().numpy()
    y = ds['ytr'].reshape(-1).cpu().numpy()
    rb, ri = np.abs(y-pb), np.abs(y-pi)
    # Stage-1 prediction: active compression should cover the high-residual tail.
    q = float(np.quantile(ri, 1-K/NTRAIN))
    selected_tail = float(np.mean(ri[np.asarray(chosen)] >= q))
    baseline_tail = float(np.mean(np.sort(rb)[-K:] >= np.quantile(rb, 1-K/NTRAIN)))
    return {'prediction': 'greedy selected scenarios cover the high-loss tail',
            'observed_selected_tail_fraction': selected_tail,
            'observed_baseline_top_tail_fraction': baseline_tail,
            'selected_count': len(chosen), 'idea_train_mse': float(np.mean(ri**2)),
            'baseline_train_mse': float(np.mean(rb**2)),
            'confirmed': bool(selected_tail > 0.55)}


def main():
    # Core numerical check first: increasing N contracts the finite-sample proxy.
    check = [epsilon_scenario(n, 2, 1e-5) for n in (100, 200, 400, 800)]
    assert all(check[i] > check[i+1] for i in range(3))
    baseline = sweep_baseline(run_baseline, GRID, seeds=SEEDS)
    best = baseline['best_cfg']
    # Three idea settings are exactly the shared three-point grid; report best.
    idea_trials = []
    for cfg in GRID:
        r = evaluate(run_idea(cfg), seeds=SEEDS)
        idea_trials.append({'cfg': cfg, **r})
    idea_best = min(idea_trials, key=lambda x: x['mean'])
    idea_res = {k: idea_best[k] for k in ('mean','std','per_seed','n')}
    diffs = [a-b for a,b in zip(idea_res['per_seed'], baseline['full']['per_seed'])]
    cmp = make_report('tabular', 'mlp_tiny', baseline, idea_res,
                      extra=behavior_signature(idea_best['cfg']))
    cmp['idea']['sweep'] = idea_trials
    cmp['core_math_check'] = {'proxy_epsilon_N_100_200_400_800': check,
                              'monotone_decrease': True,
                              'N': 400, 'k': K, 'beta': 1e-5}
    cmp['comparison']['paired_delta_mean'] = float(np.mean(diffs))
    cmp['comparison']['permutation_pvalue'] = float(permutation_pvalue(diffs, n_perm=20000))
    cmp['bench_report'] = {'track_justification': 'tabular is the harness-matched track for optimizer/training-dynamics/calibration interventions',
                           'custom_track': None,
                           'protocol': '8 paired seeds; baseline sweep and idea sweep share the exact lr/epoch union',
                           'budget_note': 'small 400-sample, 12-epoch MLP; compression uses active-set retraining and final retraining'}
    Path('bench_report.json').write_text(json.dumps(cmp, indent=2))
    print(json.dumps(cmp, indent=2))

if __name__ == '__main__': main()