import sys, json, copy, random
from pathlib import Path
import numpy as np
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_model, train_model, evaluate, sweep_baseline, make_report
import torch

TRACK = 'tabular'
MODEL = 'mlp_tiny'
N = 8
GENERATIONS = 16
SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))
# Shared union: every idea learning rate is also evaluated by baseline.
GRID = [
    {'lr': 1e-3, 'weight_decay': 0.0},
    {'lr': 3e-3, 'weight_decay': 0.0},
    {'lr': 1e-2, 'weight_decay': 0.0},
]


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)


def baseline_one(cfg, seed):
    seed_all(seed)
    d = get_dataset(TRACK, seed, n_train=400, n_test=400)
    net = make_model(MODEL, d['input_shape'], d['out_dim'])
    net, metric, history = train_model(
        net, d, epochs=GENERATIONS, lr=cfg['lr'], batch=128,
        weight_decay=cfg['weight_decay'], log=lambda *_: None)
    return float(metric) if metric is not None else float('inf')


def flat(net):
    return torch.cat([p.detach().reshape(-1) for p in net.parameters()])


def put(net, v):
    pos = 0
    with torch.no_grad():
        for p in net.parameters():
            n = p.numel()
            p.copy_(v[pos:pos+n].reshape_as(p))
            pos += n


def train_candidate(net, vec, x, y, lr):
    """One identical SGD update for each population member."""
    put(net, vec)
    net.zero_grad(set_to_none=True)
    loss = torch.mean((net(x) - y) ** 2)
    loss.backward()
    with torch.no_grad():
        out = torch.cat([(p - lr * p.grad).reshape(-1) for p in net.parameters()])
    return out, float(loss.detach().cpu())


def idea_one(cfg, seed, return_trace=False):
    seed_all(seed)
    d = get_dataset(TRACK, seed, n_train=400, n_test=400)
    net = make_model(MODEL, d['input_shape'], d['out_dim'])
    x, y = d['xtr'], d['ytr']
    center = flat(net).clone()
    # A priori fixed controls from the proposal.
    dmin, dmax, gamma, target = 1e-7, 3e-4, 0.20, 2e-4
    kappa = 1.5
    rng = np.random.default_rng(seed + 991)
    population = center[None, :] + 0.02 * torch.randn((N, center.numel()))
    records = []
    for step in range(GENERATIONS):
        rewards = []
        updated = []
        for i in range(N):
            # Candidate training is the population optimizer's training step.
            v, train_loss = train_candidate(net, population[i], x, y, cfg['lr'])
            updated.append(v)
            rewards.append(-train_loss)
        population = torch.stack(updated)
        r = np.asarray(rewards, dtype=np.float64)
        # Stable normalized selection weights.
        w = np.exp(kappa * (r - r.max()))
        w /= w.sum()
        z = population.detach().cpu().numpy()
        mu = np.sum(w[:, None] * z, axis=0)
        c = np.sum(w[:, None] * (z - mu) ** 2, axis=0)
        q = (z - mu) ** 2 - c
        rbar = float(np.sum(w * r))
        r2 = np.sum(w[:, None] * (r[:, None] - rbar) * q, axis=0) / (np.sum(w[:, None] * q*q, axis=0) + 1e-8)
        r2 = np.clip(r2, -8, 8)
        rv = float(np.sum(w * (r-rbar)**2))
        D = np.clip(dmin + gamma * rv / (1 + np.abs(r2)), dmin, dmax)
        ess = float(1 / np.sum(w*w))
        # Diversity safeguard, reflecting +2D in the variance equation.
        D = np.maximum(D, np.maximum(0, (target-c) / 2.0))
        D = np.clip(D, dmin, dmax)
        # Resample selected candidates, then apply coordinatewise diffusion.
        idx = rng.choice(N, size=N, p=w)
        noise = torch.as_tensor(rng.normal(size=(N, center.numel())), dtype=population.dtype)
        population = population[idx] + noise * torch.as_tensor(np.sqrt(2*D), dtype=population.dtype)[None, :]
        if step == GENERATIONS-1 or step % 4 == 0:
            records.append({'step': step, 'cov_trace': float(np.mean(c)), 'ess': ess,
                            'reward_var': rv, 'mutation_D': float(np.mean(D)),
                            'curvature_median': float(np.median(r2))})
    # Evaluate the trained selected system, not a readout of baseline weights.
    final_w = np.exp(kappa * (r-r.max())); final_w /= final_w.sum()
    best = population[int(np.argmax(final_w))]
    put(net, best)
    with torch.no_grad():
        metric = float(torch.mean((net(d['xte'])-d['yte'])**2).item())
    return (metric, records) if return_trace else metric


def mechanism_signature(cfg, seed=0):
    metric, trace = idea_one(cfg, seed, True)
    # Re-test the claimed additive diffusion relation on trained NN parameters:
    # measured coordinate variance increment from mutation versus predicted 2D.
    if len(trace) >= 2:
        a, b = trace[-2], trace[-1]
        observed = b['cov_trace'] - a['cov_trace']
        predicted = 2.0 * b['mutation_D']
    else:
        observed = predicted = float('nan')
    rel = abs(observed-predicted)/(abs(predicted)+1e-12)
    return {'metric': metric, 'observed_cov_increment': observed,
            'predicted_2D_increment': predicted, 'relative_error': rel,
            'curvature_median_final': trace[-1]['curvature_median'],
            'confirmed': bool(np.isfinite(rel) and rel < 0.50)}


def main():
    # Baseline sweep uses the same three lr values and a final eight-seed rerun.
    base = sweep_baseline(
        lambda cfg: (lambda seed: baseline_one(cfg, seed)), GRID, seeds=SWEEP_SEEDS)
    idea_trials = []
    for cfg in GRID:
        res = evaluate(lambda seed, cfg=cfg: idea_one(cfg, seed), seeds=SEEDS)
        idea_trials.append({'cfg': cfg, 'result': res})
    best_trial = min(idea_trials, key=lambda x: x['result']['mean'])
    sig = mechanism_signature(best_trial['cfg'], seed=0)
    report = make_report(TRACK, MODEL, base, best_trial['result'], {
        'mechanism_signature': sig,
        'track_choice': 'tabular is the built-in optimizer/training-dynamics match; Friedman regression keeps architecture and task fixed.',
        'idea_sweep': idea_trials,
        'budget': {'baseline_epochs': GENERATIONS, 'idea_generations': GENERATIONS, 'population': N}
    })
    report['idea']['best_cfg'] = best_trial['cfg']
    report['mechanism_signature']['trained_model_metric'] = sig.pop('metric')
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()