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

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


def disagreement(x):
    z = x - x.mean(axis=0, keepdims=True)
    return float(np.mean(z*z))


def graph_laplacian(n, p, seed):
    r = np.random.default_rng(seed)
    A = (r.random((n,n)) < p).astype(float)
    A = np.triu(A, 1); A = A + A.T
    # Ensure connectedness by adding a ring.
    for i in range(n):
        A[i, (i+1)%n] = A[(i+1)%n, i] = 1.0
    return np.diag(A.sum(1)) - A


def toy_sweeps():
    # Prediction 1: exact complete-consensus contraction D'/D=(1-alpha)^2.
    x = rng.normal(size=(30, 4))
    alphas = np.linspace(0.0, 1.8, 19)
    complete = []
    for a in alphas:
        y = x + a*(x.mean(0, keepdims=True)-x)
        complete.append(disagreement(y)/disagreement(x))
    pred1 = (1-alphas)**2
    err1 = float(np.max(np.abs(np.asarray(complete)-pred1)))

    # Prediction 2: Laplacian diffusion is stable iff alpha*lambda_max<2.
    L = graph_laplacian(30, .16, SEED)
    ev = np.linalg.eigvalsh(L)
    lmax, l2 = float(ev[-1]), float(ev[1])
    # Use the highest-frequency eigenmode; arbitrary x can hide the operator boundary.
    vmax = np.linalg.eigh(L)[1][:,-1:]
    agrid = np.linspace(0, 2.4/lmax, 97)
    lap_ratios = []
    for a in agrid:
        y = vmax - a*L@vmax
        lap_ratios.append(disagreement(y)/disagreement(vmax))
    observed_boundary = float(agrid[np.where(np.asarray(lap_ratios) > 1)[0][0]]) if np.any(np.asarray(lap_ratios)>1) else float('nan')
    predicted_boundary = 2/lmax
    # Prediction 3: along the slow Fiedler mode, one-step ratio is (1-alpha*l2)^2.
    fiedler = np.linalg.eigh(L)[1][:,1:2]
    f_ratios=[]; f_pred=[]
    for a in alphas:
        q=fiedler + 0.0
        f_ratios.append(disagreement(q-a*L@q)/disagreement(q))
        f_pred.append((1-a*l2)**2)
    err3=float(np.max(np.abs(np.asarray(f_ratios)-np.asarray(f_pred))))

    return {
        'complete_consensus': {'alpha': alphas.tolist(), 'observed_ratio': complete,
                               'predicted_ratio_(1-alpha)^2': pred1.tolist(), 'max_abs_error': err1},
        'laplacian_diffusion': {'lambda2': l2, 'lambda_max': lmax,
                                'predicted_boundary_2/lambda_max': predicted_boundary,
                                'observed_first_grid_ratio_gt_1': observed_boundary,
                                'alpha_grid': agrid.tolist(), 'ratios': lap_ratios},
        'fiedler_mode': {'predicted_ratio_(1-alpha*lambda2)^2': f_pred,
                         'observed_ratio': f_ratios, 'max_abs_error': err3}
    }


def normalized_adj(A):
    d=A.sum(1)+1e-6
    return A/np.sqrt(d[:,None]*d[None,:])


def learned_outage_test():
    # Tiny scalar node task: target is local feature plus graph-wide mean.
    # Train on random edge masks; evaluate on a held-out, much sparser mask.
    import torch
    torch.manual_seed(SEED); np.random.seed(SEED)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        n,d=24,5
        base=graph_laplacian(n,.22,SEED+4); baseA=np.diag(np.diag(base))-base
        X=torch.tensor(np.random.randn(n,d),dtype=torch.float32,device=device)
        target=(X[:,0:1] + .7*X.mean(0,keepdim=True).repeat(n,1)[:,0:1]).detach()
        class Net(torch.nn.Module):
            def __init__(self, consensus):
                super().__init__(); self.consensus=consensus
                self.w=torch.nn.Linear(d,12); self.out=torch.nn.Linear(12,1)
                self.msg=torch.nn.Linear(12,12,bias=False)
            def forward(self,A):
                H=torch.relu(self.w(X)); S=torch.tensor(normalized_adj(A),dtype=torch.float32,device=device)
                for _ in range(4):
                    m=S@H
                    H=H + torch.sigmoid(m.mean(1,keepdim=True))*self.msg(m)
                    if self.consensus: H=H + .10*(H.mean(0,keepdim=True)-H)
                    H=torch.relu(H)
                return self.out(H),H
        models=[Net(False),Net(True)]
        results={}
        for model in models:
            model.to(device); opt=torch.optim.Adam(model.parameters(),lr=.015)
            for step in range(500):
                mask=(np.random.random((n,n))<.75).astype(float); mask=np.triu(mask,1); mask=mask+mask.T
                mask=np.maximum(mask,baseA)
                pred,_=model(mask); loss=((pred-target)**2).mean()
                opt.zero_grad(); loss.backward(); opt.step()
            vals=[]; dis=[]
            for frac in [.75,.55,.35,.20]:
                errs=[]; ds=[]
                for k in range(30):
                    mask=(np.random.random((n,n))<frac).astype(float); mask=np.triu(mask,1); mask=mask+mask.T
                    mask=np.maximum(mask, np.eye(n)*0)
                    with torch.no_grad(): pred,H=model(mask)
                    errs.append(float(((pred-target)**2).mean().cpu())); ds.append(disagreement(H.cpu().numpy()))
                vals.append(float(np.mean(errs))); dis.append(float(np.mean(ds)))
            results['consensus' if model.consensus else 'baseline']={'mse_by_edge_keep':vals,'representation_D':dis}
        results['device']=device
        return results
    except Exception as e:
        return {'error':repr(e),'device':'cpu_fallback'}


def main():
    out={'seed':SEED,'toy':toy_sweeps(),'learned_outage':learned_outage_test()}
    Path('results.json').write_text(json.dumps(out,indent=2))
    t=out['toy']; print(json.dumps({'complete_max_error':t['complete_consensus']['max_abs_error'],
      'lap_pred_boundary':t['laplacian_diffusion']['predicted_boundary_2/lambda_max'],
      'lap_observed_boundary':t['laplacian_diffusion']['observed_first_grid_ratio_gt_1'],
      'fiedler_max_error':t['fiedler_mode']['max_abs_error'],
      'learned':out['learned_outage']},indent=2))
if __name__=='__main__': main()