import json, random, sys
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import train_model, evaluate, sweep_baseline, make_report

META = {'name': 'influence_graph_regression', 'domain': 'graph-nn',
        'description': 'Fixed-topology synthetic graph regression where the target depends on one-hop nonlinear neighbor aggregation.'}


def get_dataset(seed, n_train=400, n_test=120):
    rng = np.random.RandomState(seed)
    n, f = 10, 3
    A = rng.uniform(0.05, 1.0, size=(n, n)).astype('float32')
    A[rng.rand(n, n) < 0.52] = 0.0
    np.fill_diagonal(A, 0.0)
    A += np.eye(n, dtype='float32') * 0.15
    A /= A.sum(1, keepdims=True)
    def make(m, rs):
        z = rs.normal(size=(m, n, f)).astype('float32')
        scalar = np.tanh(z[:, :, 0] * 0.9 + z[:, :, 1] * 0.35)
        neigh = np.einsum('ij,bjf->bif', A, z)
        y = (0.65 * neigh[:, :, 0] + 0.30 * np.tanh(neigh[:, :, 1]) + 0.20 * scalar).mean(1)
        y += 0.04 * rs.normal(size=m)
        topo = np.broadcast_to(A[None, :, :], (m, n, n)).copy()
        return np.concatenate([z, topo], axis=2).astype('float32'), y.astype('float32')[:, None]
    xtr, ytr = make(n_train, rng)
    xte, yte = make(n_test, np.random.RandomState(seed + 5000))
    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte, 'task': 'regression',
            'metric': 'mse', 'out_dim': 1, 'input_shape': (n, n + f)}


def torch_dataset(seed, n_train=400, n_test=120):
    d = get_dataset(seed, n_train, n_test)
    out = dict(d)
    for k in ('xtr', 'ytr', 'xte', 'yte'):
        out[k] = torch.as_tensor(d[k], dtype=torch.float32)
    return out


class GraphMessageNet(nn.Module):
    def __init__(self, idea=False, kappa=2.0, eps=0.05, xmin=1., xmax=8.):
        super().__init__()
        self.idea, self.kappa, self.eps = idea, kappa, eps
        self.xmin, self.xmax = xmin, xmax
        self.node = nn.Linear(3, 24)
        self.msg_proj = nn.Linear(1, 24, bias=False)
        self.out = nn.Sequential(nn.Linear(24, 24), nn.ReLU(), nn.Linear(24, 1))
        self.last = {}
    def forward(self, x):
        n = x.shape[1]
        feat, A = x[:, :, :3], x[:, :, 3:3+n]
        h = torch.relu(self.node(feat))
        raw = torch.tanh(feat[:, :, 0] * 0.7 + feat[:, :, 1] * 0.2)
        influence = A.mean(1)[0]
        if self.idea:
            xi = torch.clamp(self.kappa / (self.eps + influence), self.xmin, self.xmax)
            msg = torch.clamp(raw * xi[None, :], -1., 1.)
        else:
            xi = torch.ones_like(influence)
            msg = raw
        agg = torch.einsum('bij,bj->bi', A, msg)
        pooled = h.mean(1) + self.msg_proj(agg.unsqueeze(-1)).mean(1)
        self.last = {'raw': raw.detach(), 'msg': msg.detach(), 'xi': xi.detach(), 'influence': influence.detach()}
        return self.out(pooled)


def run_one(seed, idea, lr, kappa=2.0, epochs=18):
    torch.manual_seed(10000 + seed); np.random.seed(10000 + seed); random.seed(10000 + seed)
    ds = torch_dataset(seed)
    model = GraphMessageNet(idea=idea, kappa=kappa)
    model, metric, _ = train_model(model, ds, epochs=epochs, lr=lr, batch=64,
                                   weight_decay=1e-4, log=lambda *_: None)
    if model is None or metric is None:
        return float('nan'), None
    dev = next(model.parameters()).device
    with torch.no_grad(): _ = model(ds['xte'].to(dev))
    return float(metric), model


def main():
    grid = [{'lr': 1e-3, 'kappa': 2.0}, {'lr': 3e-3, 'kappa': 2.0}, {'lr': 1e-2, 'kappa': 2.0}]
    def base_factory(cfg):
        return lambda seed: run_one(seed, False, cfg['lr'])[0]
    baseline = sweep_baseline(base_factory, grid)
    best_lr = baseline['best_cfg']['lr']
    idea_cfgs = [{'lr': best_lr, 'kappa': 1.0}, {'lr': best_lr, 'kappa': 2.0}, {'lr': best_lr, 'kappa': 4.0}]
    idea_trials = []
    for cfg in idea_cfgs:
        r = evaluate(lambda seed, c=cfg: run_one(seed, True, c['lr'], c['kappa'])[0])
        idea_trials.append({'cfg': cfg, **r})
    idea = min(idea_trials, key=lambda r: r['mean'])
    sig_rows = []
    for seed in range(8):
        _, m = run_one(seed, True, idea['cfg']['lr'], idea['cfg']['kappa'])
        z = m.last
        raw = z['raw'].cpu().numpy().ravel(); msg = z['msg'].cpu().numpy().ravel()
        xi0 = z['xi'].cpu().numpy().ravel(); a0 = z['influence'].cpu().numpy().ravel()
        reps = raw.size // xi0.size
        xi = np.tile(xi0, reps); a = np.tile(a0, reps)
        unsat = np.abs(raw * xi) < .98
        slope = float(np.mean(np.abs(msg[unsat]) / (np.abs(raw[unsat]) + 1e-8))) if unsat.any() else 0.
        sig_rows.append({'slope_observed': slope, 'xi_mean': float(xi.mean()),
                         'sat_fraction': float((np.abs(raw * xi) >= .98).mean()),
                         'bounded_max': float(np.abs(msg).max()),
                         'corr_influence_xi': float(np.corrcoef(a, xi)[0, 1])})
    predicted = float(np.mean([r['xi_mean'] for r in sig_rows])); observed = float(np.mean([r['slope_observed'] for r in sig_rows]))
    relerr = abs(observed-predicted)/(abs(predicted)+1e-8)
    signature = {'prediction': 'trained unsaturated message gain tracks xi=kappa/(epsilon+a), and messages stay in [-1,1]',
                 'predicted_mean_xi': predicted, 'observed_mean_gain': observed,
                 'relative_gain_error': relerr, 'mean_saturation_fraction': float(np.mean([r['sat_fraction'] for r in sig_rows])),
                 'max_bounded_message': float(max(r['bounded_max'] for r in sig_rows)), 'per_seed': sig_rows,
                 'confirmed': bool(relerr < .15 and max(r['bounded_max'] for r in sig_rows) <= 1.00001)}
    rep = make_report('influence_graph_regression', 'custom_graph_message_net', baseline,
                      {'cfg': idea['cfg'], 'mean': idea['mean'], 'std': idea['std'], 'per_seed': idea['per_seed'], 'n': idea['n']}, signature)
    rep['custom_track'] = {'name': META['name'], 'file': 'graph_strategic_bench.py', 'domain': META['domain']}
    Path('bench_report.json').write_text(json.dumps(rep, indent=2)); print(json.dumps(rep, indent=2))

if __name__ == '__main__': main()