import json
import sys
from pathlib import Path
import numpy as np
import torch

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

TRACK = 'tabular'
MODEL = 'mlp_tiny'
EPOCHS = 20
BATCH = 128
# Shared union: every lr and scale tested by the idea is also tested by baseline.
GRID = [
    {'lr': 1e-3, 'scale': 0.7},
    {'lr': 3e-3, 'scale': 1.0},
    {'lr': 1e-2, 'scale': 1.3},
]


def coulomb_bank(n, beta=2.0, eta=0.02, steps=180, eps=1e-5, seed=0):
    rng = np.random.default_rng(seed)
    a = np.linspace(0, 2*np.pi, n, endpoint=False) + rng.normal(0, .03, n)
    z = np.c_[np.cos(a), np.sin(a)] * np.sqrt(n / 2.0)
    for _ in range(steps):
        d = z[:, None, :] - z[None, :, :]
        r2 = np.sum(d*d, axis=-1) + eps
        np.fill_diagonal(r2, np.inf)
        rep = np.sum(d / r2[:, :, None], axis=1) / n
        z += eta * (-z + rep) + np.sqrt(2 * eta / beta) * rng.normal(size=z.shape)
    # Normalize to unit per-coordinate RMS, matching a conventional small init.
    z /= max(np.sqrt(np.mean(z*z)), 1e-8)
    return z


def set_first_layer(net, mode, scale, seed):
    # A 2-D Coulomb bank is lifted to the 10-D input rows by a fixed random
    # orthonormal projection. The only system difference is this initialization.
    layer = net[0]
    with torch.no_grad():
        if mode == 'coulomb':
            z = coulomb_bank(layer.out_features, seed=seed)
            rng = np.random.default_rng(seed + 991)
            q, _ = np.linalg.qr(rng.normal(size=(layer.in_features, 2)))
            w = z @ q.T
            layer.weight.copy_(torch.tensor(w * (scale * 0.12), dtype=layer.weight.dtype))
            layer.bias.zero_()
        else:
            torch.manual_seed(seed + 991)
            torch.nn.init.normal_(layer.weight, mean=0.0, std=scale * 0.12)
            torch.nn.init.zeros_(layer.bias)


def train_one(seed, cfg, mode, return_net=False):
    np.random.seed(seed)
    torch.manual_seed(seed)
    ds = get_dataset(TRACK, seed, n_train=400, n_test=400)
    net = make_model(MODEL, ds['input_shape'], ds['out_dim'])
    set_first_layer(net, mode, cfg['scale'], seed)
    net, metric, hist = train_model(net, ds, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH,
                                    log=lambda *_: None)
    if return_net:
        return metric, net
    return metric


def baseline_factory(cfg):
    return lambda seed: train_one(seed, cfg, 'gaussian')


def stats_from_net(net):
    w = net[0].weight.detach().cpu().numpy()
    d = np.sqrt(((w[:, None] - w[None, :])**2).sum(-1))
    np.fill_diagonal(d, np.inf)
    return float(np.median(d.min(axis=1))), float(np.mean(d.min(axis=1))), float(np.std(w, axis=0).mean())


def main():
    # Baseline is tuned on the protocol's four-seed sweep, then re-evaluated on 8.
    base = sweep_baseline(baseline_factory, GRID)
    best = base['best_cfg']
    idea_grid = [best, {'lr': 1e-3, 'scale': 0.7}, {'lr': 1e-2, 'scale': 1.3}]
    # Deduplicate while preserving the required three-config idea budget.
    uniq = []
    for c in idea_grid:
        if c not in uniq: uniq.append(c)
    idea_cfg = min(uniq, key=lambda c: np.mean([train_one(s, c, 'coulomb') for s in (0,1,2,3)]))
    idea_vals = [train_one(s, idea_cfg, 'coulomb') for s in DEFAULT_SEEDS]
    # Paired trained-model signature: compare final first-layer nearest-neighbor spacing.
    bnn, inn = [], []
    for s in DEFAULT_SEEDS:
        _, bn = train_one(s, best, 'gaussian', True)
        _, cn = train_one(s, idea_cfg, 'coulomb', True)
        bnn.append(stats_from_net(bn))
        inn.append(stats_from_net(cn))
    bmean = [float(np.mean(x)) for x in zip(*bnn)]
    imean = [float(np.mean(x)) for x in zip(*inn)]
    # The stage-1 prediction is higher non-collapse spacing after training;
    # confirmation requires a positive, practically nontrivial observed ratio.
    ratio = imean[0] / max(bmean[0], 1e-12)
    sig = {'quantity': 'trained first-layer median nearest-neighbor distance',
           'prediction': 'Coulomb repulsion preserves larger inter-row spacing as N=64 bank grows',
           'baseline_observed': bmean[0], 'idea_observed': imean[0],
           'observed_ratio_idea_over_baseline': ratio,
           'confirmed': bool(ratio > 1.05),
           'n_models': 8}
    report = make_report(TRACK, MODEL, base, {'cfg': idea_cfg, 'mean': float(np.mean(idea_vals)),
                                             'std': float(np.std(idea_vals)), 'per_seed': [float(x) for x in idea_vals],
                                             'n': len(idea_vals)},
                         {'mechanism_signature': sig,
                          'track_choice': 'tabular: initialization is explicitly covered by the benchmark mapping; same mlp_tiny and training loop used.',
                          'idea_sweep': [{'cfg': c} for c in uniq],
                          'trained_geometry': {'baseline_median_nn': bmean[0], 'idea_median_nn': imean[0],
                                              'baseline_mean_nn': bmean[1], 'idea_mean_nn': imean[1]}})
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()