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

SEED=1234
rng=np.random.default_rng(SEED)
torch.manual_seed(SEED)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type=='cuda':
        torch.empty(1, device=device)
except Exception:
    device=torch.device('cpu')


def random_rotations(n, rng):
    # Haar rotations from QR, with proper determinant
    a=rng.normal(size=(n,3,3)); out=[]
    for x in a:
        q,r=np.linalg.qr(x); q=q@np.diag(np.sign(np.diag(r)))
        if np.linalg.det(q)<0: q[:,0]*=-1
        out.append(q)
    return np.asarray(out)

def lattice_offsets(shells):
    # shells are squared Euclidean distances on a cubic periodic lattice
    alloff=[]
    for d in range(1,shells+1):
        off=[]
        for x in range(-d,d+1):
            for y in range(-d,d+1):
                for z in range(-d,d+1):
                    if (x,y,z)!=(0,0,0) and x*x+y*y+z*z==d*d:
                        off.append((x,y,z))
        alloff.append(off)
    return alloff

def p2(x): return 0.5*(3*x*x-1)

def corr_features(u,L,shells=3):
    # u: [N,3], row-major cubic lattice
    a=u.reshape(L,L,L,3); chans=[]
    for offsets in lattice_offsets(shells):
        c=np.zeros((L,L,L))
        for dx,dy,dz in offsets:
            shifted=np.roll(a,(dx,dy,dz),(0,1,2))
            c += p2(np.sum(a*shifted,axis=-1))
        c /= len(offsets)
        chans.append(c)
    return np.stack(chans,axis=0).astype(np.float32).reshape(-1)

def S_order(u):
    Q=np.einsum('ni,nj->ij',u,u)/len(u)
    Q=0.5*(3*Q-np.eye(3))
    return float(np.linalg.eigvalsh(Q)[-1])

def make_config(L, phase, rng):
    n=L**3
    if phase=='iso':
        u=rng.normal(size=(n,3)); u/=np.linalg.norm(u,axis=1,keepdims=True)
    else:
        d=rng.normal(size=3); d/=np.linalg.norm(d)
        # Gaussian tangent noise gives a controllable nematic cloud
        u=d+rng.normal(scale=.23,size=(n,3)); u/=np.linalg.norm(u,axis=1,keepdims=True)
        # random apolar signs make the physical representation genuinely apolar
        u*=rng.choice([-1.,1.],size=(n,1))
    R=random_rotations(1,rng)[0]
    return u@R.T

def transform(u,R,flip_prob,rng):
    v=u@R.T
    signs=np.where(rng.random(len(u))<flip_prob,-1.,1.)[:,None]
    return v*signs

class AE(nn.Module):
    def __init__(self, dim, zdim=2):
        super().__init__(); self.enc=nn.Sequential(nn.Linear(dim,64),nn.Tanh(),nn.Linear(64,zdim)); self.dec=nn.Sequential(nn.Linear(zdim,64),nn.Tanh(),nn.Linear(64,dim))
    def forward(self,x):
        z=self.enc(x); return self.dec(z),z

def train_ae(X, epochs=100):
    model=AE(X.shape[1]).to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    xt=torch.tensor(X,device=device)
    for _ in range(epochs):
        opt.zero_grad(); rec,_=model(xt); loss=((rec-xt)**2).mean(); loss.backward(); opt.step()
    with torch.no_grad():
        rec_t,z_t=model(xt); z=z_t.cpu().numpy(); rec=rec_t.cpu().numpy()
    return float(np.mean((rec-X)**2)), z

def abs_corr(z,s):
    return [float(abs(np.corrcoef(z[:,k],s)[0,1])) for k in range(z.shape[1])]

def main():
    L=5; n=2*120; us=[]; labels=[]; ss=[]
    for phase in ['iso','nematic']:
        for _ in range(120):
            u=make_config(L,phase,rng); us.append(u); labels.append(phase=='nematic'); ss.append(S_order(u))
    us=np.asarray(us); labels=np.asarray(labels); ss=np.asarray(ss)
    # Core prediction 1: common rotations and independent apolar flips leave c exactly unchanged.
    base=us[0]; inv=[]; raw=[]
    for p in [0,.25,.5,.75,1.0]:
        errs=[]; rawerrs=[]
        for _ in range(20):
            v=transform(base,random_rotations(1,rng)[0],p,rng)
            errs.append(np.linalg.norm(corr_features(base,L,3)-corr_features(v,L,3))/math.sqrt(len(corr_features(base,L,3))))
            rawerrs.append(np.linalg.norm(base-v)/math.sqrt(base.size))
        inv.append({'flip_probability':p,'observed_rms_feature_error':float(np.mean(errs)),'raw_rms_vector_error':float(np.mean(rawerrs))})
    # Core prediction 2: invariant error is independent of shell count (all exact pairwise dot products).
    shell=[]
    for R in [1,2,3,4]:
        vals=[]
        for _ in range(20):
            v=transform(base,random_rotations(1,rng)[0],rng.random(),rng)
            a=corr_features(base,L,R); b=corr_features(v,L,R)
            vals.append(np.max(np.abs(a-b)))
        shell.append({'shells':R,'observed_max_abs_error':float(max(vals)),'predicted':0.0})
    # A direct P2 sweep quantifies evenness: sign of dot product never matters.
    x=np.linspace(-1,1,10001); even_err=float(np.max(np.abs(p2(x)-p2(-x))))
    # Mini phase experiment: raw vectors versus invariant local channels.
    rawX=us.reshape(len(us),-1).astype(np.float32)
    invX=np.asarray([corr_features(u,L,3) for u in us])
    raw_mse,zraw=train_ae(rawX); inv_mse,zinv=train_ae(invX)
    result={'seed':SEED,'device':str(device),'L':L,'n_samples':len(us),
      'math_predictions':{
       'rotation_and_sign_invariance':'feature RMS error should be numerical zero for every flip probability',
       'shell_invariance':'max absolute error should remain numerical zero as shells increase',
       'P2_evenness':'P2(x)-P2(-x) should be exactly zero'},
      'invariance_sweep':inv,'shell_sweep':shell,'p2_even_max_error':even_err,
      'phase_experiment':{'raw_reconstruction_mse':raw_mse,'invariant_reconstruction_mse':inv_mse,
        'raw_abs_latent_S_correlations':abs_corr(zraw,ss),'invariant_abs_latent_S_correlations':abs_corr(zinv,ss),
        'raw_best_abs_corr':max(abs_corr(zraw,ss)),'invariant_best_abs_corr':max(abs_corr(zinv,ss)),
        'mean_S_iso':float(ss[~labels].mean()),'mean_S_nematic':float(ss[labels].mean())}}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()