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

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

SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
N_REPLICAS = 4
EPOCHS = 12
BATCH = 128
# The union of all lr and method knobs is shared by both sides.
GRID = [{'lr': lr, 'gamma': gamma} for lr in (0.001, 0.003, 0.006)
        for gamma in (0.05, 0.10)]
BASE_W = np.array([[0., 1.0, .7, .4], [1.0, 0., .8, .5],
                   [.7, .8, 0., 1.1], [.4, .5, 1.1, 0.]], dtype=np.float32)
SIGNATURE = []


def laplacian(w, stale_degree=False):
    degree = BASE_W.sum(axis=1) if stale_degree else w.sum(axis=1)
    return np.diag(degree) - w


def run_replicas(seed, cfg, preserve, collect=False):
    """Train four identical MLP systems on disjoint data shards.
    preserve=True is the proposed rebuilt-degree communication rule.
    """
    global SIGNATURE
    torch.manual_seed(10000 + int(seed))
    np.random.seed(20000 + int(seed))
    d = get_dataset('tabular', int(seed), n_train=400, n_test=200)
    xtr, ytr = d['xtr'].float(), d['ytr'].float()
    xte, yte = d['xte'].float(), d['yte'].float()
    nets = [make_model('mlp_tiny', d['input_shape'], d['out_dim']) for _ in range(N_REPLICAS)]
    # Same initialization is intentional; local data make disagreement naturally.
    state = nets[0].state_dict()
    for net in nets[1:]: net.load_state_dict(state)
    params = [[p for p in net.parameters()] for net in nets]
    rng = np.random.default_rng(700000 + int(seed))
    shard_idx = [np.arange(r * len(xtr)//N_REPLICAS, (r+1)*len(xtr)//N_REPLICAS)
                 for r in range(N_REPLICAS)]
    loss_fn = nn.MSELoss()
    for epoch in range(EPOCHS):
        # Identical permutation/dropout schedule for paired systems (seed controls it).
        perms = [rng.permutation(ix) for ix in shard_idx]
        for pos in range(0, len(perms[0]), BATCH):
            grads = []
            for r, net in enumerate(nets):
                ix = perms[r][pos:pos+BATCH]
                net.zero_grad(set_to_none=True)
                loss = loss_fn(net(xtr[ix]), ytr[ix])
                loss.backward()
                grads.append([p.grad.detach().clone() for p in net.parameters()])
            mask = (rng.random((N_REPLICAS, N_REPLICAS)) > .30).astype(np.float32)
            mask = np.triu(mask, 1); mask = mask + mask.T
            w = BASE_W * mask
            L = laplacian(w, stale_degree=not preserve)
            with torch.no_grad():
                flat = torch.cat([p.detach().reshape(-1) for p in nets[0].parameters()])
                theta = torch.stack([torch.cat([p.detach().reshape(-1) for p in net.parameters()]) for net in nets])
                mean_theta = theta.mean(0)
                defect = float(np.linalg.norm(L.sum(axis=1)))
                predicted = float(cfg['lr'] * cfg['gamma'] * np.linalg.norm(
                    L.sum(axis=1)[:, None] * mean_theta.cpu().numpy()[None, :]))
                coupling = torch.einsum('ij,jp->ip', torch.tensor(L, dtype=theta.dtype), theta)
                observed = float(cfg['lr'] * cfg['gamma'] * coupling.mean(0).norm().item())
                if collect: SIGNATURE.append((defect, predicted, observed))
                for r, net in enumerate(nets):
                    off = torch.einsum('j,jp->p', torch.tensor(L[r], dtype=theta.dtype), theta)
                    cursor = 0
                    for p, g in zip(params[r], grads[r]):
                        n = p.numel()
                        p.add_(-cfg['lr'] * (g + cfg['gamma'] * off[cursor:cursor+n].view_as(p)))
                        cursor += n
    with torch.no_grad():
        pred = torch.stack([net(xte) for net in nets]).mean(0)
        metric = float(loss_fn(pred, yte).item())
    return metric


def make_train(cfg, preserve, collect=False):
    return lambda seed: run_replicas(seed, cfg, preserve, collect=collect)


def main():
    # Cheap core structural sanity check FIRST, before neural training.
    rng = np.random.default_rng(1127)
    residuals, bad = [], []
    for _ in range(200):
        m = np.triu((rng.random((4,4)) > .3).astype(np.float32), 1); m += m.T
        wt = BASE_W * m
        residuals.append(np.linalg.norm(laplacian(wt) @ np.ones(4)))
        bad.append(np.linalg.norm(laplacian(wt, True) @ np.ones(4)))
    sanity = {'preserve_max_L1': float(max(residuals)), 'baseline_mean_L1': float(np.mean(bad))}
    base = sweep_baseline(lambda c: make_train(c, False), GRID, seeds=SWEEP_SEEDS)
    # Evaluate every idea setting on full paired seeds; choose best, all lrs were in baseline grid.
    idea_trials = []
    for cfg in GRID:
        r = evaluate(make_train(cfg, True), seeds=SEEDS)
        idea_trials.append({'cfg': cfg, **r})
    best = min(idea_trials, key=lambda z: z['mean'])
    # Collect mechanism numbers from the trained best baseline/idea behavior.
    SIGNATURE.clear(); base_sig = evaluate(make_train(best['cfg'], False, collect=True), seeds=SEEDS)
    arr = np.asarray(SIGNATURE, dtype=float)
    corr = float(np.corrcoef(arr[:,1], arr[:,2])[0,1]) if len(arr) > 2 and np.std(arr[:,1]) > 0 else float('nan')
    sig = {'prediction': 'common-mode coupling forcing proportional to row-sum defect',
           'trained_samples': int(len(arr)), 'predicted_mean': float(arr[:,1].mean()),
           'observed_mean': float(arr[:,2].mean()), 'predicted_observed_corr': corr,
           'preserving_L1_max': sanity['preserve_max_L1'],
           'confirmed': bool(sanity['preserve_max_L1'] < 1e-5 and corr > 0.8)}
    report = make_report('tabular', 'mlp_tiny', base, best,
                         {'mechanism_signature': sig,
                          'sanity_check': sanity,
                          'idea_trials': idea_trials,
                          'protocol': {'paired_seeds': list(SEEDS), 'epochs': EPOCHS,
                                       'replicas': N_REPLICAS, 'batch': BATCH,
                                       'intervention': 'rebuild degree after dropout'}})
    with open('bench_report.json', 'w') as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()