import json, math, random
import numpy as np
import torch
from torch import nn

SEED = 17
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.zeros(1, device='cuda')
except Exception:
    device = 'cpu'

torch.set_num_threads(4)

def orbit_tangent(x):
    return torch.stack((-x[:, 1], x[:, 0]), dim=1)

def projector(R, eps=1e-6):
    # R is [batch,n,r]; G=I. Handles rank-deficient R through damping.
    b, n, r = R.shape
    eye = torch.eye(r, device=R.device, dtype=R.dtype).expand(b, r, r)
    return R @ torch.linalg.solve(R.transpose(1,2) @ R + eps*eye,
                                   R.transpose(1,2))

def verify_math():
    x = torch.tensor([[1.2, -0.7]], dtype=torch.float64)
    r = orbit_tangent(x).unsqueeze(-1)
    P = projector(r, 1e-10)[0]
    rv = r[0,:,0:1]
    radial = x[0,:,None]
    z = torch.zeros_like(P)
    origin = torch.zeros(1,2,1,dtype=torch.float64)
    Po = projector(origin, 1e-6)[0]
    return {
        'projector_idempotence': float(torch.linalg.norm(P@P-P)),
        'orbit_projection_error': float(torch.linalg.norm(P@rv-rv)),
        'normal_projection_norm': float(torch.linalg.norm(P@radial)),
        'rank_regular': int(torch.linalg.matrix_rank(r[0])),
        'rank_singular_at_origin': int(torch.linalg.matrix_rank(origin[0])),
        'singular_projector_norm': float(torch.linalg.norm(Po))
    }

def rotate(x, theta):
    c, s = torch.cos(theta), torch.sin(theta)
    return torch.stack((c*x[:,0]-s*x[:,1], s*x[:,0]+c*x[:,1]), 1)

def make_data(n, seed, radius_noise=.10):
    gen = torch.Generator().manual_seed(seed)
    y = torch.randint(0,2,(n,),generator=gen)
    radius = (1.0 + (2*y.float()-1)*.22 + radius_noise*torch.randn(n,generator=gen))
    theta = 2*math.pi*torch.rand(n,generator=gen)
    x = torch.stack((radius*torch.cos(theta), radius*torch.sin(theta)),1)
    return x, y

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.body=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,16),nn.Tanh())
        self.head=nn.Linear(16,2)
    def forward(self,x): return self.head(self.body(x))

def orbit_penalty(model, x):
    # Exact JVP of f in the supplied orbit direction, without materializing J.
    x = x.detach().requires_grad_(True)
    v = orbit_tangent(x)
    h = model(x)
    # Use a per-output VJP loop to obtain exact Jv without materializing J.
    # to obtain exact Jv for the 2-dimensional representation.
    cols=[]
    for k in range(h.shape[1]):
        g=torch.autograd.grad(h[:,k].sum(), x, create_graph=True, retain_graph=True)[0]
        cols.append((g*v).sum(1))
    jv=torch.stack(cols,1)
    return (jv*jv).sum(1).mean()

def train(reg, seed=17, steps=350):
    torch.manual_seed(seed)
    model=Encoder().to(device)
    opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    x,y=make_data(768,seed)
    x,y=x.to(device),y.to(device)
    for step in range(steps):
        ix=torch.randint(0,len(x),(96,),device=device)
        xb,yb=x[ix],y[ix]
        logits=model(xb)
        loss=nn.functional.cross_entropy(logits,yb)
        if reg: loss=loss + .15*orbit_penalty(model,xb)
        opt.zero_grad(); loss.backward(); opt.step()
    return model

@torch.no_grad()
def evaluate(model, seed=101):
    x,y=make_data(3000,seed)
    x,y=x.to(device),y.to(device)
    clean=(model(x).argmax(1)==y).float().mean().item()
    # unseen rotations, including angles outside a narrow training orientation
    th=(torch.rand(len(x),device=device)*2-1)*math.pi
    xr=rotate(x,th)
    robust=(model(xr).argmax(1)==y).float().mean().item()
    drift=(model(xr)-model(x)).norm(dim=1).mean().item()
    return clean,robust,drift

def main():
    check=verify_math()
    base=train(False); idea=train(True)
    b=evaluate(base); q=evaluate(idea)
    out={'device':device,'math_check':check,
         'baseline':{'clean_accuracy':b[0],'rotated_accuracy':b[1],'representation_drift':b[2]},
         'quotient_regularizer':{'clean_accuracy':q[0],'rotated_accuracy':q[1],'representation_drift':q[2]}}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()