import json, random, 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, sweep_baseline, make_report

TAU = 0.20
ALPHA = (0.5, 0.5)
HIDDEN = 64
SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))

class SwitchedSSM(nn.Module):
    def __init__(self, input_shape, out_dim, comm_lambda=0.0):
        super().__init__()
        self.d = HIDDEN
        self.lam = float(comm_lambda)
        eye = torch.eye(self.d)
        self.A = nn.Parameter(-0.7 * eye[None].repeat(2, 1, 1) + 0.03 * torch.randn(2, self.d, self.d))
        self.U = nn.Parameter(0.08 * torch.randn(2, 3, self.d))
        self.head = nn.Linear(self.d, out_dim)

    def forward(self, x):
        z = x.reshape(x.shape[0], -1, 3)
        h = torch.zeros(x.shape[0], self.d, device=x.device, dtype=x.dtype)
        for t in range(z.shape[1]):
            for i, a in enumerate(ALPHA):
                E = torch.matrix_exp(a * TAU * self.A[i])
                h = h @ E.T + (a * z[:, t, :]) @ self.U[i]
        return self.head(h)

    def comm_penalty(self):
        B0 = ALPHA[0] * TAU * self.A[0]
        B1 = ALPHA[1] * TAU * self.A[1]
        C = B0 @ B1 - B1 @ B0
        return (C * C).sum()

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)

def train_idea(model, ds, epochs, lr):
    try:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
        model = model.to(device)
        x, y = ds['xtr'].to(device), ds['ytr'].to(device)
        opt = torch.optim.Adam(model.parameters(), lr=lr)
        lossf = nn.MSELoss()
        for _ in range(epochs):
            model.train()
            for ix in torch.randperm(len(x), device=device).split(128):
                loss = lossf(model(x[ix]), y[ix]) + model.lam * model.comm_penalty()
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
        model.eval()
        with torch.no_grad(): metric = float(lossf(model(ds['xte'].to(device)), ds['yte'].to(device)).cpu())
        return model, metric
    except Exception:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        model = model.cpu(); x, y = ds['xtr'], ds['ytr']; opt = torch.optim.Adam(model.parameters(), lr=lr); lossf = nn.MSELoss()
        for _ in range(epochs):
            for ix in torch.randperm(len(x)).split(128):
                loss = lossf(model(x[ix]), y[ix]) + model.lam * model.comm_penalty()
                opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
        with torch.no_grad(): metric = float(lossf(model(ds['xte']), ds['yte']))
        return model, metric

def train_base(cfg, seed, keep=False):
    seed_all(seed); ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
    model = SwitchedSSM(ds['input_shape'], ds['out_dim'], 0.0)
    # Uses the standard harness path for the matched no-penalty switched baseline.
    model, metric, _ = train_model(model, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128, weight_decay=0.0)
    return (model, metric) if keep else metric

def train_idea_cfg(cfg, seed, keep=False):
    seed_all(seed); ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
    model = SwitchedSSM(ds['input_shape'], ds['out_dim'], cfg['comm_lambda'])
    model, metric = train_idea(model, ds, cfg['epochs'], cfg['lr'])
    return (model, metric) if keep else metric

def base_factory(cfg):
    return lambda seed: train_base(cfg, seed)

def metric_and_models(fn, cfg):
    vals=[]; models=[]
    for s in SEEDS:
        m,v = fn(cfg, s, True); vals.append(v); models.append(m)
    return {'config':cfg, 'per_seed':vals, 'mean':float(np.mean(vals)), 'models':models}

def diagnostics(models):
    comm=[]; gain=[]
    for m in models:
        with torch.no_grad():
            A=m.A.detach().cpu(); B0=ALPHA[0]*TAU*A[0]; B1=ALPHA[1]*TAU*A[1]
            C=B0@B1-B1@B0; Phi=torch.matrix_exp(B1)@torch.matrix_exp(B0)
            comm.append(float((C*C).sum())); gain.append(float(torch.linalg.svdvals(Phi).max()))
    return {'observed_commutator_mean':float(np.mean(comm)), 'observed_cycle_gain_mean':float(np.mean(gain)), 'n_models':len(models)}

def main():
    # Equal union: every idea lr is included in baseline tuning.
    grid=[{'lr':lr,'epochs':20} for lr in (0.003,0.006,0.009)]
    tuned=sweep_baseline(base_factory, grid, seeds=SWEEP_SEEDS)
    best_cfg=tuned['best_cfg']
    base_full=metric_and_models(lambda c,s,k: train_base(c,s,k), best_cfg)
    idea_runs=[]
    for lam in (0.0, 0.01, 0.05):
        cfg=dict(best_cfg, comm_lambda=lam)
        idea_runs.append(metric_and_models(lambda c,s,k: train_idea_cfg(c,s,k), cfg))
    idea=min(idea_runs, key=lambda r:r['mean'])
    bd=diagnostics(base_full['models']); id=diagnostics(idea['models'])
    report=make_report('dynamics','rnn_small',
        {'best_cfg':best_cfg, 'sweep':tuned['sweep'], 'full':{'per_seed':base_full['per_seed'],'mean':base_full['mean']}},
        {'config':idea['config'],'per_seed':idea['per_seed'],'mean':idea['mean']},
        {'track_justification':'Dynamics directly tests switched latent stability/control and Lyapunov contraction.',
         'mechanism_signature':{'prediction':'commutator regularization lowers trained ordered-cycle noncommutativity and cycle gain',
             'baseline_observed':bd, 'idea_observed':id,
             'commutator_ratio':id['observed_commutator_mean']/max(bd['observed_commutator_mean'],1e-12),
             'cycle_gain_delta':id['observed_cycle_gain_mean']-bd['observed_cycle_gain_mean'],
             'confirmed': bool(id['observed_commutator_mean'] < bd['observed_commutator_mean'] and id['observed_cycle_gain_mean'] <= bd['observed_cycle_gain_mean']*1.02)}})
    report['idea_sweep']=[{k:r[k] for k in ('config','mean','per_seed')} for r in idea_runs]
    Path('bench_report.json').write_text(json.dumps(report, indent=2, allow_nan=False))
    print(json.dumps(report, indent=2))

if __name__=='__main__': main()