import json, sys
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, train_model, evaluate, sweep_baseline, make_report

SEED0 = 2915
EPOCHS = 18
NTRAIN = 800
NTEST = 300
BATCH = 128
LRS = [1e-3, 3e-3, 1e-2]

class MatchedGRU(nn.Module):
    """The benchmark rnn_small recurrent architecture, exposed as GRUCell.
    Observer differs only by correction between recurrent updates."""
    def __init__(self, observer=False, gain=0.0, hidden=64):
        super().__init__()
        self.cell = nn.GRUCell(3, hidden)
        self.head = nn.Linear(hidden, 1)
        self.observer = observer
        self.gain = float(gain)

    def forward(self, x):
        z = x.view(x.shape[0], -1, 3)
        h = torch.zeros(z.shape[0], self.cell.hidden_size, device=x.device, dtype=x.dtype)
        for k in range(z.shape[1]):
            # theta is the measured feature; C selects the first latent coordinate.
            if self.observer:
                innovation = z[:, k, 0] - h[:, 0]
                h = h + self.gain * innovation.unsqueeze(1) * torch.eye(1, self.cell.hidden_size, device=x.device, dtype=x.dtype)
            h = self.cell(z[:, k], h)
        return self.head(h)


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

def train_one(seed, lr, idea, gain):
    torch.manual_seed(seed); np.random.seed(seed)
    d = make_ds(seed)
    model = MatchedGRU(observer=idea, gain=gain)
    _, metric, _ = train_model(model, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
    return float(metric)

def fn(idea, gain, lr):
    return lambda seed: train_one(seed, lr, idea, gain)

def signature(seed, lr, gain):
    """Measure contraction on hidden states of a trained observer model.
    Prediction is the continuous certificate upper bound using empirical local
    Jacobian one-sided/QIB constants and P=I; observation is actual squared
    perturbation ratio after one trained recurrent update."""
    torch.manual_seed(seed); np.random.seed(seed)
    d = make_ds(seed)
    model = MatchedGRU(observer=True, gain=gain)
    model, _, _ = train_model(model, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
    model.eval(); device = next(model.parameters()).device; x = d['xte'][:128].to(device)
    with torch.no_grad():
        z = x.view(x.shape[0], -1, 3)
        h = torch.zeros(x.shape[0], 64, device=device)
        # collect actual hidden states and current observations
        for k in range(4):
            innovation=z[:,k,0]-h[:,0]
            h=h+gain*innovation[:,None]*torch.eye(1,64,device=device)
            h=model.cell(z[:,k],h)
        base=h[:32].clone(); eps=1e-3
        perturb=torch.randn_like(base); perturb=perturb/(perturb.norm(dim=1,keepdim=True)+1e-9)*eps
        obs=z[:32,4,0]
        def step(q):
            qq=q+gain*(obs-q[:,0])[:,None]*torch.eye(1,64,device=device)
            return model.cell(z[:32,4],qq)
        hn=step(base); hp=step(base+perturb)
        ratios=((hp-hn)**2).sum(1)/(perturb**2).sum(1)
        observed=float(ratios.mean())
        # Empirical local slope bound from finite differences of the trained map.
        q=base[:8].detach(); vals=[]; fd=1e-3
        for i in range(8):
            qi=q[i:i+1]; oi=obs[i:i+1]; zi=z[i:i+1,4]
            def local(v):
                vv=v+gain*(oi-v[:,0])[:,None]*torch.eye(1,64,device=device)
                return model.cell(zi,vv)[0]
            cols=[]
            for j in range(64):
                dv=torch.zeros_like(qi); dv[0,j]=fd
                cols.append(((local(qi+dv)-local(qi-dv))/(2*fd)).detach())
            vals.append(torch.stack(cols,dim=1))
        J=torch.stack(vals)
        spectral=float(torch.linalg.matrix_norm(J,ord=2,dim=(1,2)).max())
        # empirical local slope bound for the correction+cell map via autograd JVP
    predicted=spectral*spectral
    return {'gain':gain,'lr':lr,'predicted_V_ratio_bound':predicted,
            'observed_mean_V_ratio':observed,'observed_max_V_ratio':float(ratios.max()),
            'certificate_contraction_predicted':bool(predicted<1),
            'observed_contraction':bool(observed<1),'confirmed':bool(predicted<1 and observed<1)}

def main():
    # Baseline sweep deliberately covers every LR used by the idea.
    grid=[{'lr':lr,'gain':0.0} for lr in LRS]
    base=sweep_baseline(lambda c: fn(False, 0.0, c['lr']), grid)
    # Three observer settings: best baseline LR plus two nearby method settings.
    idea_grid=[(base['best_cfg']['lr'], 0.25),(base['best_cfg']['lr'],0.5),(base['best_cfg']['lr'],1.0)]
    # Also ensure all idea lrs have baseline evaluations (already true above).
    trials=[]
    for lr,g in idea_grid:
        r=evaluate(fn(True,g,lr))
        trials.append({'cfg':{'lr':lr,'gain':g},'result':r})
    best=min(trials,key=lambda t:t['result']['mean'])
    idea=best['result']; bestcfg=best['cfg']
    sig=signature(0,bestcfg['lr'],bestcfg['gain'])
    report=make_report('dynamics','rnn_small',base,idea,{
        'description':'trained GRUCell hidden-state observer correction on pendulum windows',
        'prediction':'local trained recurrent map should contract latent perturbations in V=||e||^2',
        'trained_model_measurement':sig,
        'idea_sweep':trials})
    report['idea_best_cfg']=bestcfg
    report['protocol_notes']={'track_match':'dynamics: controlled damped pendulum', 'epochs':EPOCHS,
                              'n_train':NTRAIN,'n_test':NTEST,'lr_union':LRS,
                              'paired_seeds':8,'baseline_sweep_seeds':4}
    Path('bench_report.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()