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

SEED = 428
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device.type == "cuda":
        torch.cuda.empty_cache()
except Exception:
    device = torch.device("cpu")


def orthogonal(n, seed):
    rng = np.random.default_rng(seed)
    q, r = np.linalg.qr(rng.normal(size=(n, n)))
    s = np.sign(np.diag(r)); s[s == 0] = 1
    return q @ np.diag(s)


def math_check(d=5, r=3, n_power=12):
    q = orthogonal(d+r, SEED)
    D = np.diag(np.r_[np.ones(d), -np.ones(r)])
    up, um = q, q @ D
    T, F = up[:d, :d], up[:d, d:]
    G, H = up[d:, :d], up[d:, d:]
    block_minus = np.block([[T, -F], [G, -H]])
    checks = {
        "Uplus_orth_error": float(np.linalg.norm(up.T @ up - np.eye(d+r))),
        "Uminus_orth_error": float(np.linalg.norm(um.T @ um - np.eye(d+r))),
        "block_sign_construction_error": float(np.linalg.norm(um-block_minus)),
        "difference_structure_error": float(np.linalg.norm((up-um)-np.block([[np.zeros((d,d)),2*F],[np.zeros((r,d)),2*H]]))),
        "power_difference_identity_error": 0.0,
        "max_norm_error_over_powers": 0.0,
    }
    diff = np.zeros_like(up)
    left = np.eye(d+r)
    right = np.eye(d+r)
    for k in range(n_power):
        # identity for N=k: U+^(k+1)-U-^(k+1)
        lhs = np.linalg.matrix_power(up,k+1)-np.linalg.matrix_power(um,k+1)
        rhs = up @ (np.linalg.matrix_power(up,k)-np.linalg.matrix_power(um,k)) + (up-um) @ np.linalg.matrix_power(um,k)
        checks["power_difference_identity_error"] = max(checks["power_difference_identity_error"], float(np.linalg.norm(lhs-rhs)))
        v = np.random.default_rng(SEED+k).normal(size=d+r)
        for mat in (up, um):
            checks["max_norm_error_over_powers"] = max(checks["max_norm_error_over_powers"], abs(np.linalg.norm(np.linalg.matrix_power(mat,k) @ v)-np.linalg.norm(v)))
    return checks

class TanhRNN(nn.Module):
    def __init__(self, h):
        super().__init__(); self.h=h
        self.A=nn.Parameter(torch.randn(h,h)*0.12); self.B=nn.Parameter(torch.randn(h,1)*0.12); self.C=nn.Parameter(torch.randn(1,h)*0.12); self.bias=nn.Parameter(torch.zeros(h))
    def forward(self, x, return_norms=False):
        z=torch.zeros(x.size(0),self.h,device=x.device); norms=[]
        for t in range(x.size(1)):
            z=torch.tanh(z@self.A.T + x[:,t:t+1]@self.B.T + self.bias); norms.append(z.norm(dim=1).mean())
        return (z@self.C.T).squeeze(-1), torch.stack(norms) if return_norms else None

class FixedUnitary(nn.Module):
    def __init__(self, U, h):
        super().__init__(); self.register_buffer('U',torch.tensor(U,dtype=torch.float32)); self.B=nn.Parameter(torch.randn(h,1)*0.12); self.C=nn.Parameter(torch.randn(1,h)*0.12)
    def forward(self,x,return_norms=False):
        z=torch.zeros(x.size(0),self.U.size(0),device=x.device); norms=[]
        for t in range(x.size(1)):
            z=z@self.U.T+x[:,t:t+1]@self.B.T; norms.append(z.norm(dim=1).mean())
        return (z@self.C.T).squeeze(-1), torch.stack(norms) if return_norms else None

def make_data(batch, seq, device, generator):
    # The first token carries a random +/- bit; subsequent tokens are distractors.
    y=torch.randint(0,2,(batch,),device=device,generator=generator).float()*2-1
    x=torch.randn(batch,seq,1,device=device,generator=generator)*0.35
    x[:,0,0]=y
    return x,y

def train_eval(model, seq=40, steps=350, h=32):
    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    g=torch.Generator(device=device); g.manual_seed(SEED+77)
    losses=[]
    for _ in range(steps):
        x,y=make_data(48,seq,device,g); pred,_=model(x)
        loss=torch.nn.functional.softplus(-pred*y).mean(); opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10); opt.step(); losses.append(float(loss.detach().cpu()))
    with torch.no_grad():
        x,y=make_data(512,seq,device,g); pred,norms=model(x,True); acc=float(((pred>0)==(y>0)).float().mean().cpu()); test_loss=float(torch.nn.functional.softplus(-pred*y).mean().cpu()); norm_last=float(norms[-1].cpu()); norm_max=float(norms.max().cpu())
    # Differentiable probe is deliberately separate from no_grad evaluation.
    x,y=make_data(32,seq,device,g); x.requires_grad_(True); pred,_=model(x); grad=torch.autograd.grad(pred.mean(),x,retain_graph=False)[0]; gfirst=float(grad[:,0].norm(dim=1).mean().cpu()); glast=float(grad[:,-1].norm(dim=1).mean().cpu())
    return {"final_train_loss":float(np.mean(losses[-25:])),"test_loss":test_loss,"accuracy":acc,"hidden_norm_last":norm_last,"hidden_norm_max":norm_max,"input_grad_first":gfirst,"input_grad_last":glast}

def main():
    global device
    h=32; q=orthogonal(h,SEED+1); d=h//2; D=np.diag(np.r_[np.ones(d),-np.ones(h-d)])
    t=time.time(); math_results=math_check()
    def run():
        results={"device":str(device),"math":math_results,"experiment":{}}
        results["experiment"]["tanh_rnn"]=train_eval(TanhRNN(h))
        results["experiment"]["single_unitary"]=train_eval(FixedUnitary(q,h))
        results["experiment"]["dual_unitary"]=train_eval(FixedUnitary(q@D,h))
        results["elapsed_sec"]=time.time()-t
        return results
    try:
        results=run()
    except (RuntimeError, torch.cuda.OutOfMemoryError) as exc:
        if device.type != "cuda":
            raise
        print("CUDA failed; retrying on CPU:", repr(exc))
        try: torch.cuda.empty_cache()
        except Exception: pass
        device=torch.device("cpu")
        results=run()
    with open("results.json","w") as f: json.dump(results,f,indent=2)
    print(json.dumps(results,indent=2))

if __name__=="__main__": main()