import json, math, random
from pathlib import Path
import numpy as np

SEED=1150
np.random.seed(SEED); random.seed(SEED)


def sbm(n1=60,n2=60,p_in=.22,p_out=.025):
    n=n1+n2
    y=np.r_[np.zeros(n1,dtype=int),np.ones(n2,dtype=int)]
    R=np.random.rand(n,n)
    same=y[:,None]==y[None,:]
    A=((R < np.where(same,p_in,p_out))).astype(float)
    A=np.triu(A,1); A=A+A.T; np.fill_diagonal(A,0)
    # ensure connected enough; this seed graph is connected
    return A,y

def matrices(A,v):
    d=A.sum(1); L=np.diag(d)-A
    inv=np.diag(1/np.sqrt(v)); S=inv@L@inv
    B=np.diag(1/v)@L
    return d,L,S,B

def ipr(u): return float(np.sum(u**4)/np.sum(u**2)**2)

def math_checks(A):
    d,L,S,B=matrices(A,np.ones(len(A)))
    lam,U=np.linalg.eigh(S); lmax=lam[-1]
    # Prediction 1: Euler spectral boundary eta_c=2/lambda_max.
    pred_eta=2/lmax
    etas=pred_eta*np.linspace(.80,1.20,161)
    rho=np.array([max(abs(1-e*lam)) for e in etas])
    
    # first grid point whose spectral radius exceeds one, ignoring eta=0
    crossing_idx=np.where(rho>1+1e-10)[0]
    cross=etas[crossing_idx[0]] if len(crossing_idx) else float('nan')
    # Prediction 2: exact modal decay/amplification after m steps.
    eta=.65*pred_eta; m=17; j=len(lam)-1
    x=U[:,j]; measured=np.linalg.norm(np.linalg.matrix_power(np.eye(len(A))-eta*S,m)@x)
    predicted=abs(1-eta*lam[j])**m
    # Prediction 3: volume heterogeneity/localization. use v=k^alpha, and independent lognormal.
    rows=[]
    for alpha in [-1.0,-.5,0,.5,1.0,1.5,2.0]:
        v=np.maximum(d,1e-3)**alpha
        _,_,Ss,_=matrices(A,v); ee,uu=np.linalg.eigh(Ss)
        r=d/v
        rows.append({'kind':'degree_power','alpha':alpha,'std_log_r':float(np.std(np.log(r))), 'ipr_top':ipr(uu[:,-1])})
    # Replicate independent lognormal volumes: the prediction is a positive
    # association between ratio heterogeneity and extremal-mode IPR.
    for s in [0,.25,.5,.75,1.0,1.25]:
        vals=[]
        for rep in range(8):
            rng=np.random.RandomState(SEED+1000+rep)
            v=np.exp(s*rng.randn(len(A))); _,_,Ss,_=matrices(A,v); ee,uu=np.linalg.eigh(Ss)
            r=d/v
            vals.append((np.std(np.log(r)),ipr(uu[:,-1])))
        rows.append({'kind':'lognormal_replicated','s':s,
                     'std_log_r_mean':float(np.mean([z[0] for z in vals])),
                     'ipr_top_mean':float(np.mean([z[1] for z in vals])),
                     'ipr_top_std':float(np.std([z[1] for z in vals]))})
    rep_rows=[z for z in rows if z['kind']=='lognormal_replicated']
    localization_corr=float(np.corrcoef([z['std_log_r_mean'] for z in rep_rows], [z['ipr_top_mean'] for z in rep_rows])[0,1])
    # Similarity/eigenvalue sanity
    simerr=np.max(np.abs(np.sort(np.real(np.linalg.eigvals(B)))-lam))
    return {'lambda_max':float(lmax),'predicted_eta_c':pred_eta,'observed_eta_crossing':float(cross),
            'boundary_relative_error':float(abs(cross-pred_eta)/pred_eta),
            'decay_eta':eta,'decay_steps':m,'decay_predicted':float(predicted),'decay_measured':float(measured),
            'similarity_eigenvalue_max_error':float(simerr),'localization':rows,'localization_correlation':localization_corr}

def train_experiment(A,y,device='cpu'):
    try:
        import torch
        torch.manual_seed(SEED)
        if device=='cuda': torch.cuda.manual_seed_all(SEED)
        dev=torch.device(device)
        n=len(y); X=torch.eye(n,device=dev); Y=torch.tensor(y,device=dev)
        At=torch.tensor(A,dtype=torch.float32,device=dev)
        d=At.sum(1); L=torch.diag(d)-At
        # standard normalized GCN propagation, with self loops
        An=At+torch.eye(n,device=dev); dn=An.sum(1); Pn=An/dn[:,None]
        # mass diffusion: stable step and self-residual, same 2-layer width
        v=torch.tensor(np.maximum(d.cpu().numpy(),1e-3)**1.0,dtype=torch.float32,device=dev)
        B=L/v[:,None]
        lam=torch.linalg.eigvalsh(torch.diag(1/torch.sqrt(v))@L@torch.diag(1/torch.sqrt(v)))[-1]
        eta=.8*2/lam
        Pm=torch.eye(n,device=dev)-eta*B
        def run(kind):
            torch.manual_seed(SEED+3)
            W1=torch.nn.Parameter(torch.randn(n,16,device=dev)*.08); W2=torch.nn.Parameter(torch.randn(16,2,device=dev)*.08)
            opt=torch.optim.Adam([W1,W2],lr=.03,weight_decay=1e-3)
            idx=torch.randperm(n,device=dev); tr=idx[:80]; te=idx[80:]
            vals=[]
            for step in range(250):
                opt.zero_grad()
                P=Pn if kind=='gcn' else Pm
                h=torch.relu(P@X@W1); z=P@h@W2
                loss=torch.nn.functional.cross_entropy(z[tr],Y[tr]); loss.backward(); opt.step()
                if step in (49,99,249): vals.append(float(loss.detach().cpu()))
            with torch.no_grad():
                h=torch.relu(P@X@W1); z=P@h@W2; acc=float((z[te].argmax(1)==Y[te]).float().mean().cpu())
            return {'loss_steps_50_100_250':vals,'test_accuracy':acc}
        return {'device':str(dev),'gcn':run('gcn'),'mass_diffusion':run('mass')}
    except Exception as e:
        return {'error':repr(e)}

def main():
    A,y=sbm(); checks=math_checks(A)
    try:
        import torch
        device='cuda' if torch.cuda.is_available() else 'cpu'
    except Exception: device='cpu'
    exp=train_experiment(A,y,device)
    out={'seed':SEED,'n':len(y),'math':checks,'experiment':exp}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()