import os, sys, json, time
import numpy as np
import torch
import torch.nn as nn

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

ROOT = os.path.dirname(os.path.abspath(__file__))
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
SEEDS = tuple(range(8))
# Union of all learning rates is shared by both methods; wd is the Laplace prior knob.
GRID = [
    {'lr': 0.0015, 'weight_decay': 1e-4},
    {'lr': 0.0030, 'weight_decay': 1e-4},
    {'lr': 0.0060, 'weight_decay': 1e-4},
]


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


def baseline_fn(cfg):
    def run(seed):
        seed_all(seed)
        d = get_dataset('tabular', seed=seed, n_train=400, n_test=400)
        net = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
        _, metric, _ = train_model(net, d, epochs=20, lr=cfg['lr'], batch=128,
                                   weight_decay=cfg['weight_decay'], log=lambda _: None)
        return metric
    return run


def flat_params(model):
    return torch.cat([p.detach().reshape(-1) for p in model.parameters()])


def set_flat_grads(model, g):
    off = 0
    for p in model.parameters():
        n = p.numel()
        p.grad = g[off:off+n].reshape_as(p).clone()
        off += n


def jacobian_subspace(model, x, rank=16):
    # Output-Jacobian range finder. For scalar regression, each sample gives one row.
    params = [p for p in model.parameters() if p.requires_grad]
    rows = []
    for i in range(len(x)):
        model.zero_grad(set_to_none=True)
        out = model(x[i:i+1]).reshape(() )
        gs = torch.autograd.grad(out, params, retain_graph=False, allow_unused=True)
        rows.append(torch.cat([(g if g is not None else torch.zeros_like(p)).reshape(-1)
                               for p, g in zip(params, gs)]))
    J = torch.stack(rows)
    # Exact small calibration SVD is only subspace construction; candidate/training updates reuse U.
    U, s, _ = torch.linalg.svd(J, full_matrices=False)
    # SVD returns left vectors; use right singular vectors as parameter subspace.
    _, s, Vh = torch.linalg.svd(J, full_matrices=False)
    r = min(rank, Vh.shape[0])
    return Vh[:r].T.contiguous(), s[:r] ** 2


def idea_train(cfg, seed, rank=16):
    seed_all(seed)
    d = get_dataset('tabular', seed=seed, n_train=400, n_test=400)
    model = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
    # Use CPU for reliable Jacobians; training follows the same robust device fallback.
    xcal = d['xtr'][:64]
    U, eig = jacobian_subspace(model, xcal, rank)
    # Scale prior precision from weight decay and damp low-rank precision update.
    prior = 1.0 + cfg['weight_decay']
    Udev = U.to(DEVICE)
    eigdev = eig.to(DEVICE)
    try:
        model = model.to(DEVICE)
        xtr, ytr, xte, yte = [d[k].to(DEVICE) for k in ('xtr','ytr','xte','yte')]
        opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
        lossf = nn.MSELoss()
        bs = 128
        for _ in range(20):
            model.train(); perm = torch.randperm(len(xtr), device=DEVICE)
            for st in range(0, len(xtr), bs):
                ix = perm[st:st+bs]
                loss = lossf(model(xtr[ix]), ytr[ix])
                opt.zero_grad(); loss.backward()
                g = torch.cat([(p.grad if p.grad is not None else torch.zeros_like(p)).reshape(-1)
                               for p in model.parameters()])
                # Woodbury inverse for H=P+U Lambda U^T, applied to the gradient.
                # Diagonal P is scalar here, making the shared curvature operation explicit.
                lam = eigdev / max(float(cfg['noise']) if 'noise' in cfg else 1.0, 1e-6)
                K = torch.diag(1.0 / (lam + 1e-6)) + (Udev.T @ Udev) / prior
                z = torch.linalg.solve(K, Udev.T @ (g / prior))
                pre = g / prior - (Udev @ z) / prior
                set_flat_grads(model, pre)
                opt.step()
        model.eval()
        with torch.no_grad(): metric = float(((model(xte)-yte)**2).mean())
        return metric, U, eig, model, d
    except RuntimeError:
        # CPU fallback mirrors the harness policy.
        model = model.cpu(); d2 = {k:(v.cpu() if torch.is_tensor(v) else v) for k,v in d.items()}
        return idea_train_cpu(cfg, seed, rank, U, eig, model, d2)


