import json, math, os
import numpy as np
from scipy.sparse import coo_matrix, csr_matrix
from scipy.sparse.linalg import spsolve

SEED = 1387
rng = np.random.default_rng(SEED)


def grid_graph(dim, R):
    """Finite integer lattice ball (L_inf ball), with Dirichlet boundary outside radius R."""
    pts = [tuple(x) for x in np.ndindex(*([2*R+1]*dim))]
    pts = [tuple(a-R for a in p) for p in pts]
    idx = {p:i for i,p in enumerate(pts)}
    edges=[]
    for p in pts:
        for d in range(dim):
            q=list(p); q[d]+=1; q=tuple(q)
            if q in idx: edges.append((idx[p],idx[q],1.0))
    return pts, idx, edges


def capacity_grid(dim, R, n=3, distance_weight=False):
    pts, idx, edges = grid_graph(dim,R)
    origin=idx[tuple([0]*dim)]
    # Dirichlet problem: f(origin)=1, all nodes on outer L_inf shell are 0.
    fixed={origin:1.0}
    for p,i in idx.items():
        if max(abs(x) for x in p)==R: fixed[i]=0.0
    free=[i for i in range(len(pts)) if i not in fixed]
    pos={i:j for j,i in enumerate(free)}
    rows=[]; cols=[]; vals=[]; b=np.zeros(len(free))
    adj=[[] for _ in pts]
    for i,j,_ in edges:
        # lattice edge geometric distance is one, hence w=d^(n-2)=1
        w=1.0
        adj[i].append((j,w)); adj[j].append((i,w))
    for i in free:
        r=pos[i]; deg=sum(w for _,w in adj[i]); rows.append(r); cols.append(r); vals.append(deg)
        for j,w in adj[i]:
            if j in fixed: b[r]+=w*fixed[j]
            else: rows.append(r); cols.append(pos[j]); vals.append(-w)
    f=np.zeros(len(pts)); f[list(fixed)]=list(fixed.values())
    if free: f[free]=spsolve(coo_matrix((vals,(rows,cols)),shape=(len(free),len(free))).tocsr(),b)
    # half ordered energy = sum undirected edge w*(du)^2
    energy=0.0
    for i,j,w in edges: energy += w*(f[i]-f[j])**2
    return float(energy)


def path_transition(N, self_loop=0.0):
    W=np.zeros((N,N))
    for i in range(N-1): W[i,i+1]=W[i+1,i]=1.0
    if self_loop: W += self_loop*np.eye(N)
    return W/W.sum(axis=1,keepdims=True)


