import json, random, time
import numpy as np
import torch
from torch import nn
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

SEED = 530
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)


def grid_graph(h, w):
    n = h*w
    edges = []
    for r in range(h):
        for c in range(w):
            i = r*w+c
            if r+1 < h: edges.append((i, (r+1)*w+c))
            if c+1 < w: edges.append((i, r*w+c+1))
    B = np.zeros((n, len(edges)), dtype=np.float64)
    for e, (i, j) in enumerate(edges):
        B[i, e] = 1.0; B[j, e] = -1.0
    return B, B @ B.T, edges


def positive_mass_projection(x, mass):
    x = np.maximum(x, 0.0)
    s = x.sum(axis=1, keepdims=True)
    out = x / np.maximum(s, 1e-12) * mass
    dead = (s[:, 0] <= 1e-12)
    out[dead] = mass / x.shape[1]
    return out


def density_step(rho, B, L, edges, dt, D, rng):
    # Finite-volume edge density and divergence-form stochastic flux.
    edge_rho = np.stack([(rho[:, i] + rho[:, j]) / 2.0 for i, j in edges], axis=1)
    xi = rng.standard_normal((rho.shape[0], len(edges)))
    stochastic_flux = np.sqrt(np.maximum(edge_rho, 0.0)) * xi
    nxt = rho - dt * D * (rho @ L.T) + np.sqrt(2.0 * D * dt) * (stochastic_flux @ B.T)
    return positive_mass_projection(nxt, rho.shape[1])


def math_check():
    B, L, edges = grid_graph(8, 8)
    eig = np.linalg.eigvalsh(L)
    rng = np.random.default_rng(SEED)
    rho = np.ones((4096, 64))
    dt = 0.20 / eig.max()  # D dt lambda_max = 0.2, safely below 1
    sums=[]; mins=[]; snapshots=[]
    for t in range(80):
        rho = density_step(rho, B, L, edges, dt, 1.0, rng)
        sums.append(np.max(np.abs(rho.sum(1)-64.0)))
        mins.append(rho.min())
        if t >= 40: snapshots.append(rho.copy())
    x = np.concatenate(snapshots, axis=0) - 1.0
    # Spatial covariance by Manhattan distance, compared to iid Gaussian with same variance.
    coords=[(i//8,i%8) for i in range(64)]
    bins={d:[] for d in range(1,9)}
    for i,(ri,ci) in enumerate(coords):
        for j,(rj,cj) in enumerate(coords):
            d=abs(ri-rj)+abs(ci-cj)
            if 1 <= d <= 8: bins[d].append((i,j))
    cov={d: float(np.mean([np.mean(x[:,i]*x[:,j]) for i,j in pairs])) for d,pairs in bins.items()}
    var=float(np.mean(x*x))
    iid={d: (0.0 if d>0 else var) for d in bins}
    return {'lambda_max':float(eig.max()), 'dt':float(dt),
            'max_mass_error':float(max(sums)), 'minimum_density':float(min(mins)),
            'variance':var, 'covariance_by_manhattan_distance':cov,
            'iid_control_covariance_by_distance':iid,
            'mass_and_positivity_pass': bool(max(sums)<1e-10 and min(mins)>=-1e-12)}


class TokenNet(nn.Module):
    def __init__(self, mode, n=64, dim=16, alpha=0.30, seed=SEED):
        super().__init__(); self.mode=mode; self.n=n; self.dim=dim; self.alpha=alpha
        self.embed=nn.Linear(1, dim); self.fc=nn.Sequential(nn.LayerNorm(n*dim), nn.Linear(n*dim,64), nn.GELU(), nn.Linear(64,10))
        B,L,edges=grid_graph(8,8); self.B=torch.tensor(B,dtype=torch.float32); self.L=torch.tensor(L,dtype=torch.float32); self.edges=edges
        self.rng=np.random.default_rng(seed+17)
    def forward(self, x):
        z=self.embed(x.reshape(x.shape[0], self.n, 1))
        if self.training and self.mode != 'none':
            if self.mode == 'iid':
                noise=torch.randn_like(z)
                z=z + self.alpha*noise
            else:
                b=x.shape[0]; rho=np.ones((b,self.n), dtype=np.float64)
                # Autonomous density is deliberately detached from model autograd.
                rho=density_step(rho, self.B.numpy(), self.L.numpy(), self.edges, 0.20/float(torch.linalg.eigvalsh(self.L).max()), 1.0, self.rng)
                amplitude=torch.tensor((rho-1.0)/np.sqrt(0.65),dtype=z.dtype,device=z.device).unsqueeze(-1)
                z=z + self.alpha*amplitude*torch.randn_like(z)
        return self.fc(z.flatten(1))


def train_eval(mode, Xtr, ytr, Xte, yte, device):
    torch.manual_seed(SEED); model=TokenNet(mode).to(device); opt=torch.optim.AdamW(model.parameters(),lr=2e-3,weight_decay=1e-4)
    bs=128; losses=[]; t0=time.time()
    for epoch in range(15):
        model.train(); perm=torch.randperm(len(Xtr),device=device); total=0.
        for st in range(0,len(Xtr),bs):
            ix=perm[st:st+bs]; out=model(Xtr[ix]); loss=nn.functional.cross_entropy(out,ytr[ix]); opt.zero_grad(); loss.backward(); opt.step(); total += loss.item()*len(ix)
        losses.append(total/len(Xtr))
    model.eval()
    with torch.no_grad(): pred=model(Xte).argmax(1); acc=float((pred==yte).float().mean())
    return {'test_accuracy':acc,'final_train_loss':losses[-1],'seconds':time.time()-t0,'loss_curve':losses}


def main():
    check=math_check()
    d=load_digits(); X=d.images.astype('float32')/16.0; y=d.target.astype('int64')
    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    try:
        Xtr=torch.tensor(Xtr).to(device); Xte=torch.tensor(Xte).to(device); ytr=torch.tensor(ytr).to(device); yte=torch.tensor(yte).to(device)
        results={m:train_eval(m,Xtr,ytr,Xte,yte,device) for m in ('none','iid','conserved')}
    except Exception as e:
        device=torch.device('cpu'); Xtr=torch.tensor(Xtr).cpu(); Xte=torch.tensor(Xte).cpu(); ytr=torch.tensor(ytr).cpu(); yte=torch.tensor(yte).cpu()
        results={m:train_eval(m,Xtr,ytr,Xte,yte,device) for m in ('none','iid','conserved')}; results['device_fallback']=str(e)
    out={'seed':SEED,'device':str(device),'math_check':check,'training':results}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()