import sys, json, math
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, evaluate, sweep_baseline, make_report

N = 8
EDGE = np.triu_indices(N, 1)
SEEDS = tuple(range(8))


def tree_loss(w):
    b, n, _ = w.shape
    deg = w.sum(-1)
    lap = torch.diag_embed(deg) - w
    shift = torch.ones((b, n, n), device=w.device, dtype=w.dtype) / n
    sign, ld = torch.linalg.slogdet(lap + shift)
    if torch.any(sign <= 0):
        raise RuntimeError('non-positive shifted Laplacian')
    return (-ld + math.log(n)).mean()


def math_check():
    rng = np.random.default_rng(91)
    a = rng.uniform(.1, 1., (N, N)); w = (a + a.T) / 2
    np.fill_diagonal(w, 0)
    lap = np.diag(w.sum(1)) - w
    ev = np.linalg.eigvalsh(lap)
    pdet = np.prod(ev[1:])
    rdet = np.linalg.det(lap + np.ones((N, N)) / N)
    cof = np.linalg.det(lap[:-1, :-1])
    return {'relative_rank_one_pdet_error': float(abs(pdet-rdet)/pdet),
            'relative_matrix_tree_cofactor_error': float(abs(pdet/N-cof)/cof)}


class GraphNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.body = nn.Sequential(nn.Flatten(), nn.Linear(16, 64), nn.ReLU(),
                                  nn.Linear(64, 64), nn.ReLU())
        self.edge = nn.Linear(64, 28)

    def forward(self, x):
        z = self.body(x)
        e = F.softplus(self.edge(z))
        w = torch.zeros((x.shape[0], N, N), device=x.device, dtype=x.dtype)
        w[:, EDGE[0], EDGE[1]] = e
        w[:, EDGE[1], EDGE[0]] = e
        return w, e.mean(-1, keepdim=True)


def spectral_metrics(w):
    fiedler, comps, pdet = [], [], []
    for a in w:
        lap = np.diag(a.sum(1)) - a
        ev = np.linalg.eigvalsh(lap)
        fiedler.append(float(ev[1]))
        comps.append(int(np.sum(ev < 1e-7)))
        pdet.append(float(np.prod(np.maximum(ev[1:], 1e-30))))
    return {'fiedler_mean': float(np.mean(fiedler)),
            'components_mean': float(np.mean(comps)),
            'pseudo_det_mean': float(np.mean(pdet))}


def train_one(seed, alpha, lr, epochs=35, return_model=False):
    torch.manual_seed(seed); np.random.seed(seed)
    d = get_dataset('soft_graph_connectivity', seed, 400, 120)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        if device == 'cuda': torch.zeros(1, device=device).sum().item()
    except Exception:
        device = 'cpu'
    model = GraphNet().to(device)
    xtr, ytr = torch.as_tensor(d['xtr'], dtype=torch.float32, device=device), torch.as_tensor(d['ytr'], dtype=torch.float32, device=device)
    xte, yte = torch.as_tensor(d['xte'], dtype=torch.float32, device=device), torch.as_tensor(d['yte'], dtype=torch.float32, device=device)
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    for _ in range(epochs):
        model.train()
        for start in range(0, len(xtr), 128):
            w, pred = model(xtr[start:start+128])
            loss = F.mse_loss(pred, ytr[start:start+128])
            if alpha:
                loss = loss + alpha * tree_loss(w)
            opt.zero_grad(); loss.backward(); opt.step()
    model.eval()
    with torch.no_grad():
        w, pred = model(xte)
        metric = float(F.mse_loss(pred, yte).cpu())
    if return_model:
        return metric, spectral_metrics(w.detach().cpu().numpy()), model
    return metric


def make_fn(cfg):
    return lambda seed: train_one(seed, cfg['alpha'], cfg['lr'])


def main():
    # Union parity: every idea lr is also present in baseline sweep.
    grid = [{'alpha': 0.0, 'lr': lr} for lr in (1e-3, 3e-3, 1e-2, 2e-3, 6e-3)]
    base = sweep_baseline(make_fn, grid)
    # The baseline sweep itself selects alpha=0; idea is evaluated at same lrs.
    idea_lrs = sorted(set([base['best_cfg']['lr'], 2e-3, 6e-3]))
    idea_cfgs = [{'alpha': a, 'lr': lr} for lr in idea_lrs for a in (1e-4, 1e-3, 1e-2)]
    # Run all three on the full paired seeds, then retain best by the same 4-seed selection.
    idea_sweep = []
    for cfg in idea_cfgs:
        r4 = evaluate(make_fn(cfg), (0,1,2,3))
        idea_sweep.append({'cfg': cfg, 'mean': r4['mean']})
    best_idea_cfg = min(idea_sweep, key=lambda z: z['mean'])['cfg']
    idea = evaluate(make_fn(best_idea_cfg), SEEDS)
    base['idea_side_sweep'] = idea_sweep
    sig_base=[]; sig_idea=[]
    for s in SEEDS:
        _, mb, _ = train_one(s, 0.0, base['best_cfg']['lr'], return_model=True)
        _, mi, _ = train_one(s, best_idea_cfg['alpha'], best_idea_cfg['lr'], return_model=True)
        sig_base.append(mb); sig_idea.append(mi)
    signature = {
        'quantity': 'Fiedler eigenvalue of trained predicted adjacency',
        'baseline_predicted_fiedler_mean': float(np.mean([z['fiedler_mean'] for z in sig_base])),
        'idea_predicted_fiedler_mean': float(np.mean([z['fiedler_mean'] for z in sig_idea])),
        'baseline_components_mean': float(np.mean([z['components_mean'] for z in sig_base])),
        'idea_components_mean': float(np.mean([z['components_mean'] for z in sig_idea])),
        'prediction': 'tree regularization should increase Fiedler connectivity',
        'confirmed': bool(np.mean([z['fiedler_mean'] for z in sig_idea]) > np.mean([z['fiedler_mean'] for z in sig_base]))
    }
    report = make_report('soft_graph_connectivity', 'graphnet_custom', base, idea,
                         {'mechanism_signature': signature,
                          'custom_track': {'name': 'soft_graph_connectivity', 'file': 'graph_track.py', 'domain': 'graph-nn'},
                          'math_check': math_check(), 'idea_cfg': best_idea_cfg})
    Path('report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()