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

SEED=2704
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)

def cross_matrix(n):
    x,y,z=n
    return np.array([[0,-z,y],[z,0,-x],[-y,x,0]],float)

def rot(axis, angle):
    axis=np.asarray(axis,float); axis=axis/np.linalg.norm(axis)
    K=cross_matrix(axis)
    return np.eye(3)*math.cos(angle)+(1-math.cos(angle))*np.outer(axis,axis)+math.sin(angle)*K

def spin_seq(events, axes, angles, s0=np.array([0.31,-0.52,0.796])):
    s=np.asarray(s0,float); s=s/np.linalg.norm(s)
    for e in events: s=rot(axes[e],angles[e])@s
    return s

def mechanism_checks():
    # Prediction 1: rotations preserve norm (up to floating point).
    axes=[np.array([1.,0,0]),np.array([0,1.,0])]
    rng=np.random.default_rng(SEED)
    maxerr=0
    expm_err=0
    for _ in range(100):
        a=rng.normal(size=3); a/=np.linalg.norm(a); q=rng.uniform(-math.pi,math.pi)
        expm_err=max(expm_err, np.max(np.abs(rot(a,q)-expm(q*cross_matrix(a)))))
    for _ in range(1000):
        s=rng.normal(size=3); s/=np.linalg.norm(s)
        for e in rng.integers(0,2,100): s=rot(axes[e],rng.uniform(-math.pi,math.pi))@s
        maxerr=max(maxerr,abs(np.linalg.norm(s)-1))
    # Prediction 2: AB-vs-BA difference is zero for parallel axes; for small equal
    # angles orthogonal axes it scales quadratically, with commutator coefficient.
    orth=[]; par=[]
    s0=np.array([.31,-.52,.796]); s0/=np.linalg.norm(s0)
    for th in np.logspace(-3,-0.2,10):
        orth.append(np.linalg.norm(spin_seq([0,1],axes,[th,th],s0)-spin_seq([1,0],axes,[th,th],s0)))
        pa=[np.array([1.,0,0]),np.array([1.,0,0])]
        par.append(np.linalg.norm(spin_seq([0,1],pa,[th,th],s0)-spin_seq([1,0],pa,[th,th],s0)))
    slope=np.polyfit(np.log(np.logspace(-3,-0.2,10)[:5]),np.log(np.array(orth)[:5]),1)[0]
    # Prediction 3: sensitivity grows with nonparallelity, using fixed small angles.
    th=.08; deltas=[]; sins=[]
    for phi in np.linspace(0,.5*math.pi,9):
        a=[np.array([1.,0,0]),np.array([math.cos(phi),math.sin(phi),0.])]
        deltas.append(np.linalg.norm(spin_seq([0,1],a,[th,th],s0)-spin_seq([1,0],a,[th,th],s0)))
        sins.append(abs(math.sin(phi)))
    corr=np.corrcoef(deltas,sins)[0,1]
    # Quantitative predictions: norm error=0; exponent is 2; parallel effect=0;
    # fixed-angle order sensitivity is proportional to sin(axis angle).
    return {'predictions': {
        'rotation_exponential_max_error_predicted_0': expm_err,
        'sphere_norm_error_predicted_0': maxerr,
        'small_angle_order_effect_exponent_predicted_2': slope,
        'parallel_axis_order_effect_predicted_0': max(par),
        'axis_misalignment_dependence_corr_predicted_1': corr
    }, 'orthogonal_delta_at_theta_.08':deltas[-1]}

class CountModel(nn.Module):
    def __init__(self):
        super().__init__(); self.fc=nn.Linear(2,2)
    def forward(self,x): return self.fc(torch.stack([(x==0).float().sum(1),(x==1).float().sum(1)],1))

class GRUModel(nn.Module):
    def __init__(self, spin=False):
        super().__init__(); self.spin=spin; self.gru=nn.GRU(2,16,batch_first=True)
        if spin:
            self.axis=nn.Parameter(torch.tensor([[1.,0,0],[0,1.,0]]))
            self.theta=nn.Parameter(torch.tensor([1.,1.]))
            self.fc=nn.Linear(19,2)
        else: self.fc=nn.Linear(16,2)
    def forward(self,x):
        h,_=self.gru(torch.nn.functional.one_hot(x,2).float()); z=h[:,-1]
        if not self.spin: return self.fc(z)
        axes=self.axis/(self.axis.norm(dim=1,keepdim=True)+1e-8)
        s=torch.zeros(x.size(0),3,device=x.device); s[:,2]=1
        for t in range(x.size(1)):
            n=axes[x[:,t]]; th=torch.tanh(self.theta[x[:,t]])*math.pi
            s=s*torch.cos(th[:,None])+torch.cross(n,s,dim=1)*torch.sin(th[:,None])+n*(n*s).sum(1,keepdim=True)*(1-torch.cos(th[:,None]))
        return self.fc(torch.cat([z,s],1))

def dataset(n, L=64):
    # Equal counts and identical suffix; label is ordering of first two events.
    x=np.tile(np.array([0,1]*((L-2)//2)),(n,1)); x=np.concatenate([np.zeros((n,1),int),x],1)[:,:L]
    # correct dimensions and balanced suffix
    suffix=np.array([0,1]*((L-2)//2),int)
    X=[]; y=[]
    for i in range(n):
        lab=i%2; prefix=[0,1] if lab==0 else [1,0]
        X.append(prefix+suffix.tolist()); y.append(lab)
    return torch.tensor(np.array(X),dtype=torch.long),torch.tensor(y)

def train(model, Xtr,ytr,Xte,yte, epochs=100):
    opt=torch.optim.Adam(model.parameters(),lr=.02); lossfn=nn.CrossEntropyLoss()
    for _ in range(epochs):
        opt.zero_grad(); lossfn(model(Xtr),ytr).backward(); opt.step()
    with torch.no_grad():
        pred=model(Xte).argmax(1); acc=(pred==yte).float().mean().item(); final=lossfn(model(Xte),yte).item()
    return acc,final

def mini_experiment():
    X,y=dataset(512); Xtr,ytr=X[:384],y[:384]; Xte,yte=X[384:],y[384:]
    out={}
    for name,model in [('count',CountModel()),('gru',GRUModel()),('gru_spin',GRUModel(True))]:
        out[name]=train(model,Xtr,ytr,Xte,yte)
    return out

def main():
    checks=mechanism_checks(); exp=mini_experiment()
    result={'seed':SEED,'mechanism_checks':checks,'classification':exp}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()