import json, math, time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

SEEDS = [0, 1, 2]
N = 12

def laplacian(w):
    return torch.diag(w.sum(-1)) - w

def tree_loss(w):
    n = w.shape[-1]
    L = laplacian(w)
    sign, ld = torch.linalg.slogdet(L + torch.ones_like(L) / n)
    if sign.item() <= 0:
        raise RuntimeError('non-positive shifted Laplacian')
    return -ld + math.log(n)

def spectral_metrics(w):
    a = np.asarray(w)
    L = np.diag(a.sum(1)) - a
    ev = np.linalg.eigvalsh(L)
    return {
        'pseudo_det': float(np.prod(ev[1:])),
        'fiedler': float(ev[1]),
        'components': int(np.sum(ev < 1e-7)),
        'eigenvalues': ev.tolist(),
    }

def math_check():
    rng = np.random.default_rng(4)
    x = rng.uniform(.1, 1.2, (N, N))
    w = (x + x.T) / 2
    np.fill_diagonal(w, 0)
    L = np.diag(w.sum(1)) - w
    ev = np.linalg.eigvalsh(L)
    pdet = float(np.prod(ev[1:]))
    rankone_det = float(np.linalg.det(L + np.ones((N, N)) / N))
    cofactor = float(np.linalg.det(L[:-1, :-1]))
    eps = 1e-4
    derivative = (np.linalg.det(L + eps*np.eye(N)) - np.linalg.det(L - eps*np.eye(N))) / (2*eps)
    wt = torch.tensor(w, dtype=torch.double, requires_grad=True)
    tl = tree_loss(wt)
    tl.backward()
    return {
        'pseudo_det_vs_rank_one_det_relerr': abs(pdet-rankone_det)/pdet,
        'matrix_tree_pdet_over_N_vs_cofactor_relerr': abs(pdet/N-cofactor)/cofactor,
        'derivative_formula_relerr': abs(derivative-pdet)/pdet,
        'tree_loss_vs_stated_formula_abs': abs(float(tl)-(-math.log(pdet)+math.log(N))),
        'autograd_gradient_norm': float(wt.grad.norm()),
    }

class EdgeLearner(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.logits = nn.Parameter(torch.randn(n, n) * .15)
    def weights(self):
        s = (self.logits + self.logits.T) / 2
        w = F.softplus(s)
        return w * (1 - torch.eye(w.shape[0], device=w.device))

def make_target(seed):
    rng = np.random.default_rng(seed)
    # Two-community target: the learning objective must preserve task-relevant structure.
    y = np.zeros((N, N), dtype=np.float32)
    groups = np.arange(N) < N//2
    for i in range(N):
        for j in range(i):
            base = 1.0 if groups[i] == groups[j] else .08
            y[i,j] = y[j,i] = base + rng.normal(0, .025)
    np.fill_diagonal(y, 0)
    return torch.tensor(y)

def run_one(seed, alpha, device):
    torch.manual_seed(seed); np.random.seed(seed)
    target = make_target(seed).to(device)
    model = EdgeLearner(N).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=.08)
    losses=[]; t0=time.perf_counter()
    for step in range(250):
        w = model.weights()
        fit = ((w-target)**2).mean()
        reg = tree_loss(w)
        loss = fit + alpha*reg
        opt.zero_grad(); loss.backward(); opt.step()
        losses.append(float(loss.detach().cpu()))
    elapsed=time.perf_counter()-t0
    w=model.weights().detach().cpu().numpy()
    met=spectral_metrics(w)
    met.update({'fit_mse': float(((torch.tensor(w)-target.cpu())**2).mean()),
                'final_loss': losses[-1], 'loss_std_last50': float(np.std(losses[-50:])),
                'seconds': elapsed})
    return met

def main():
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        # Probe CUDA and fall back on any runtime/device error.
        if device == 'cuda': torch.zeros(1, device='cuda').sum().item()
    except Exception:
        device='cpu'
    result={'device':device, 'math_check':math_check(), 'runs':{}}
    for alpha in [0.0, 1e-3, 1e-2]:
        key='baseline' if alpha==0 else 'tree_alpha_'+str(alpha)
        vals=[]
        for seed in SEEDS:
            try: vals.append(run_one(seed, alpha, device))
            except Exception:
                if device=='cuda':
                    device='cpu'; result['device']='cpu'; vals.append(run_one(seed, alpha, device))
                else: raise
        result['runs'][key]=vals
    for key, vals in result['runs'].items():
        result['summary_'+key]={k:float(np.mean([v[k] for v in vals])) for k in ['fit_mse','fiedler','components','pseudo_det','loss_std_last50','seconds']}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__=='__main__': main()