import sys, json, random
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 get_dataset, train_model, evaluate, sweep_baseline, make_report

SEEDS = tuple(range(8))
DIM, K = 64, 2
TRACK = 'cycle_topology_graph'
MODEL = 'mlp_tiny'


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


class TopologicalMLP(nn.Module):
    """Same network in both arms; only latent transition dissipation differs."""
    def __init__(self, input_dim, out_dim, mode, damping):
        super().__init__()
        self.mode, self.damping = mode, damping
        self.enc = nn.Sequential(nn.Linear(input_dim, DIM), nn.ReLU(), nn.Linear(DIM, DIM))
        self.interaction = nn.Sequential(nn.Tanh(), nn.Linear(DIM, DIM), nn.Tanh())
        self.head = nn.Linear(DIM, out_dim)
        H = torch.zeros(DIM, K); H[0, 0] = 1.; H[1, 1] = 1.
        self.register_buffer('H', H)
        self.register_buffer('PH', H @ H.T)
        self.register_buffer('PP', torch.eye(DIM) - H @ H.T)

    def forward(self, x):
        z = self.enc(x)
        n = self.interaction(z)
        if self.mode == 'baseline':
            znext = z + n - self.damping * z
        else:
            h = z @ self.PH
            u = z @ self.PP
            # Projected interaction and damping preserve the harmonic branch.
            nperp = n @ self.PP
            znext = h + u + nperp - self.damping * u
        return self.head(znext)

    def signature(self, x):
        with torch.no_grad():
            x = x.to(next(self.parameters()).device)
            z = self.enc(x); n = self.interaction(z)
            h = z @ self.PH; u = z @ self.PP
            if self.mode == 'baseline':
                step = n - self.damping * z
            else:
                step = n @ self.PP - self.damping * u
            harmonic_step = step @ self.PH
            diss_step = step @ self.PP
            return float(harmonic_step.norm(dim=1).mean()), float(diss_step.norm(dim=1).mean()), float(h.norm(dim=1).mean())


def train_one(mode, cfg, seed, return_model=False):
    seed_all(seed)
    d = get_dataset(TRACK, seed, n_train=400, n_test=400)
    net = TopologicalMLP(int(np.prod(d['input_shape'])), d['out_dim'], mode, cfg['damping'])
    net, metric, hist = train_model(net, d, epochs=cfg['epochs'], lr=cfg['lr'], batch=128)
    if return_model:
        return metric, net, d
    return metric


def main():
    # Union parity: every lr and damping value is tried by baseline and idea.
    grid = [{'lr': lr, 'damping': damp, 'epochs': 18}
            for lr in (1e-3, 3e-3, 1e-2) for damp in (0.05, 0.2, 0.5)]
    base = sweep_baseline(lambda c: lambda s: train_one('baseline', c, s), grid, seeds=(0,1,2,3))
    idea_cfgs = grid
    idea_trials = []
    for cfg in idea_cfgs:
        r = evaluate(lambda s, c=cfg: train_one('idea', c, s), seeds=(0,1,2,3))
        idea_trials.append({'cfg': cfg, 'mean': r['mean']})
    best = min(idea_trials, key=lambda q: q['mean'])['cfg']
    idea = evaluate(lambda s: train_one('idea', best, s), seeds=SEEDS)
    # Trained-model signature: measured projected harmonic update and dissipative update.
    hs, ds, hm = [], [], []
    for s in SEEDS:
        _, model, data = train_one('idea', best, s, return_model=True)
        h, u, a = model.signature(data['xte']); hs.append(h); ds.append(u); hm.append(a)
    base['idea_grid'] = idea_trials
    base['baseline_grid_union'] = grid
    rep = make_report(TRACK, MODEL, base, idea, extra={
        'prediction': 'harmonic update is near zero while dissipative update remains nonzero',
        'trained_model_observed_mean_harmonic_step': float(np.mean(hs)),
        'trained_model_observed_mean_dissipative_step': float(np.mean(ds)),
        'trained_model_observed_mean_harmonic_amplitude': float(np.mean(hm)),
        'confirmed': bool(np.mean(hs) < 1e-7 and np.mean(ds) > 1e-5),
        'measurement': 'mean latent transition norms on each trained idea model test split'
    })
    rep['custom_track'] = {'name': TRACK, 'file': '/home/maxwelhelp/all/math2nn/bench/custom_tracks/cycle_topology_graph.py', 'domain': 'graph_topology'}
    Path('bench_report.json').write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))

if __name__ == '__main__':
    main()