import numpy as np
import torch
from scipy.linalg import solve
from scipy.sparse.csgraph import shortest_path
import time, json

SEED=417

def directed_dist(A):
    D=shortest_path(A, directed=True, unweighted=True)
    if not np.isfinite(D).all(): raise ValueError('graph is not strongly connected')
    return D

def curvature(D, eps=1e-6, clip=5.):
    n=len(D)
    # np.linalg.pinv implements the specified Moore-Penrose SVD cutoff convention
    mo=np.linalg.pinv(D, rcond=eps) @ (n*np.ones(n))
    mi=np.linalg.pinv(D.T, rcond=eps) @ (n*np.ones(n))
    scale=(np.abs(mo).sum()+np.abs(mi).sum())/(2*n)+1e-6
    z=np.stack([np.clip(mo/scale,-clip,clip), np.clip(mi/scale,-clip,clip)],1)
    return mo,mi,z,scale

def graph(n=96, p=0.10, seed=0):
    rng=np.random.default_rng(seed)
    y=np.repeat(np.arange(2), n//2)
    rng.shuffle(y)
    # Group 0 has relatively strong outgoing connections; group 1 receives them.
    P=np.full((2,2),p)
    P[0,1]=.22; P[1,0]=.035; P[0,0]=.10; P[1,1]=.10
    A=np.zeros((n,n),dtype=np.float32)
    for i in range(n):
        for j in range(n):
            if i!=j: A[i,j]=rng.random()<P[y[i],y[j]]
    # Guaranteed strongly-connected directed backbone, negligible as a local feature.
    for i in range(n): A[i,(i+1)%n]=1
    D=directed_dist(A)
    return A,y,D

def math_check():
    signed_sum_errors=[]; l1_errors=[]; residuals=[]; negative_cases=0
    for seed in range(20):
        A,_,D=graph(24,seed=100+seed)
        mo,mi,_,_=curvature(D)
        residuals.append(max(np.max(np.abs(D@mo-len(D))),
                            np.max(np.abs(D.T@mi-len(D)))))
        signed_sum_errors.append(abs(mo.sum()-mi.sum()))
        l1_errors.append(abs(np.abs(mo).sum()-np.abs(mi).sum()))
        if (mo < 0).any() or (mi < 0).any(): negative_cases += 1
    return {
        'instances': 20,
        'max_out_in_system_residual': float(max(residuals)),
        'max_signed_sum_difference': float(max(signed_sum_errors)),
        'mean_absolute_l1_difference': float(np.mean(l1_errors)),
        'max_absolute_l1_difference': float(max(l1_errors)),
        'negative_solution_cases': negative_cases
    }

class DirectedNet(torch.nn.Module):
    def __init__(self, in_dim):
        super().__init__()
        self.lin=torch.nn.Linear(in_dim*3,24)
        self.out=torch.nn.Linear(24,2)
    def forward(self,x,A):
        # A[i,j]=edge i->j; incoming and outgoing directed aggregations.
        out=A@x/(A.sum(1,keepdims=True)+1e-6)
        inc=A.T@x/(A.T.sum(1,keepdims=True)+1e-6)
        h=torch.relu(self.lin(torch.cat([x,out,inc],1)))
        return self.out(h)

def run_one(seed, use_curv, labeled_per_class=5):
    torch.manual_seed(seed); np.random.seed(seed)
    A,y,D=graph(seed=seed)
    n=len(y); rng=np.random.default_rng(seed+88)
    # Features deliberately carry no class information; degrees are the standard local control.
    X=rng.normal(0,1,(n,4)).astype(np.float32)
    if use_curv:
        _,_,z,_=curvature(D)
        X=np.concatenate([X,z.astype(np.float32)],1)
    At=torch.tensor(A); Xt=torch.tensor(X); yt=torch.tensor(y,dtype=torch.long)
    train=[]
    for c in [0,1]: train.extend(rng.choice(np.where(y==c)[0],labeled_per_class,replace=False))
    train=torch.tensor(train); test=torch.tensor([i for i in range(n) if i not in set(train.tolist())])
    model=DirectedNet(X.shape[1]); opt=torch.optim.Adam(model.parameters(),lr=.02,weight_decay=1e-4)
    for _ in range(250):
        opt.zero_grad(); logits=model(Xt,At); loss=torch.nn.functional.cross_entropy(logits[train],yt[train]); loss.backward(); opt.step()
    with torch.no_grad():
        logits=model(Xt,At); pred=logits.argmax(1); prob=logits.softmax(1)[:,1]
        acc=(pred[test]==yt[test]).float().mean().item()
        # Brier score is a simple calibration metric (lower is better).
        brier=((prob[test]-(yt[test]==1).float())**2).mean().item()
    return acc,brier

def main():
    t=time.perf_counter(); check=math_check(); prep=time.perf_counter()-t
    rows=[]
    for k in [5,10]:
        for mode in [False,True]:
            r=[run_one(100+s,mode,k) for s in range(5)]
            rows.append({'labels_per_class':k,'curvature':mode,'accuracy_mean':float(np.mean([x[0] for x in r])),'accuracy_std':float(np.std([x[0] for x in r])),'brier_mean':float(np.mean([x[1] for x in r]))})
    print(json.dumps({'math_check':check,'math_check_seconds':prep,'results':rows},indent=2))
if __name__=='__main__': main()