def influence_scaling():
    # Large path, start away from boundaries; fit the known 1D local CLT prediction ||p_t||_2 ~ t^-1/4.
    N=4001; P=path_transition(N, self_loop=1.0); x=np.zeros(N); x[N//2]=1
    times=np.array([4,8,16,32,64,128,256,512,1024,2048])
    vals=[]; t=0
    for target in times:
        for _ in range(target-t): x=P@x
        t=target; vals.append(np.linalg.norm(x))
    slope=np.polyfit(np.log(times),np.log(vals),1)[0]
    return {'times':times.tolist(),'l2':np.round(vals,8).tolist(),'fitted_exponent':float(slope),'predicted_exponent':-0.25}


def diffusion_checks():
    # Distance-weighted P on a tiny irregular geometric graph. Verify stochasticity and stationary conservation.
    coords=np.array([[0.,0.],[1.,0.],[0.,2.],[2.,1.],[3.,0.]])
    E=[(0,1),(0,2),(1,3),(2,3),(1,4),(3,4)]
    n=4; W=np.zeros((len(coords),len(coords)))
    for i,j in E:
        d=np.linalg.norm(coords[i]-coords[j]); w=d**(n-2); W[i,j]=W[j,i]=w
    P=W/W.sum(1,keepdims=True)
    x=rng.normal(size=(len(coords),3)); y=P@x
    return {'max_row_sum_error':float(np.max(abs(P.sum(1)-1))), 'max_nonnegative_error':float(max(0,-P.min())), 'feature_mass_before':float(x.sum()), 'feature_mass_after_degree_weighted':float((W.sum(1)[:,None]*y).sum()), 'feature_mass_expected':float((W.sum(1)[:,None]*x).sum())}


def classification_demo():
    # Same graph/parameters: noisy coordinate labels, compare unweighted GCN and d^(n-2) diffusion.
    rs=np.random.default_rng(SEED); N=180; xy=rs.uniform(-1,1,(N,2)); labels=(xy[:,0]+0.25*xy[:,1]>0).astype(np.float32)
    # kNN graph, n=4 makes longer edges receive larger weights.
    D=((xy[:,None,:]-xy[None,:,:])**2).sum(2); W0=np.zeros((N,N))
    k=8
    for i in range(N):
        for j in np.argsort(D[i])[1:k+1]: W0[i,j]=W0[j,i]=1
    Wd=W0*np.maximum(D,1e-12) # n-2=2, d^2 = D
    def norm(W): return W/W.sum(1,keepdims=True)
    Ps=[norm(W0),norm(Wd)]
    X=np.stack([xy[:,0]+rs.normal(0,.7,N),xy[:,1]+rs.normal(0,.7,N)],1)
    out=[]
    for name,P in zip(['GCN_unweighted','capacity_diffusion_weighted'],Ps):
        h=X.copy(); acc=[]; var=[]
        # fixed linear readout, labels are the task; report best depth, not train a model.
        for depth in range(1,21):
            h=.5*h+.5*(P@h)
            score=h[:,0]; pred=(score>0).astype(np.float32)
            acc.append(float((pred==labels).mean())); var.append(float(np.mean(np.sum((h-h.mean(0))**2,1))))
        out.append({'model':name,'best_accuracy':max(acc),'best_depth':int(np.argmax(acc)+1),'accuracy_depth20':acc[-1],'variance_depth1':var[0],'variance_depth20':var[-1]})
    return out


def main():
    caps={}
    for dim in [1,2,3]:
        Rs=([4,8,16,32,64,128] if dim==1 else ([4,8,16,24,32,48,64] if dim==2 else [2,3,4,5,6,8,10]))
        vals=[capacity_grid(dim,R) for R in Rs]
        caps[str(dim)]={'R':Rs,'capacity':vals}
    # Predictions: 1D Cap~R^-1; 2D Cap~1/log R. Fit slopes on sufficiently large radii.
    r1=np.array(caps['1']['R'],float); c1=np.array(caps['1']['capacity']); slope1=np.polyfit(np.log(r1),np.log(c1),1)[0]
    r2=np.array(caps['2']['R'],float); c2=np.array(caps['2']['capacity']); invlog=np.polyfit(1/np.log(r2),c2,1)
    # In 3D, compare late/early capacity; predicted nonzero limiting capacity.
    c3=np.array(caps['3']['capacity'])
    result={'seed':SEED,'prediction_checks':{
      'path_capacity_power_law':{'predicted_exponent':-1.0,'observed_exponent':float(slope1),'relative_error':float(abs(slope1+1))},
      'grid2_capacity_log_law':{'prediction':'Cap approximately affine in 1/log(R)','R2':float(np.corrcoef(1/np.log(r2),c2)[0,1]**2),'capacity_ratio_R64_over_R4':float(c2[-1]/c2[0]),'log_ratio_prediction':float(np.log(r2[0])/np.log(r2[-1]))},
      'grid3_capacity_nonvanishing':{'early':float(c3[0]),'late':float(c3[-1]),'late_over_early':float(c3[-1]/c3[0])},
      'influence_local_clt':influence_scaling(),
      'row_stochasticity':diffusion_checks()},'capacities':caps,'classification':classification_demo()}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()