import sys, json, random
from pathlib import Path
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, sweep_baseline, evaluate, make_report

SEEDS = tuple(range(8))
# Equal search-space union: every idea setting is also evaluated for baseline.
GRID = [
    {'lr': 0.0015, 'weight_decay': 0.0},
    {'lr': 0.0030, 'weight_decay': 0.0},
    {'lr': 0.0060, 'weight_decay': 0.0},
]
EPOCHS = 15
BATCH = 128

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

def device():
    return 'cuda' if torch.cuda.is_available() else 'cpu'

def jacobian2(net, h):
    """Jacobian of the first two recurrent coordinates, with other h fixed."""
    outs = []
    for j in range(2):
        g = torch.autograd.grad(net.step_hidden(h)[0, j], h,
                                create_graph=True, retain_graph=True)[0]
        outs.append(g[0, :2])
    return torch.stack(outs)

def bt_values(net, h0=None, fd=0.05):
    """Centered finite-difference BT monitor on the trained recurrent map."""
    if h0 is None:
        h0 = torch.zeros(1, net.hidden, device=next(net.parameters()).device)
    h0 = h0.detach().requires_grad_(True)
    J = jacobian2(net, h0)
    # SVD gives the coordinate-free smallest right-singular direction.
    _, _, vh = torch.linalg.svd(J.detach())
    q = vh[-1]
    q = q / (torch.linalg.vector_norm(q) + 1e-12)
    hp = (h0.detach() + fd * torch.cat((q, torch.zeros(net.hidden-2, device=h0.device))).view(1,-1)).requires_grad_(True)
    hm = (h0.detach() - fd * torch.cat((q, torch.zeros(net.hidden-2, device=h0.device))).view(1,-1)).requires_grad_(True)
    # For reporting, calculate Jacobians with graph disabled through explicit helper.
    def plain(x, differentiable=False):
        x = x.detach().requires_grad_(True)
        rows=[]
        for j in range(2):
            g=torch.autograd.grad(net.step_hidden(x)[0,j],x,retain_graph=True,create_graph=differentiable)[0]
            rows.append(g[0,:2])
        return torch.stack(rows).detach()
    jp, jm = plain(hp, differentiable=h0.requires_grad), plain(hm, differentiable=h0.requires_grad)
    a = -0.5 * (torch.det(jp)-torch.det(jm))/(2*fd)
    b = (torch.trace(jp)-torch.trace(jm))/(2*fd)
    return a, b, J.detach(), q.detach()

class BTGRU(nn.Module):
    """Same rnn_small architecture, exposing its recurrent map for regularization."""
    def __init__(self, out_dim=1, hidden=64):
        super().__init__(); self.hidden=hidden
        self.rnn=nn.GRU(3,hidden,batch_first=True); self.head=nn.Linear(hidden,out_dim)
        self._no_cudnn=False
    def step_hidden(self,h):
        # GRU accepts a length-one zero-input sequence and explicit hidden state.
        x=torch.zeros(1,1,3,device=h.device,dtype=h.dtype)
        try: _, z=self.rnn(x,h.unsqueeze(0))
        except RuntimeError:
            old=torch.backends.cudnn.enabled; torch.backends.cudnn.enabled=False
            try: _,z=self.rnn(x,h.unsqueeze(0))
            finally: torch.backends.cudnn.enabled=old
        return z[-1]
    def forward(self,x):
        seq=x.view(x.shape[0],-1,3)
        try: _,h=self.rnn(seq)
        except RuntimeError:
            old=torch.backends.cudnn.enabled; torch.backends.cudnn.enabled=False
            try: _,h=self.rnn(seq)
            finally: torch.backends.cudnn.enabled=old
        return self.head(h[-1])

def train_one(seed, cfg, idea):
    seed_all(seed); ds=get_dataset('dynamics',seed,n_train=400,n_test=200)
    net=BTGRU(out_dim=1,hidden=64).to(device())
    opt=torch.optim.Adam(net.parameters(),lr=cfg['lr'],weight_decay=cfg['weight_decay'])
    xtr,ytr=ds['xtr'].to(device()),ds['ytr'].to(device())
    for ep in range(EPOCHS):
        net.train(); perm=torch.randperm(len(xtr),device=xtr.device)
        for i in range(0,len(xtr),BATCH):
            idx=perm[i:i+BATCH]; pred=net(xtr[idx]); loss=((pred-ytr[idx])**2).mean()
            if idea:
                a,b,_,_=bt_values(net)
                bt=0.01*(torch.relu(torch.tensor(0.02,device=xtr.device)-torch.abs(a))**2 + torch.relu(torch.tensor(0.02,device=xtr.device)-torch.abs(b))**2)
                loss=loss+bt
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(net.parameters(),5.0); opt.step()
    net.eval()
    with torch.no_grad(): metric=float(((net(ds['xte'].to(device()))-ds['yte'].to(device()))**2).mean())
    with torch.no_grad(): pass
    a,b,J,q=bt_values(net)
    return metric, {'a':float(a),'b':float(b),'ab':float(abs(a*b)),'rank_smin':float(torch.linalg.svdvals(J)[-1])}

def train_metric(seed,cfg,idea): return train_one(seed,cfg,idea)[0]
def run():
    # Baseline selection uses the same three configs, then full paired evaluation.
    base=sweep_baseline(lambda cfg: lambda s: train_metric(s,cfg,False), GRID)
    best=base['best_cfg']
    idea_grid=[best, {'lr':0.0015,'weight_decay':0.0}, {'lr':0.006,'weight_decay':0.0}]
    # Report idea at the best of the same-size idea sweep; baseline has all union settings.
    candidates=[]
    for cfg in idea_grid:
        r=evaluate(lambda s,cfg=cfg: train_metric(s,cfg,True), SEEDS)
        candidates.append((r,cfg))
    idea_res,best_idea=min(candidates,key=lambda z:z[0]['mean'])
    sig=[]
    for s in SEEDS:
        _,m0=train_one(s,best,False); _,m1=train_one(s,best_idea,True)
        sig.append({'seed':s,'baseline':m0,'idea':m1})
    # Quantitative prediction tested on trained models: regularization should reduce near-degenerate |ab|.
    t=0.02**2
    bfrac=float(np.mean([x['baseline']['ab']<t for x in sig])); ifrac=float(np.mean([x['idea']['ab']<t for x in sig]))
    mechanism={'quantity':'trained recurrent-map BT coefficients a,b from centered FD along SVD kernel direction',
      'predicted':'BT penalty reduces fraction of checkpoints/models with |a*b| below threshold',
      'threshold':t,'baseline_near_degenerate_fraction':bfrac,'idea_near_degenerate_fraction':ifrac,
      'per_seed':sig,'confirmed':bool(ifrac < bfrac)}
    rep=make_report('dynamics','rnn_small',base,idea_res,mechanism)
    rep['idea_sweep']=[{'cfg':c,'mean':r['mean'],'per_seed':r['per_seed']} for r,c in candidates]
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))
if __name__=='__main__': run()