import json, math, random
from pathlib import Path
import numpy as np
from scipy.linalg import expm, norm, eigvals
import torch
import torch.nn as nn

SEED = 449
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def bch_check():
    # For two generators, the residual after retaining the commutator term
    # should scale cubically with the step size.
    rng = np.random.default_rng(SEED)
    X = rng.normal(size=(3,3)); Y = rng.normal(size=(3,3))
    X -= 1.7*np.eye(3); Y -= 1.7*np.eye(3)
    C = X @ Y - Y @ X
    rows = []
    for s in [0.02, 0.04, 0.08, 0.16, 0.32]:
        exact = expm(s*Y) @ expm(s*X)
        avg = expm(s*(X+Y) - 0.5*s*s*C)
        first = expm(s*(X+Y))
        rows.append((s, norm(exact-first, 'fro'), norm(exact-avg, 'fro')))
    x = np.log([r[0] for r in rows[-3:]])
    y1 = np.log([r[1] for r in rows[-3:]])
    y2 = np.log([r[2] for r in rows[-3:]])
    return {"rows": [[float(v) for v in r] for r in rows],
            "first_order_loglog_slope": float(np.polyfit(x,y1,1)[0]),
            "commutator_corrected_loglog_slope": float(np.polyfit(x,y2,1)[0])}


def matrix_metrics(A, tau=0.7, alpha=(0.5,0.5)):
    B = [alpha[i]*tau*A[i] for i in range(2)]
    Phi = expm(B[1]) @ expm(B[0])
    avg = expm(B[0]+B[1])
    comm = norm(B[0]@B[1]-B[1]@B[0], 'fro')**2
    mismatch = norm(Phi-avg, 'fro') / max(norm(avg,'fro'), 1e-12)
    return float(comm), float(mismatch), Phi


def regularizer_optimization():
    # Directly optimize two matrices toward a fixed noncommuting task matrix,
    # with and without the proposed commutator penalty.
    torch.manual_seed(SEED+1)
    target = torch.tensor(np.array([[.72,.35],[-.15,.61]], dtype=np.float32))
    out = {}
    for name, lam in [("unregularized",0.0),("commutator_regularized",2.0)]:
        A = nn.Parameter(torch.randn(2,2)*.7)
        D = nn.Parameter(torch.randn(2,2)*.7)
        opt = torch.optim.Adam([A,D], lr=.04)
        for _ in range(500):
            B0=.35*A; B1=.35*D
            phi=torch.matrix_exp(B1)@torch.matrix_exp(B0)
            loss=(phi-target).pow(2).mean()+lam*(B0@B1-B1@B0).pow(2).sum()
            opt.zero_grad(); loss.backward(); opt.step()
        a=A.detach().numpy(); d=D.detach().numpy()
        comm, mismatch, _=matrix_metrics([a,d],tau=.7)
        fit=float(((expm(.35*d)@expm(.35*a))-target.numpy())**2 .mean()) if False else float(np.mean((expm(.35*d)@expm(.35*a)-target.numpy())**2))
        out[name]={"task_fit_mse":fit,"commutator":comm,"ordered_averaged_relative_error":mismatch}
    return out


class SwitchedSSM(nn.Module):
    def __init__(self, d=8, modes=2, regularized=False, single=False):
        super().__init__(); self.d=d; self.m=modes; self.regularized=regularized; self.single=single
        n=1 if single else modes
        self.A=nn.Parameter(torch.randn(n,d,d)*.08 - .35*torch.eye(d).unsqueeze(0).repeat(n,1,1))
        self.U=nn.Parameter(torch.randn(modes,d)*.12)
        self.v=nn.Parameter(torch.randn(d)*.12)
        self.readout=nn.Linear(d,1)
    def forward(self,x):
        h=torch.zeros(x.shape[0],self.d,device=x.device)
        alpha=1.0/self.m; tau=.7
        E=[]
        for i in range(self.m):
            A=self.A[0] if self.single else self.A[i]
            E.append(torch.matrix_exp(alpha*tau*A))
        for t in range(x.shape[1]):
            for i in range(self.m):
                h=(h.unsqueeze(1) @ E[i].T).squeeze(1)
                h=h+alpha*x[:,t:t+1]*self.U[i]
        return self.readout(h).squeeze(-1)
    def penalties(self):
        if self.single: return torch.tensor(0.,device=self.A.device)
        B=[.5*.7*a for a in self.A]
        return sum(((B[i]@B[j]-B[j]@B[i])**2).sum() for i in range(self.m) for j in range(i))
    def diagnostics(self):
        with torch.no_grad():
            As=self.A.detach().cpu().numpy(); As=As if not self.single else np.repeat(As,2,axis=0)
            comm, mismatch, phi=matrix_metrics(As,tau=.7)
            gain=max(abs(eigvals(phi)))
        return {"commutator":comm,"ordered_averaged_relative_error":mismatch,"cycle_spectral_gain":float(gain)}


def sequence_experiment():
    torch.manual_seed(SEED+2)
    dev=DEVICE
    ntrain, ntest, T=192,96,18
    rng=np.random.default_rng(SEED+2)
    X=rng.normal(size=(ntest+ntrain,T)).astype('float32')
    # Stable long-memory target: exponentially discounted sequence sum.
    weights=(.88**np.arange(T-1,-1,-1)).astype('float32')
    y=(X*weights).sum(1).astype('float32')
    xt=torch.tensor(X[:ntrain],device=dev); yt=torch.tensor(y[:ntrain],device=dev)
    xe=torch.tensor(X[ntrain:],device=dev); ye=torch.tensor(y[ntrain:],device=dev)
    results={}
    for name, reg, single in [("single_generator",0.,True),("switched_unregularized",0.,False),("switched_commutator_regularized",.08,False)]:
        torch.manual_seed(SEED+10+len(results)); model=SwitchedSSM(8,2,reg>0,single).to(dev)
        opt=torch.optim.Adam(model.parameters(),lr=.025)
        for step in range(100):
            pred=model(xt); loss=((pred-yt)**2).mean()+reg*model.penalties()
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5.0); opt.step()
        with torch.no_grad(): val=float(((model(xe)-ye)**2).mean().cpu())
        results[name]={"validation_mse":val,**model.diagnostics()}
    return results


def main():
    global DEVICE
    try:
        b=bch_check(); ro=regularizer_optimization(); seq=sequence_experiment()
    except Exception as e:
        # CUDA failures are explicitly allowed to fall back to CPU.
        if DEVICE=="cuda":
            DEVICE="cpu"; torch.cuda.empty_cache(); b=bch_check(); ro=regularizer_optimization(); seq=sequence_experiment()
        else: raise
    result={"device":DEVICE,"bch":b,"regularizer_optimization":ro,"sequence_experiment":seq}
    Path("results.json").write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))
if __name__=="__main__": main()