import json, time, random
import numpy as np
import torch
import torch.nn as nn

SEED = 388
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
try:
    if device == 'cuda': torch.cuda.empty_cache()
except Exception:
    device = 'cpu'

def antihermitian(W):
    # i times a real symmetric matrix is anti-Hermitian
    S = (W + W.transpose(-1, -2)) / 2
    return 1j * S.to(torch.complex64)

def target_map(x, h0):
    n = h0.numel()
    H1 = antihermitian(torch.arange(n*n, device=x.device, dtype=torch.float32).reshape(n,n)/17)
    H2 = antihermitian(torch.flip(torch.arange(n*n, device=x.device, dtype=torch.float32), (0,)).reshape(n,n)/19)
    out = h0[None].expand(x.shape[0], -1).to(torch.complex64)
    for c, H in [(0.55*torch.sin(x[:,0]), H1), (0.4*torch.cos(x[:,1]), H2)]:
        E = torch.matrix_exp(c[:,None,None] * H[None])
        out = torch.bmm(E, out[...,None])[...,0]
    return out

class StructuredKLU(nn.Module):
    def __init__(self, d=3, n=4, K=4):
        super().__init__(); self.d=d; self.n=n; self.K=K
        self.H = nn.Parameter(torch.randn(K,n,n)*0.12)
        self.J = nn.Parameter(torch.randn(K,n,n)*0.12)
        self.psi = nn.ModuleList([nn.ModuleList([nn.Sequential(nn.Linear(1,8),nn.Tanh(),nn.Linear(8,1)) for _ in range(d)]) for _ in range(K)])
        self.coeff = nn.ModuleList([nn.Sequential(nn.Linear(1,8),nn.Tanh(),nn.Linear(8,2),nn.Tanh()) for _ in range(K)])
    def forward(self, x, h, return_u=False):
        B=x.shape[0]; out=h.to(torch.complex64)
        U=torch.eye(self.n, device=x.device, dtype=torch.complex64)[None].expand(B,-1,-1).clone()
        for k in range(self.K):
            s=sum(self.psi[k][j](x[:,j:j+1]) for j in range(self.d))
            ab=self.coeff[k](s)
            G=ab[:,0,None,None]*antihermitian(self.H[k])[None] + ab[:,1,None,None]*antihermitian(self.J[k])[None]
            E=torch.matrix_exp(G)
            out=torch.bmm(E,out[...,None])[...,0]
            U=torch.bmm(E,U)
        return (out,U) if return_u else out

class DenseController(nn.Module):
    def __init__(self, d=3, n=4, K=4):
        super().__init__(); self.n=n; self.K=K
        self.net=nn.Sequential(nn.Linear(d,32),nn.Tanh(),nn.Linear(32,2*K*n*n))
    def forward(self,x,h,return_u=False):
        B=x.shape[0]; q=self.net(x).reshape(B,self.K,2,self.n,self.n)
        out=h.to(torch.complex64); U=torch.eye(self.n,device=x.device,dtype=torch.complex64)[None].expand(B,-1,-1).clone()
        for k in range(self.K):
            G=antihermitian(q[:,k,0])+1j*(q[:,k,1]-q[:,k,1].transpose(-1,-2))/2
            E=torch.matrix_exp(G*0.06)
            out=torch.bmm(E,out[...,None])[...,0]; U=torch.bmm(E,U)
        return (out,U) if return_u else out

def params(m): return sum(p.numel() for p in m.parameters() if p.requires_grad)

def main():
    n,d,K=4,3,4; Ntr,Nte=768,256
    g=torch.Generator(device='cpu').manual_seed(SEED)
    xtr=(torch.rand(Ntr,d,generator=g)*2-1).to(device); xte=(torch.rand(Nte,d,generator=g)*2-1).to(device)
    h0=torch.tensor([1+0j,.2+.3j,-.4+.1j,.5-.2j],device=device)
    ytr=target_map(xtr,h0).detach(); yte=target_map(xte,h0).detach()
    # independent mathematical sanity check, including a longer product
    W=torch.randn(32,4,4,device=device); G=antihermitian(W); E=torch.matrix_exp(G)
    one=torch.eye(4,device=device,dtype=torch.complex64)
    single_res=(E.conj().transpose(-1,-2)@E-one).norm().item()
    P=one
    for i in range(32): P=torch.matrix_exp(G[i])@P
    product_res=(P.conj().transpose(-1,-2)@P-one).norm().item()
    results={'device':device,'math_single_residual':single_res,'math_32_product_residual':product_res}
    for name, cls in [('KLU',StructuredKLU),('dense',DenseController)]:
        torch.manual_seed(SEED+ (0 if name=='KLU' else 1))
        model=cls(d,n,K).to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
        t0=time.perf_counter(); losses=[]; gradvals=[]
        for step in range(301):
            ix=torch.randint(0,Ntr,(64,),device=device)
            pred=model(xtr[ix],h0.expand(64,-1))
            loss=((pred-ytr[ix]).abs()**2).mean()
            opt.zero_grad(); loss.backward(); gradvals.append(float(torch.nn.utils.clip_grad_norm_(model.parameters(),100))); opt.step()
            if step in (0,100,300): losses.append(float(loss.detach().cpu()))
        with torch.no_grad():
            pred,U=model(xte,h0.expand(Nte,-1),return_u=True)
            test=float(((pred-yte).abs()**2).mean().cpu())
            norm_drift=float((pred.abs().norm(dim=1)-h0.abs().norm()).abs().mean().cpu())
            unit=float(((U.conj().transpose(-1,-2)@U-torch.eye(n,device=device,dtype=torch.complex64)).norm(dim=(1,2))).mean().cpu())
        results[name]={'parameters':params(model),'train_loss_0_100_300':losses,'test_mse':test,'norm_drift':norm_drift,'unitarity_residual':unit,'seconds':time.perf_counter()-t0,'mean_grad_norm':float(np.mean(gradvals))}
    with open('results.json','w') as f: json.dump(results,f,indent=2)
    print(json.dumps(results,indent=2))
if __name__=='__main__': main()