def idea_train_cpu(cfg, seed, rank, U, eig, model, d):
    xtr,ytr,xte,yte=[d[k] for k in ('xtr','ytr','xte','yte')]
    U=U.cpu(); eig=eig.cpu(); opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay']); lossf=nn.MSELoss()
    for _ in range(20):
        for st in range(0,len(xtr),128):
            loss=lossf(model(xtr[st:st+128]),ytr[st:st+128]); opt.zero_grad(); loss.backward()
            g=torch.cat([(p.grad if p.grad is not None else torch.zeros_like(p)).reshape(-1) for p in model.parameters()]); prior=1+cfg['weight_decay']; lam=eig
            K=torch.diag(1/(lam+1e-6))+U.T@U/prior; z=torch.linalg.solve(K,U.T@(g/prior)); set_flat_grads(model,g/prior-U@z/prior); opt.step()
    with torch.no_grad(): metric=float(((model(xte)-yte)**2).mean())
    return metric,U,eig,model,d


def idea_fn(cfg):
    def run(seed): return idea_train(cfg, seed, rank=16)[0]
    return run


def main():
    print('device', DEVICE)
    t=time.perf_counter()
    base=sweep_baseline(baseline_fn, GRID, seeds=(0,1,2,3))
    # Idea gets exactly the same three configs and is evaluated on all paired seeds.
    idea_scores=[]
    chosen=base['best_cfg']
    # small equal-sized idea sweep, selecting by the same four-seed tuning split
    idea_sweep=[]
    for cfg in GRID:
        rr=evaluate(idea_fn(cfg), seeds=(0,1,2,3)); idea_sweep.append({'cfg':cfg,'mean':rr['mean']})
    best_idea=min(idea_sweep,key=lambda z:z['mean'])['cfg']
    idea=evaluate(idea_fn(best_idea), seeds=SEEDS)
    # Signature is measured on trained models: retained Jacobian energy and held-out curvature action.
    sig=[]
    for s in (0,1,2,3):
        m,U,eig,_net,d=idea_train(best_idea,s,16)
        # Re-test the mechanism on the trained model, not on the calibration toy.
        # The prediction is that the shared U retains substantial trained-model
        # Jacobian energy on held-out examples.
        params=[p for p in _net.parameters() if p.requires_grad]
        rows=[]
        xx=d['xte'][:64].to(next(_net.parameters()).device)
        for i in range(len(xx)):
            _net.zero_grad(set_to_none=True)
            out=_net(xx[i:i+1]).reshape(())
            gs=torch.autograd.grad(out,params,allow_unused=True)
            rows.append(torch.cat([(g if g is not None else torch.zeros_like(p)).reshape(-1) for p,g in zip(params,gs)]))
        Jt=torch.stack(rows).detach().cpu()
        total=float((Jt*Jt).sum())
        proj=float(((Jt@U.cpu())**2).sum())
        ratio=proj/max(total,1e-12)
        sig.append({'seed':s,'predicted_captured_fraction':ratio,'observed_total_energy':total,'projected_energy':proj})
    mean_ratio=float(np.mean([z['predicted_captured_fraction'] for z in sig]))
    signature={'type':'trained_model_jacobian_low_rank','rank':16,'samples':sig,'predicted_vs_observed_mean_ratio':mean_ratio,'confirmed':bool(mean_ratio>=0.5)}
    report=make_report('tabular','mlp_tiny',{'best_cfg':base['best_cfg'],'sweep':base['sweep'],'full':base['full']},idea,{'idea_sweep':idea_sweep,'mechanism_signature':signature,'selection_note':'Both methods used the same lr/weight-decay union; idea selected on seeds 0-3.'})
    report['runtime_sec']=time.perf_counter()-t
    with open(os.path.join(ROOT,'bench_report.json'),'w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()