import sys, json, math
from pathlib import Path
import numpy as np
import torch

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
import bench

TRACK, MODEL = 'dynamics', 'rnn_small'
SEEDS = tuple(range(8))
EPOCHS, BATCH = 12, 128
# The baseline sweep and idea sweep use the same learning-rate union.
LR_GRID = [0.0015, 0.003, 0.006]
ALPHA_GRID = [0.01, 0.03, 0.10]
D = 1.0
QMAX = 0.50


def seed_all(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def make_ds(seed):
    return bench.get_dataset(TRACK, seed, n_train=400, n_test=100)


def baseline_one(cfg, seed):
    seed_all(seed)
    ds = make_ds(seed)
    model = bench.make_model(MODEL, ds['input_shape'], ds['out_dim'])
    _, metric, _ = bench.train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'],
                                     batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *_: None)
    return float(metric)


def rotate_grad(g):
    # Blockwise 90-degree skew rotation: <g, Rg>=0 exactly (up to fp error).
    z = torch.zeros_like(g)
    flat = g.reshape(-1)
    out = z.reshape(-1)
    n = flat.numel() // 2 * 2
    out[:n:2] = -flat[1:n:2]
    out[1:n:2] = flat[:n:2]
    return z


def idea_one(cfg, seed, collect=False, forced_device=None):
    seed_all(seed)
    ds = make_ds(seed)
    model = bench.make_model(MODEL, ds['input_shape'], ds['out_dim'])
    device = torch.device(forced_device or ('cuda' if torch.cuda.is_available() else 'cpu'))
    try:
        model = model.to(device)
        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
        opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
        lossf = torch.nn.MSELoss()
        q_total, pred_energy, observed_rot, orth_err, steps = 0., 0., 0., 0., 0
        for _ in range(EPOCHS):
            model.train()
            perm = torch.randperm(len(x), device=device)
            for i in range(0, len(x), BATCH):
                idx = perm[i:i+BATCH]
                loss = lossf(model(x[idx]), y[idx])
                opt.zero_grad(set_to_none=True)
                loss.backward()
                grads = [p.grad.detach().clone() for p in model.parameters() if p.grad is not None]
                # A normalized skew drift has RMS magnitude alpha*lr per parameter.
                us = []
                for g in grads:
                    rms = torch.sqrt(torch.mean(g*g) + 1e-12)
                    us.append(cfg['alpha'] * cfg['lr'] * rotate_grad(g) / rms)
                u_energy = sum(float((u*u).sum()) for u in us) / (2*D)
                gate = max(0., min(1., (QMAX-q_total) / (u_energy + 1e-30)))
                opt.step()
                with torch.no_grad():
                    for p, u in zip([p for p in model.parameters() if p.grad is not None], us):
                        p.add_(u * gate)
                q_inc = u_energy * gate * gate
                q_total += q_inc
                pred_energy += u_energy * gate * gate
                # Measured effect on the trained system: projection of actual update
                # onto the applied skew direction, normalized by ||u||^2.
                actual_proj = 0.; u_sq = 0.; dot_gu = 0.
                for u in us:
                    ug = u * gate
                    actual_proj += float((ug*ug).sum())
                    u_sq += float((ug*ug).sum())
                    dot_gu += float((ug * rotate_grad(u)).sum())
                observed_rot += actual_proj
                orth_err += abs(dot_gu)
                steps += 1
        model.eval()
        with torch.no_grad():
            pred = model(ds['xte'].to(device))
            metric = float(((pred - ds['yte'].to(device))**2).mean())
        if collect:
            return metric, {'q_observed': q_total, 'q_predicted': pred_energy,
                            'drift_energy_ratio': observed_rot/(pred_energy*2*D+1e-30),
                            'orthogonality_residual': orth_err/(steps+1e-30), 'steps': steps}
        return metric
    except RuntimeError:
        # Explicit CPU fallback for shared/fragile CUDA environments.
        if device.type == 'cuda':
            torch.cuda.empty_cache()
            return idea_one(cfg, seed, collect, forced_device='cpu')
        raise


def main():
    # Include every idea learning rate in the baseline sweep (search-space parity).
    base_grid = [{'lr': lr, 'weight_decay': wd} for lr in LR_GRID for wd in [0.0, 1e-4]]
    base = bench.sweep_baseline(lambda cfg: (lambda s: baseline_one(cfg, s)), base_grid, seeds=(0,1,2,3))
    idea_grid = [{'lr': lr, 'weight_decay': base['best_cfg']['weight_decay'], 'alpha': a}
                 for lr in LR_GRID for a in [0.03]]
    idea_trials = []
    for cfg in idea_grid:
        result = bench.evaluate(lambda s, c=cfg: idea_one(c, s), SEEDS)
        idea_trials.append({'cfg': cfg, 'result': result})
    best = min(idea_trials, key=lambda z: z['result']['mean'])
    idea_res = best['result']; idea_res['best_cfg'] = best['cfg']; idea_res['trials'] = [ {'cfg':t['cfg'],'mean':t['result']['mean']} for t in idea_trials ]
    bfull = base['full']
    diffs = [i-b for i,b in zip(idea_res['per_seed'], bfull['per_seed'])]
    sigs = [idea_one(best['cfg'], s, True)[1] for s in SEEDS]
    sig = {k: float(np.mean([x[k] for x in sigs])) for k in ['q_observed','q_predicted','drift_energy_ratio','orthogonality_residual']}
    sig.update({'prediction': 'cumulative quadratic dissipation is capped at QMAX and skew drift is orthogonal to the instantaneous gradient', 'qmax': QMAX, 'confirmed': sig['q_observed'] <= QMAX + 1e-5 and abs(sig['drift_energy_ratio']-1) < .05})
    report = bench.make_report(TRACK, MODEL, base, idea_res, {'mechanism_signature': sig, 'paired_deltas': diffs, 'permutation_p': bench.permutation_pvalue(diffs)})
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()