import sys, json
import numpy as np
import torch
import torch.nn as nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
import bench
from bench import train_model, evaluate, sweep_baseline, make_report

TRACK = 'geometric_graph_diffusion'

def graph_operator(x):
    # x[...,2:4] are the observed node coordinates in the registered track.
    xy = x[..., 2:4]
    d2 = ((xy[:, :, None, :] - xy[:, None, :, :]) ** 2).sum(-1)
    b, n, _ = d2.shape
    adj = torch.zeros_like(d2)
    # Five nearest neighbors, symmetrized, matching the registered generator.
    ind = torch.argsort(d2, dim=-1)[:, :, 1:6]
    adj.scatter_(2, ind, 1.0)
    adj = torch.maximum(adj, adj.transpose(1, 2))
    w = adj * torch.clamp(d2, min=1e-5)
    p = w / torch.clamp(w.sum(-1, keepdim=True), min=1e-8)
    return 0.55 * torch.eye(n, device=x.device).unsqueeze(0) + 0.45 * p

class GraphDiffusionNet(nn.Module):
    # The two systems share every learnable parameter and differ only in alpha.
    def __init__(self, alpha):
        super().__init__()
        self.alpha = float(alpha)
        self.fc1 = nn.Linear(5, 32)
        self.fc2 = nn.Linear(32, 1)
    def forward(self, x):
        q = graph_operator(x)
        if self.alpha:
            n = x.shape[1]
            q = (1.0-self.alpha)*q + self.alpha*torch.ones_like(q)/n
        h = torch.relu(self.fc1(x))
        h = torch.bmm(q, h)
        z = self.fc2(h).squeeze(-1)
        # Root marker makes this a scalar graph readout, as required by target.
        root = x[..., 4]
        return (z * root).sum(1, keepdim=True)

def run(cfg, seed, return_model=False):
    torch.manual_seed(9000 + int(seed)); np.random.seed(9000 + int(seed))
    ds = bench.get_dataset(TRACK, int(seed), n_train=400, n_test=160)
    net = GraphDiffusionNet(cfg['alpha'])
    net, metric, history = train_model(net, ds, epochs=cfg['epochs'], lr=cfg['lr'],
                                       batch=128, weight_decay=cfg['wd'], log=lambda *_: None)
    return (float(metric), net, ds) if return_model else float(metric)

def factory(cfg):
    return lambda seed: run(cfg, seed)

def main():
    # Shared union of all learning-rate/regularization settings on both sides.
    common = [(lr, wd) for lr in (0.003, 0.01, 0.03) for wd in (0.0, 0.0001)]
    baseline_grid = [{'alpha': 0.0, 'lr': lr, 'wd': wd, 'epochs': 20} for lr,wd in common]
    baseline = sweep_baseline(factory, baseline_grid, seeds=(0,1,2,3))
    best = baseline['best_cfg']
    idea_grid = []
    for alpha in (0.03, 0.1, 0.3):
        c = dict(best); c['alpha'] = alpha; idea_grid.append(c)
    idea_candidates = [{'cfg': c, 'result': evaluate(factory(c))} for c in idea_grid]
    chosen = min(idea_candidates, key=lambda z: z['result']['mean'])
    idea = chosen['result']

    # Behavioural signature from trained systems, not an analytical toy graph.
    rows=[]
    for alpha in (0.0, 0.03, 0.1, 0.3):
        cfg=dict(best); cfg['alpha']=alpha
        _, net, ds = run(cfg, 0, return_model=True)
        net.eval(); x=ds['xte'].to(next(net.parameters()).device)
        with torch.no_grad():
            q=graph_operator(x)
            h=torch.relu(net.fc1(x)); hc=h-h.mean(1,keepdim=True)
            before=hc.norm().item()
            n=x.shape[1]; qt=(1-alpha)*q + alpha*torch.ones_like(q)/n
            after=torch.bmm(qt,hc).norm().item()
        rows.append({'alpha':alpha,'observed_mean_centered_norm_ratio':after/max(before,1e-12), 'predicted':1-alpha})
    err=max(abs(r['observed_mean_centered_norm_ratio']-r['predicted']) for r in rows)
    extra={'prediction':'uniform teleportation contracts mean-zero feature modes approximately by 1-alpha',
           'trained_model_rows':rows,'max_abs_error':float(err),'confirmed':bool(err < 0.10),
           'idea_sweep':idea_candidates,
           'track_match':'registered geometric_graph_diffusion: node-field regression on geometric kNN graphs'}
    report=make_report(TRACK,'graph_diffusion_net',baseline,idea,extra)
    report['idea_candidates']=idea_candidates
    with open('bench_report.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))

if __name__=='__main__': main()