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

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

HERE = Path(__file__).parent
reload_custom_tracks()
TRACK = 'correlated_token_moe_regression'
SEEDS = tuple(range(8))
# Shared union: all idea learning rates are also baseline sweep values.
GRID = [{'lr': 1e-3, 'epochs': 20}, {'lr': 3e-3, 'epochs': 20}, {'lr': 1e-2, 'epochs': 20}]

class Baseline(nn.Module):
    """Standard label-sensitive flattened MLP, matched width/depth."""
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(64, 64), nn.Tanh(),
                                 nn.Linear(64, 64), nn.Tanh(),
                                 nn.Linear(64, 1))
    def forward(self, x):
        return self.net(x.reshape(x.shape[0], -1))

class Equivariant(nn.Module):
    """Permutation-equivariant token encoder followed by invariant pooling."""
    def __init__(self):
        super().__init__()
        self.token = nn.Sequential(nn.Linear(4, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh())
        self.pair = nn.Sequential(nn.Linear(4, 32), nn.Tanh(), nn.Linear(32, 16), nn.Tanh())
        self.readout = nn.Sequential(nn.Linear(4 + 64 + 16, 64), nn.Tanh(),
                                     nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1))
    def forward(self, x):
        h = self.token(x)
        mean_x = x.mean(dim=1)
        dif = x[:, :, None, :] - x[:, None, :, :]
        messages = self.pair(dif).mean(dim=2).mean(dim=1)
        return self.readout(torch.cat([mean_x, h.mean(dim=1), messages], dim=1))

def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)

def run(kind, cfg, seed, keep=False):
    seed_all(seed)
    ds = get_dataset(TRACK, seed, 400, 200)
    model = Baseline() if kind == 'baseline' else Equivariant()
    net, metric, _ = train_model(model, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128)
    if net is None:
        raise RuntimeError('benchmark training failed')
    return (float(metric), net, ds) if keep else float(metric)

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

def main():
    baseline = sweep_baseline(factory('baseline'), GRID)
    idea_configs = []
    for cfg in GRID:
        res = evaluate(lambda seed, cfg=cfg: run('idea', cfg, seed), SEEDS)
        idea_configs.append((cfg, res))
    best_cfg, idea = min(idea_configs, key=lambda z: z[1]['mean'])

    base_swap, idea_swap = [], []
    for seed in SEEDS:
        _, bnet, ds = run('baseline', baseline['best_cfg'], seed, True)
        _, inet, _ = run('idea', best_cfg, seed, True)
        x = ds['xte'][:64]
        xp = x[:, [1, 0] + list(range(2, x.shape[1]))]
        bd = next(bnet.parameters()).device; idv = next(inet.parameters()).device
        with torch.no_grad():
            base_swap.append(float((bnet(x.to(bd)) - bnet(xp.to(bd))).abs().mean()))
            idea_swap.append(float((inet(x.to(idv)) - inet(xp.to(idv))).abs().mean()))
    signature = {
        'quantity': 'mean absolute prediction change under swapping two trained-task tokens',
        'prediction': 'equivariant/invariant system has zero change while flattened baseline is nonzero',
        'baseline_per_seed': base_swap,
        'idea_per_seed': idea_swap,
        'baseline_mean': float(np.mean(base_swap)),
        'idea_mean': float(np.mean(idea_swap)),
        'confirmed': bool(np.mean(idea_swap) < 1e-6 and np.mean(base_swap) > 1e-5)
    }
    report = make_report(TRACK, 'matched_custom', baseline, idea, signature)
    report['idea_sweep'] = [{'cfg': cfg, 'mean': res['mean'], 'std': res['std'], 'per_seed': res['per_seed']} for cfg, res in idea_configs]
    report['track_justification'] = 'The registered correlated-token track contains exchangeable multi-token correlations; built-in tracks do not.'
    (HERE / 'bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()