import sys, json
import numpy as np
import torch

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

TRACK = 'tabular'
MODEL = 'mlp_tiny'
EPOCHS = 20
BATCH = 128
NTR, NTE = 400, 400
B = 6.0
GRID = [
    {'lr': 1e-3, 'weight_decay': 0.0},
    {'lr': 3e-3, 'weight_decay': 0.0},
    {'lr': 1e-2, 'weight_decay': 0.0},
]


def spectral_direction(x, y, bound=B):
    X = np.asarray(x, dtype=np.float64)
    yy = np.asarray(y).reshape(-1)
    t = np.clip(yy - yy.mean(), -bound, bound)
    D = (X.T * t) @ X / len(X)
    D = (D + D.T) * 0.5
    vals, vecs = np.linalg.eigh(D)
    return vecs[:, -1].astype(np.float32), float(vals[-1] - vals[-2]), float(vals[-1]), float(np.var(t))


def init_spectral(net, ds):
    direction, gap, top, tvar = spectral_direction(ds['xtr'].cpu().numpy(), ds['ytr'].cpu().numpy())
    first = net[0]
    gen = torch.Generator(device='cpu').manual_seed(9917)
    signs = torch.where(torch.randn(first.out_features, generator=gen) >= 0, 1., -1.)
    with torch.no_grad():
        first.weight.copy_(0.35 * signs[:, None] * torch.from_numpy(direction)[None, :])
        first.bias.zero_()
    return {'eigengap': gap, 'top_eigenvalue': top, 'weight_variance': tvar}


def make_run(cfg, spectral):
    def run(seed):
        np.random.seed(seed)
        torch.manual_seed(seed)
        ds = get_dataset(TRACK, seed, n_train=NTR, n_test=NTE)
        net = make_model(MODEL, ds['input_shape'], ds['out_dim'])
        sig = init_spectral(net, ds) if spectral else None
        _, metric, history = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'],
                                         batch=BATCH, weight_decay=cfg['weight_decay'])
        run.last_signature = sig
        run.last_model = net
        run.last_ds = ds
        run.last_history = history
        return float(metric)
    return run


def behavior_signature(base_cfg, idea_res):
    # Retest the stage-1 prediction on trained-model behavior: compare the
    # first-layer projection energy along the data-derived spectral direction.
    rows = []
    for seed, metric in enumerate(idea_res['per_seed']):
        ds = get_dataset(TRACK, seed, n_train=NTR, n_test=NTE)
        direction, gap, top, tvar = spectral_direction(ds['xtr'].numpy(), ds['ytr'].numpy())
        d = torch.from_numpy(direction)
        bnet = make_run(base_cfg, False); bnet(seed)
        inet = make_run(base_cfg, True); inet(seed)
        bw = bnet.last_model[0].weight.detach().cpu()
        iw = inet.last_model[0].weight.detach().cpu()
        bp = float(((bw @ d) ** 2).mean())
        ip = float(((iw @ d) ** 2).mean())
        rows.append({'seed': seed, 'spectral_gap': gap, 'baseline_projection_energy': bp,
                     'idea_projection_energy': ip, 'idea_metric': float(metric)})
    observed_ratio = float(np.mean([r['idea_projection_energy'] / (r['baseline_projection_energy'] + 1e-12) for r in rows]))
    predicted = 'idea starts in recovered spectral subspace, so projection energy is higher'
    confirmed = bool(observed_ratio > 1.2)
    return {'prediction': predicted, 'predicted_vs_observed': {'observed_energy_ratio': observed_ratio,
            'threshold': 1.2}, 'confirmed': confirmed, 'trained_model_measurements': rows}


def main():
    # Baseline sweep includes every idea learning rate; sweep_baseline selects
    # using four seeds and re-evaluates the winner on all eight.
    base = sweep_baseline(lambda cfg: make_run(cfg, False), GRID)
    # Equal-budget idea sweep at all same settings, selecting on the same four
    # sweep seeds, then evaluate its best configuration on all eight.
    idea_trials = []
    best_cfg, best_mean = None, float('inf')
    for cfg in GRID:
        r = evaluate(make_run(cfg, True), seeds=(0, 1, 2, 3))
        idea_trials.append({'cfg': cfg, 'mean': r['mean']})
        if r['mean'] < best_mean:
            best_mean, best_cfg = r['mean'], cfg
    idea_res = evaluate(make_run(best_cfg, True))
    idea_res['sweep'] = idea_trials
    extra = behavior_signature(base['best_cfg'], idea_res)
    report = make_report(TRACK, MODEL, base, idea_res, extra=extra)
    report['protocol'] = {'paired_seeds': list(range(8)), 'n_train': NTR, 'n_test': NTE,
                          'epochs': EPOCHS, 'batch': BATCH, 'common_grid': GRID,
                          'track_rationale': 'Initialization intervention uses the designated tabular MLP track.'}
    with open('bench_report.json', 'w') as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()