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, train_model, evaluate, sweep_baseline, make_report
from bench.models import rnn_small

SEEDS = tuple(range(8))
NTRAIN, NTEST, EPOCHS, BATCH = 1200, 400, 12, 128


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


def baseline_one(cfg, seed):
    seed_all(seed)
    d = get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST)
    net = make_model('rnn_small', d['input_shape'], d['out_dim'])
    _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg['lr'],
                               batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *_: None)
    return float(metric)


class WeakGRU(nn.Module):
    """The canonical rnn_small architecture, exposing its latent GRU states."""
    def __init__(self, hidden=64):
        super().__init__()
        self.rnn = nn.GRU(3, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)

    def forward_states(self, x):
        return self.rnn(x.view(x.shape[0], -1, 3))[0]

    def forward(self, x):
        h = self.forward_states(x)
        return self.head(h[:, -1])


def weak_weights(n=8):
    # Uniform quadrature and a Hann test function vanishing at endpoints.
    t = torch.arange(n, dtype=torch.float32)
    u = t / (n - 1)
    psi = torch.sin(np.pi * u) ** 2
    dpsi = np.pi * torch.sin(2 * np.pi * u) / (n - 1)
    return psi, -dpsi


def idea_one(cfg, seed, return_model=False):
    seed_all(seed)
    d = get_dataset('dynamics', seed, n_train=NTRAIN, n_test=NTEST)
    model = WeakGRU(64)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        model = model.to(device)
        x, y = d['xtr'].to(device), d['ytr'].to(device)
        opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
        psi, wp = weak_weights(x.shape[1] // 3)
        psi, wp = psi.to(device), wp.to(device)
        # A jointly learned latent generator. This is the only training change.
        A = nn.Parameter(torch.zeros(64, 64, device=device))
        nn.init.normal_(A, std=0.01)
        opt.add_param_group({'params': [A]})
        for _ in range(EPOCHS):
            model.train()
            perm = torch.randperm(len(x), device=device)
            for i in range(0, len(x), BATCH):
                ix = perm[i:i+BATCH]
                states = model.forward_states(x[ix])
                pred = model.head(states[:, -1])
                task = ((pred - y[ix]) ** 2).mean()
                G = (states * psi[None, :, None]).sum(1)
                # B_j = - integral psi' z; for one Hann window this is a vector.
                Bv = (states * wp[None, :, None]).sum(1)
                # A convention: dz/dt = A z, so weak prediction is G A^T.
                weak_pred = G @ A.T
                weak = ((Bv - weak_pred) ** 2).mean() / (states.detach().var() + 1e-4)
                loss = task + cfg['alpha'] * weak
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(list(model.parameters()) + [A], 5.0); opt.step()
        model.eval()
        with torch.no_grad():
            metric = float(((model(d['xte'].to(device)) - d['yte'].to(device)) ** 2).mean())
        if return_model: return metric, model, A.detach(), d, device
        return metric
    except RuntimeError:
        # Explicit CPU fallback for a shared or exhausted CUDA slot.
        model = WeakGRU(64)
        x, y = d['xtr'], d['ytr']; opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'], weight_decay=cfg['weight_decay'])
        psi, wp = weak_weights(x.shape[1] // 3); A = nn.Parameter(torch.randn(64,64)*0.01); opt.add_param_group({'params':[A]})
        for _ in range(EPOCHS):
            for i in range(0, len(x), BATCH):
                states=model.forward_states(x[i:i+BATCH]); pred=model.head(states[:,-1]); task=((pred-y[i:i+BATCH])**2).mean()
                G=(states*psi[None,:,None]).sum(1); Bv=(states*wp[None,:,None]).sum(1); weak=((Bv-G@A.T)**2).mean()/(states.detach().var()+1e-4)
                opt.zero_grad(); (task+cfg['alpha']*weak).backward(); opt.step()
        with torch.no_grad(): metric=float(((model(d['xte'])-d['yte'])**2).mean())
        if return_model: return metric, model, A.detach(), d, 'cpu'
        return metric


def signature():
    cfg={'lr':0.003,'weight_decay':0.0,'alpha':0.03}
    metric, model, A, d, dev = idea_one(cfg, 0, True)
    with torch.no_grad():
        states=model.forward_states(d['xte'].to(dev)); psi,wp=weak_weights(states.shape[1]); psi,wp=psi.to(dev),wp.to(dev)
        G=(states*psi[None,:,None]).sum(1); Bv=(states*wp[None,:,None]).sum(1); residual=Bv-G@A.T
        observed=float(residual.pow(2).mean().sqrt()); predicted=float((Bv.detach().var()+1e-8).sqrt())
    return {'model_test_mse':metric, 'weak_residual_rms_observed':observed,
            'weak_scale_predicted_from_observed_B':predicted,
            'relative_residual':observed/(predicted+1e-12),
            'confirmed': bool(observed < predicted)}


def main():
    # Union parity: every idea learning rate is present in baseline sweep.
    grid=[{'lr':lr,'weight_decay':wd} for lr in (0.001,0.003,0.01) for wd in (0.0,1e-4)]
    base=sweep_baseline(lambda c: lambda s: baseline_one(c,s), grid)
    best=base['best_cfg']
    idea_grid=[{'lr':best['lr'],'weight_decay':best['weight_decay'],'alpha':a} for a in (0.01,0.03,0.1)]
    idea_runs=[]
    for c in idea_grid:
        r=evaluate(lambda s, c=c: idea_one(c,s), SEEDS)
        idea_runs.append((c,r))
    idea_cfg, idea=min(idea_runs, key=lambda cr: cr[1]['mean'])
    rep=make_report('dynamics','rnn_small',base,idea,{'weak_koopman':signature(), 'selected_cfg':idea_cfg,
        'baseline_grid':grid, 'idea_grid':idea_grid, 'track_justification':'Actuated pendulum rollout has explicit dynamical stability/control structure.'})
    rep['idea_sweep']=[{'cfg':c,'result':r} for c,r in idea_runs]
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()