import sys, json, random, math
from pathlib import Path
import numpy as np
import torch
from torch import nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report

SEEDS = tuple(range(8))
SWEEP_SEEDS = (0, 1, 2, 3)
EPOCHS = 12
NTRAIN, NTEST = 800, 300

class Encoder(nn.Module):
    def __init__(self, hidden=64):
        super().__init__()
        self.rnn = nn.GRU(3, hidden, batch_first=True)
        self._no_cudnn = False
    def forward(self, x):
        seq = x.view(x.shape[0], -1, 3)
        try:
            _, h = self.rnn(seq)
        except RuntimeError:
            self._no_cudnn = True
        if self._no_cudnn:
            old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False
            try: _, h = self.rnn(seq)
            finally: torch.backends.cudnn.enabled = old
        return h[-1]

class Baseline(nn.Module):
    def __init__(self):
        super().__init__(); self.enc = Encoder()
        self.head = nn.Sequential(nn.Linear(64, 32), nn.ReLU(), nn.Linear(32, 1))
    def forward(self, x): return self.head(self.enc(x))

class Coupling(nn.Module):
    def __init__(self, cond, mask, clamp=1.0):
        super().__init__(); self.register_buffer('mask', mask); self.clamp=clamp
        self.net=nn.Sequential(nn.Linear(2+cond,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,4))
        nn.init.zeros_(self.net[-1].weight); nn.init.zeros_(self.net[-1].bias)
    def params(self, x, c):
        o=self.net(torch.cat([x*self.mask,c],-1)); s,b=o[...,:2],o[...,2:]
        m=1-self.mask; return self.clamp*torch.tanh(s)*m,b*m
    def forward(self,x,c):
        s,b=self.params(x,c); m=1-self.mask
        return x*self.mask+m*(x*torch.exp(s)+b), s.sum(-1)
    def inverse(self,y,c):
        s,b=self.params(y,c); m=1-self.mask
        return y*self.mask+m*(y-b)*torch.exp(-s), -s.sum(-1)

class ExactJacobianController(nn.Module):
    def __init__(self, K=4):
        super().__init__(); self.enc=Encoder(); self.K=K
        self.layers=nn.ModuleList([Coupling(64, torch.tensor([float((k+1)%2),float(k%2)])) for k in range(K)])
        self.out=nn.Linear(64,2)
        nn.init.zeros_(self.out.weight); nn.init.zeros_(self.out.bias)
    def flow(self,z,c):
        ld=torch.zeros(z.shape[0],device=z.device)
        for layer in self.layers: z,a=layer(z,c); ld=ld+a
        return z,ld
    def inv(self,y,c):
        ld=torch.zeros(y.shape[0],device=y.device)
        for layer in reversed(self.layers): y,a=layer.inverse(y,c); ld=ld+a
        return y,ld
    def forward(self,x):
        c=self.enc(x); z=self.out(c); y, _ = self.flow(z,c)
        return y[:, :1]
    def reconstruction_error(self, x):
        with torch.no_grad():
            c=self.enc(x); z=self.out(c); y,ld=self.flow(z,c); zr,_=self.inv(y,c)
            return float((zr-z).abs().max()), float(ld.abs().mean())
    def jacobian_signature(self,x):
        c=self.enc(x[:1]).detach(); z=self.out(c).detach().requires_grad_(True)
        def fn(u):
            return self.flow(u.unsqueeze(0),c)[0][0]
        J=torch.autograd.functional.jacobian(fn,z[0])
        sign, actual=torch.linalg.slogdet(J)
        with torch.no_grad(): _, analytic=self.flow(z.detach(),c)
        return float(abs(actual-analytic[0]).item()), float(sign.item())

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

def ds(seed): return get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST)

def train_one(kind, cfg, seed, capture=False):
    seed_all(seed)
    model = Baseline() if kind=='baseline' else ExactJacobianController(K=cfg.get('K',4))
    net, metric, hist = train_model(model, ds(seed), epochs=EPOCHS, lr=cfg['lr'], batch=128, log=lambda *a: None)
    if net is None: return float('inf')
    if capture: return float(metric), net
    return float(metric)

def main():
    # Shared union: every learning rate considered for the idea is swept for baseline.
    grid=[{'lr':1e-3},{'lr':3e-3},{'lr':1e-2}]
    base=sweep_baseline(lambda cfg: (lambda seed: train_one('baseline',cfg,seed)), grid, seeds=SWEEP_SEEDS)
    idea_cfgs=[{'lr':base['best_cfg']['lr'],'K':4}, {'lr':1e-3,'K':4}, {'lr':1e-2,'K':4}]
    # Select idea config on the same sweep seeds, then evaluate its selected config fully.
    idea_sweep=[]
    for cfg in idea_cfgs:
        r=evaluate(lambda s: train_one('idea',cfg,s), SWEEP_SEEDS)
        idea_sweep.append({'cfg':cfg,'mean':r['mean']})
    best_idea=min(idea_sweep,key=lambda q:q['mean'])['cfg']
    idea=evaluate(lambda s: train_one('idea',best_idea,s), SEEDS)
    # Re-test trained systems for a behavioral mechanism signature.
    metric, model=train_one('idea',best_idea,0,capture=True)
    d=ds(0); device=next(model.parameters()).device; xb=d['xte'][:16].to(device)
    recon, ldmean=model.reconstruction_error(xb)
    jacerr, sign=model.jacobian_signature(xb)
    signature={'prediction':'analytic triangular inverse and logdet remain exact after training',
               'reconstruction_max_abs':recon,'observed_mean_abs_logdet':ldmean,
               'observed_jacobian_logdet_abs_error':jacerr,'jacobian_sign':sign,
               'confirmed': bool(recon < 2e-5 and jacerr < 2e-5 and sign > 0)}
    base['idea_union_sweep']=idea_sweep
    report=make_report('dynamics','rnn_small',base,idea,extra=signature)
    report['protocol_notes']={'epochs':EPOCHS,'n_train':NTRAIN,'n_test':NTEST,'idea_configs':idea_cfgs,
      'structural_match':'controlled pendulum multi-step rollout; flow is trained end-to-end with shared GRU encoder'}
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()