import os, sys, json, time
import numpy as np
import torch
import torch.nn as nn

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

SEEDS = tuple(range(8))
SWEEP_SEEDS = (0, 1, 2, 3)


def seed_all(seed):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def sm(x):
    return torch.softmax(x, dim=-1)


class RelationalGNN(nn.Module):
    """Shared small GNN; only route() differs between baseline and idea."""
    def __init__(self, mode='baseline', tau=1.0, beta4=0.0, steps=2):
        super().__init__()
        self.mode, self.tau, self.beta4, self.steps = mode, float(tau), float(beta4), int(steps)
        self.node = nn.Sequential(nn.Linear(1, 24), nn.ReLU(), nn.Linear(24, 24))
        self.edge = nn.Sequential(nn.Linear(2 * 24 + 1, 32), nn.ReLU(), nn.Linear(32, 3))
        self.rel = nn.ModuleList([nn.Linear(24, 24, bias=False) for _ in range(3)])
        self.head = nn.Sequential(nn.Linear(24, 24), nn.ReLU(), nn.Linear(24, 2))
        self.beta = nn.Parameter(torch.zeros(3))

    def unary(self, x):
        h = self.node(x[..., 8:9])
        hi, hj = h.unsqueeze(2).expand(-1,-1,8,-1), h.unsqueeze(1).expand(-1,8,-1,-1)
        ef = torch.cat([hi, hj, x[..., :8].unsqueeze(-1)], dim=-1)
        return self.edge(ef), h

    def route(self, logits):
        p = sm(logits / self.tau)
        if self.mode == 'baseline' or self.beta4 == 0:
            return p
        n = p.shape[1]
        q = p
        for _ in range(self.steps):
            r = torch.zeros_like(q)
            for a in range(3):
                other = [z for z in range(3) if z != a]
                r[..., a] = (q[:, :, :, other[0]].unsqueeze(2) * q[:, :, :, other[1]].unsqueeze(1) +
                             q[:, :, :, other[1]].unsqueeze(2) * q[:, :, :, other[0]].unsqueeze(1)).sum(dim=3)
            score = (2.0 * self.beta.view(1,1,1,3) + (self.beta4 / n) * r) / self.tau
            q = sm(score)
        return q

    def forward(self, x):
        logits, h = self.unary(x)
        p = self.route(logits)
        a = x[..., :8]
        msgs = []
        for z in range(3):
            mz = self.rel[z](h)
            msgs.append(torch.einsum('bij,bjd->bid', a * p[..., z], mz))
        pooled = h + sum(msgs) / 8.0
        return self.head(pooled.mean(dim=1))

    @torch.no_grad()
    def route_stats(self, x, beta4=None):
        old = self.beta4
        if beta4 is not None: self.beta4 = float(beta4)
        logits, _ = self.unary(x)
        p = self.route(logits)
        self.beta4 = old
        n = p.shape[1]
        tri = []
        for i in range(n):
            for j in range(i+1,n):
                for k in range(j+1,n):
                    v = 0.
                    for a in range(3):
                        o = [z for z in range(3) if z != a]
                        v = v + p[:,i,j,a] * (p[:,i,k,o[0]]*p[:,j,k,o[1]] + p[:,i,k,o[1]]*p[:,j,k,o[0]])
                    tri.append(v)
        density = torch.stack(tri, 1).mean().item()
        return {'rainbow_density': density, 'entropy': float((-p*torch.log(p.clamp_min(1e-8))).sum(-1).mean().item()),
                'mean_motif_logit': float((self.beta4/n * torch.abs(torch.stack(tri,1))).mean().item())}


def make_ds(seed, ntr=400, nte=160):
    d = get_dataset(seed, ntr, nte)
    for k in ('xtr','ytr','xte','yte'):
        d[k] = torch.from_numpy(d[k])
    return d


def run_cfg(cfg, seed, mode):
    seed_all(seed)
    ds = make_ds(seed)
    net = RelationalGNN(mode=mode, tau=cfg['tau'], beta4=cfg.get('beta4', 0), steps=2)
    _, metric, _ = train_model(net, ds, epochs=18, lr=cfg['lr'], batch=64, log=lambda *_: None)
    return metric


def main():
    grid = [{'lr': lr, 'tau': tau} for lr in (0.0015, 0.003, 0.006) for tau in (0.5, 1.0, 2.0)]
    baseline = sweep_baseline(lambda c: lambda s: run_cfg(c, s, 'baseline'), grid, seeds=SWEEP_SEEDS)
    bc = baseline['best_cfg']
    idea_grid = [dict(bc, beta4=0.75), dict(bc, beta4=1.5), dict(bc, beta4=3.0)]
    idea_runs = []
    best_cfg, best_mean = None, float('inf')
    for c in idea_grid:
        r = evaluate(lambda s, c=c: run_cfg(c, s, 'idea'), seeds=SEEDS)
        idea_runs.append({'cfg': c, 'result': r})
        if r['mean'] < best_mean: best_mean, best_cfg = r['mean'], c
    idea = next(z['result'] for z in idea_runs if z['cfg'] == best_cfg)
    sig_rows = []
    for s in SEEDS:
        seed_all(s); ds = make_ds(s)
        net = RelationalGNN(mode='idea', tau=best_cfg['tau'], beta4=best_cfg['beta4'], steps=2)
        net, _, _ = train_model(net, ds, epochs=18, lr=best_cfg['lr'], batch=64, log=lambda *_: None)
        net.eval(); x = ds['xte'].to(next(net.parameters()).device)
        z0, z1 = net.route_stats(x, 0.0), net.route_stats(x, best_cfg['beta4'])
        sig_rows.append({'seed': s, 'density0': z0['rainbow_density'], 'density_beta4': z1['rainbow_density'],
                         'delta_density': z1['rainbow_density']-z0['rainbow_density'],
                         'entropy_beta4': z1['entropy']})
    dmean = float(np.mean([z['delta_density'] for z in sig_rows]))
    signature = {'prediction': 'trained positive beta4 increases rainbow density while beta4=0 is unary',
                 'observed_mean_density_change': dmean, 'per_seed': sig_rows,
                 'beta4_zero_max_change': 0.0,
                 'confirmed': bool(dmean > 0.001)}
    rep = make_report('relational_graph_classification', 'local_relational_gnn', baseline, idea,
                      {'mechanism_signature': signature, 'idea_sweep': idea_runs,
                       'protocol_note': 'Built-in relational custom track used; no structurally matching built-in track exists.'})
    rep['wall_clock_note'] = 'canonical train_model; equal 18 epochs, batch 64; timing not primary metric'
    with open('bench_report.json', 'w') as f: json.dump(rep, f, indent=2)
    print(json.dumps(rep, indent=2))

if __name__ == '__main__': main()