import json, math, time, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

SEED = 475
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    if device == 'cuda':
        torch.set_default_device('cuda')
        torch.zeros(1, device='cuda')
except Exception:
    device = 'cpu'

# Real CMV-style specialization. M covers (0,1),(2,3),...; L covers
# (1,2),(3,4),...,(n-1,0). Scatter avoids in-place autograd mutations.
def apply_pair_factor(x, starts, theta, period):
    n = x.shape[-1]
    starts = torch.as_tensor(starts, device=x.device, dtype=torch.long)
    ends = (starts + 1) % n
    t = theta[torch.arange(len(starts), device=x.device) % period]
    c, s = torch.cos(t), torch.sin(t)
    u, v = x[:, starts], x[:, ends]
    vals = torch.stack((c[None, :] * u - s[None, :] * v,
                        s[None, :] * u + c[None, :] * v), dim=-1)
    idx = torch.stack((starts, ends), dim=-1).reshape(-1)
    # Each coordinate occurs once for each disjoint factor.
    return torch.zeros_like(x).scatter(1, idx[None, :].expand(x.shape[0], -1),
                                        vals.reshape(x.shape[0], -1))

def cmv_apply(x, theta, period):
    n = x.shape[-1]
    y = apply_pair_factor(x, list(range(0, n, 2)), theta, period)
    return apply_pair_factor(y, list(range(1, n, 2)), theta, period)

def cmv_matrix(n, theta, period):
    return cmv_apply(torch.eye(n, device=theta.device, dtype=theta.dtype), theta, period)

class CMVRNN(nn.Module):
    def __init__(self, n, inp, period):
        super().__init__(); self.n=n
        self.theta=nn.Parameter(torch.empty(period).uniform_(-math.pi, math.pi))
        self.inp=nn.Linear(inp,n); self.out=nn.Linear(n,1)
    def forward(self, x):
        h=torch.zeros(x.shape[0],self.n,device=x.device)
        for t in range(x.shape[1]):
            h=torch.tanh(cmv_apply(h,self.theta,self.theta.numel())+self.inp(x[:,t]))
        return self.out(h).squeeze(-1)

class DenseRNN(nn.Module):
    def __init__(self,n,inp,kind='orthogonal'):
        super().__init__(); self.n=n
        q,_=torch.linalg.qr(torch.randn(n,n)); self.W=nn.Parameter(q if kind=='orthogonal' else 1.05*q)
        self.inp=nn.Linear(inp,n); self.out=nn.Linear(n,1)
    def forward(self,x):
        h=torch.zeros(x.shape[0],self.n,device=x.device)
        for t in range(x.shape[1]): h=torch.tanh(h@self.W.T+self.inp(x[:,t]))
        return self.out(h).squeeze(-1)

def math_check():
    n=32; b=128; x=torch.randn(b,n)
    th=torch.linspace(-1.2,1.1,8,device=x.device)
    U=cmv_matrix(n,th,8)
    err=float((U.T@U-torch.eye(n,device=x.device)).abs().max())
    ratios=[]; v=x
    for _ in range(100):
        v=cmv_apply(v,th,8); ratios.append((v.norm(dim=1)/x.norm(dim=1)).detach().cpu().numpy())
    ratios=np.concatenate(ratios)
    q,_=torch.linalg.qr(torch.randn(n,n,device=x.device)); expansive=1.05*q
    y=x.clone(); dense_rat=[]
    for _ in range(100):
        y=y@expansive.T; dense_rat.append((y.norm(dim=1)/x.norm(dim=1)).cpu().numpy())
    return {'orthogonality_max_error':err,
            'cmv_norm_ratio_max_abs_error':float(np.max(np.abs(ratios-1))),
            'cmv_norm_ratio_mean':float(ratios.mean()),
            'expansive_dense_norm_ratio_at_100':float(np.mean(dense_rat[-1]))}

def train(model, steps=120, seq=40, batch=64):
    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    lossfn=nn.BCEWithLogitsLoss(); vals=[]; start=time.time()
    for step in range(steps):
        x=torch.randn(batch,seq,1,device=device); y=(x.sum(1).squeeze(-1)>0).float()
        opt.zero_grad(set_to_none=True); loss=lossfn(model(x),y); loss.backward(); opt.step()
        if step%20==0 or step==steps-1:
            with torch.no_grad():
                xx=torch.randn(512,seq,1,device=device); yy=(xx.sum(1).squeeze(-1)>0).float()
                acc=((model(xx)>0)==(yy>0.5)).float().mean().item()
            vals.append((step,float(loss),acc))
    return {'final_loss':vals[-1][1],'final_acc':vals[-1][2],
            'trace':vals,'seconds':time.time()-start}

def main():
    torch.set_num_threads(4)
    check=math_check(); n=32; inp=1
    models={'dense_orthogonal':DenseRNN(n,inp,'orthogonal'),
            'cmv_period_8':CMVRNN(n,inp,8), 'cmv_period_32':CMVRNN(n,inp,32)}
    results={}
    for name,m in models.items():
        results[name]={'parameters':sum(p.numel() for p in m.parameters()), **train(m)}
    out={'seed':SEED,'device':device,'math_check':check,'models':results}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()