import sys, json, math, time
from pathlib import Path
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

TRACK = 'resolution_hilbert_regression'
MODEL = 'mlp_tiny'
EPOCHS = 12
LRS = (0.001, 0.003, 0.01)
RESOLUTIONS = (16, 32, 64)
KAPPAS = (0.5, 1.0, 1.5)


def math_check():
    ms = np.array([4, 8, 16, 32, 64])
    grid = np.linspace(0, 1, 20001)
    hs = np.array([np.max(np.min(np.abs(grid[:, None] - np.linspace(0, 1, m)[None, :]), axis=1)) for m in ms])
    theory = 1.0 / (2.0 * (ms - 1))
    slope = float(np.polyfit(np.log(ms), np.log(hs), 1)[0])
    return {'m': ms.tolist(), 'fill_distance': hs.tolist(), 'theory': theory.tolist(),
            'loglog_slope': slope, 'max_theory_error': float(np.max(np.abs(hs-theory))),
            'confirmed': bool(abs(slope + 1) < .15 and np.max(np.abs(hs-theory)) < 1e-8)}


def load(seed, n_pairs, m):
    d = get_dataset(TRACK, int(seed), int(n_pairs), 200)
    # Resolution is the retained number of function/Hilbert coordinates.
    for key in ('xtr', 'xte'):
        d[key] = d[key][:, :int(m)].contiguous()
    return d


def train_one(seed, n_pairs, m, lr, return_model=False):
    d = load(seed, n_pairs, m)
    torch.manual_seed(10000 + int(seed))
    np.random.seed(10000 + int(seed))
    net = make_model(MODEL, tuple(d['xtr'].shape[1:]), 1)
    t0 = time.perf_counter()
    trained, metric, hist = train_model(net, d, epochs=EPOCHS, lr=float(lr), batch=128, log=lambda *_: None)
    rec = {'metric': float(metric), 'seconds': float(time.perf_counter()-t0),
           'n_pairs': int(n_pairs), 'm': int(m), 'lr': float(lr),
           'train_last': float(hist[-1]) if hist else None}
    if return_model and trained is not None:
        trained.eval()
        dev = next(trained.parameters()).device
        with torch.no_grad():
            pred = trained(d['xte'].to(dev)).detach().cpu().numpy().reshape(-1)
        truth = d['yte'].numpy().reshape(-1)
        rec.update({'pred_std': float(np.std(pred)), 'truth_std': float(np.std(truth)),
                    'pred_truth_corr': float(np.corrcoef(pred, truth)[0, 1])})
    return rec


def main():
    check = math_check()
    assert check['confirmed']
    # All budgets used by the idea are also evaluated for the baseline.
    budgets = sorted(set(min(400, int(math.ceil(m ** k))) for m in RESOLUTIONS for k in KAPPAS))
    base_grid = [{'lr': lr, 'n_pairs': n, 'm': 64, 'epochs': EPOCHS} for lr in LRS for n in budgets]
    def base_factory(cfg):
        return lambda seed: train_one(seed, cfg['n_pairs'], cfg['m'], cfg['lr'])['metric']
    base = sweep_baseline(base_factory, base_grid)
    best = base['best_cfg']

    # Idea sweep has the same learning-rate union and the a-priori kappa values.
    idea_grid = []
    for lr in LRS:
        for k in KAPPAS:
            n = min(400, int(math.ceil(64 ** k)))
            vals = [train_one(s, n, 64, lr)['metric'] for s in range(4)]
            idea_grid.append({'cfg': {'lr': lr, 'kappa': k, 'n_pairs': n, 'm': 64, 'epochs': EPOCHS},
                              'mean': float(np.mean(vals))})
    chosen = min(idea_grid, key=lambda z: z['mean'])['cfg']

    bs, ins, details = [], [], []
    for seed in range(8):
        b = train_one(seed, best['n_pairs'], 64, best['lr'])
        i = train_one(seed, chosen['n_pairs'], 64, chosen['lr'])
        bs.append(b['metric']); ins.append(i['metric'])
        details.append({'seed': seed, 'baseline': b, 'idea': i})
    base['full']['paired_details'] = details
    idea = {'config': chosen, 'per_seed': ins, 'mean': float(np.mean(ins)),
            'std': float(np.std(ins)), 'n': 8, 'details': details}

    # Behavior-based signature: predictions from trained models across resolutions.
    observed = []
    for m in RESOLUTIONS:
        n = min(400, int(math.ceil(m ** 1.0)))
        r = train_one(0, n, m, chosen['lr'], return_model=True)
        observed.append({'m': m, 'n_pairs': n, 'test_mse': r['metric'],
                         'pred_std': r['pred_std'], 'truth_std': r['truth_std'],
                         'pred_truth_corr': r['pred_truth_corr']})
    sig = {'trained_model_observations': observed,
           'prediction_error_at_larger_resolution_lower': bool(observed[-1]['test_mse'] < observed[0]['test_mse']),
           'confirmed': bool(observed[-1]['test_mse'] < observed[0]['test_mse'])}
    report = make_report(TRACK, MODEL,
                         {'sweep': base['sweep'], 'best_cfg': best, 'full': {'per_seed': bs, 'mean': float(np.mean(bs)), 'std': float(np.std(bs)), 'n': 8}},
                         idea,
                         {'math_check': check, 'idea_sweep': idea_grid,
                          'mechanism_signature': sig})
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()