"""Train a 32-parameter top-1 MoE diagnostic toy, not a language model.

CPU/NumPy only. Saves data, weights, traces, and an inspectable HTML heatmap.
Discrete choices/capacity masks are locally constant for differentiation.
"""
import argparse
import json
import math
import platform
from pathlib import Path

import numpy as np


def softmax(z):
    m = z.max(axis=1, keepdims=True)
    e = np.exp(z - m)
    return e / e.sum(axis=1, keepdims=True), (m[:, 0] + np.log(e.sum(axis=1)))


def dispatch(p, factor=1.25):
    n, experts = p.shape
    capacity = math.ceil(factor * n / experts)
    selected = p.argmax(axis=1)  # ties go to lowest index, intentionally explicit
    onehot = np.eye(experts)[selected]
    positions = (onehot.cumsum(axis=0) * onehot).sum(axis=1)
    accepted = positions <= capacity
    pre = onehot.sum(axis=0).astype(int)
    post = (onehot * accepted[:, None]).sum(axis=0).astype(int)
    return selected, accepted, pre, post, capacity


def forward(w, v, x, y, alpha=0.0, beta=0.0, factor=1.25):
    n, experts = len(x), w.shape[1]
    p, logz = softmax(x @ w)
    selected, accepted, pre, post, capacity = dispatch(p, factor)
    f = pre / n
    probability = p.mean(axis=0)
    gate = p[np.arange(n), selected]
    expert_out = np.zeros(n)
    for e in range(experts):
        idx = (selected == e) & accepted
        expert_out[idx] = x[idx] @ v[e]
    pred = 0.25 * x[:, 0] + gate * expert_out
    error = pred - y
    task = np.mean(error ** 2)
    bal = experts * np.dot(f, probability)  # BEFORE capacity, f stop-gradient
    zloss = np.mean(logz ** 2)
    gp = np.broadcast_to(alpha * experts * f / n, p.shape).copy()
    gp[np.arange(n), selected] += 2 * error * expert_out / n
    gz = p * (gp - (gp * p).sum(axis=1, keepdims=True))
    gz += beta * 2 * logz[:, None] * p / n
    gw = x.T @ gz
    gv = np.zeros_like(v)
    for e in range(experts):
        idx = (selected == e) & accepted
        gv[e] = x[idx].T @ (2 * error[idx] * gate[idx] / n)
    metrics = dict(mse=float(task), balance=float(bal), z_loss=float(zloss),
                   pre=pre.tolist(), post=post.tolist(), capacity=capacity,
                   drop=float(1 - accepted.mean()), max_share=float(f.max()),
                   load_cv=float(f.std() / f.mean()), dead=int((pre == 0).sum()),
                   entropy=float(-(p * np.log(np.maximum(p, 1e-300))).sum(axis=1).mean()))
    trace = dict(probabilities=p.tolist(), selected=selected.tolist(),
                 accepted=accepted.tolist(), predictions=pred.tolist())
    return task + alpha * bal + beta * zloss, (gw, gv), metrics, trace


def make_data():
    rng = np.random.default_rng(7000)
    result = {}
    for split, n in [('train', 256), ('dev', 128), ('test', 256)]:
        x = np.column_stack([rng.normal(size=(n, 3)), np.ones(n)])
        group = (x[:, 0] > 0).astype(int) + 2 * (x[:, 1] > 0).astype(int)
        y = 0.25*x[:, 0] + np.array([1., -1., 2., -2.])[group]*x[:, 2]
        y += 0.4*np.array([1., -1., -1., 1.])[group]
        result[split] = dict(x=x.tolist(), y=y.tolist(), group=group.tolist())
    return result


def evaluate(w, v, data, factor=1.25):
    x, y = np.array(data['x']), np.array(data['y'])
    metrics, traces = [], []
    for start in range(0, len(x), 64):
        _, _, m, t = forward(w, v, x[start:start+64], y[start:start+64], factor=factor)
        metrics.append(m)
        traces.append(t)
    avg = {k: float(np.mean([m[k] for m in metrics])) for k in
           ['mse', 'balance', 'z_loss', 'drop', 'max_share', 'load_cv', 'dead', 'entropy']}
    avg['pre'] = np.sum([m['pre'] for m in metrics], axis=0).tolist()
    avg['post'] = np.sum([m['post'] for m in metrics], axis=0).tolist()
    avg['batches'] = metrics
    return avg, traces


def checks():
    p, _ = softmax(np.zeros((64, 4)))
    s, a, pre, post, cap = dispatch(p)
    bal = 4 * np.dot(pre/64, p.mean(axis=0))
    assert bal == 1 and pre.tolist() == [64, 0, 0, 0]
    assert cap == 20 and post.tolist() == [20, 0, 0, 0] and (~a).sum() == 44
    # High z-loss can coexist with IDENTICAL probabilities/routes after a shift.
    p2, logz2 = softmax(np.full((64, 4), 100.0))
    assert np.array_equal(p, p2)
    # Counterexample to measuring balance only after capacity clipping.
    indices = np.repeat(np.arange(4), [28, 12, 12, 12])
    q = np.eye(4)[indices] * 0.96 + 0.01
    _, _, before, after, _ = dispatch(q, 0.5)
    assert before.tolist() == [28, 12, 12, 12] and after.tolist() == [8, 8, 8, 8]
    # Dropped-to-zero argmax would mislabel overflow as expert zero.
    original = np.full(64, 2)
    mask = np.eye(4)[original]
    mask[20:] = 0
    assert np.bincount(mask.argmax(1), minlength=4).tolist() == [44, 0, 20, 0]
    rng = np.random.default_rng(170)
    x = np.column_stack([rng.normal(size=(16, 3)), np.ones(16)])
    y = rng.normal(size=16)
    w, v = rng.normal(size=(4, 4)), rng.normal(size=(4, 4))
    args = (x, y, 0.1, 0.001)
    _, grads, _, base_trace = forward(w, v, *args)
    maximum = 0.0
    for param, grad in zip([w, v], grads):
        for index in np.ndindex(param.shape):
            old, eps = param[index], 1e-6
            param[index] = old + eps
            plus, _, _, tp = forward(w, v, *args)
            param[index] = old - eps
            minus, _, _, tm = forward(w, v, *args)
            param[index] = old
            assert tp['selected'] == tm['selected'] == base_trace['selected']
            assert tp['accepted'] == tm['accepted'] == base_trace['accepted']
            maximum = max(maximum, abs((plus-minus)/(2*eps) - grad[index]))
    assert maximum < 1e-7
    # Zero expert gradient on an entirely unused expert; residual on dropped rows.
    w = np.zeros((4, 4)); w[3, 0] = 8
    _, (_, gv), _, tr = forward(w, v, x, y)
    assert np.all(gv[1:] == 0)
    dropped = ~np.array(tr['accepted'])
    assert np.allclose(np.array(tr['predictions'])[dropped], 0.25*x[dropped, 0])
    return dict(gradient_max_abs_error=maximum, uniform_balance=bal,
                uniform_pre=pre.tolist(), uniform_post=post.tolist(), uniform_drop=44/64,
                shifted_z_loss=float(np.mean(logz2**2)), post_clip_counterexample=[before.tolist(), after.tolist()],
                checks='finite difference, tie, capacity, overflow index, shift, dead gradient, residual: passed')


def train(data, seed, name, alpha, beta, out):
    rng = np.random.default_rng(seed)
    w = rng.normal(0, 0.03, (4, 4)); w[3, 0] += 2
    v = rng.normal(0, 0.1, (4, 4))
    moments = [np.zeros_like(w), np.zeros_like(v)]
    variances = [np.zeros_like(w), np.zeros_like(v)]
    x, y = np.array(data['train']['x']), np.array(data['train']['y'])
    history = []
    for step in range(601):
        if step % 100 == 0:
            dev, _ = evaluate(w, v, data['dev'])
            train_metrics, _ = evaluate(w, v, data['train'])
            history.append(dict(step=step, dev=dev, train=train_metrics))
        if step == 600:
            break
        idx = rng.choice(len(x), 64, replace=False)
        _, grads, _, _ = forward(w, v, x[idx], y[idx], alpha, beta)
        for i, (param, grad) in enumerate(zip([w, v], grads)):
            moments[i] = 0.9*moments[i] + 0.1*grad
            variances[i] = 0.999*variances[i] + 0.001*grad**2
            mhat = moments[i] / (1 - 0.9**(step+1))
            vhat = variances[i] / (1 - 0.999**(step+1))
            param -= 0.03 * mhat / (np.sqrt(vhat) + 1e-8)
    test, traces = evaluate(w, v, data['test'])
    capacity_sweep = {str(c): evaluate(w, v, data['test'], c)[0] for c in [0.5, 1.0, 1.25, 2.0, 4.0]}
    key = f'{name}_{seed}'
    result = dict(name=name, seed=seed, alpha=alpha, beta=beta, steps=600,
                  history=history, test=test, capacity_sweep=capacity_sweep, traces=traces)
    (out/f'{key}_weights.json').write_text(json.dumps(dict(w=w.tolist(), v=v.tolist()), indent=2))
    (out/f'{key}.json').write_text(json.dumps(result, indent=2))
    return result


def html_report(results, out):
    # A numeric HTML table, not an image and not a substitute for the imagegen figures.
    sections = []
    for r in results:
        rows = []
        t = r['traces'][0]
        for j, probs in enumerate(t['probabilities']):
            cells = ''.join(f'<td style="background:rgba(0,135,170,{p:.4f})">{p:.4f}</td>' for p in probs)
            rows.append(f'<tr><th>{j}</th>{cells}<td>{t["selected"][j]}</td><td>{t["accepted"][j]}</td></tr>')
        sections.append(f'<details><summary>{r["name"]}, seed {r["seed"]}: MSE {r["test"]["mse"]:.4f}, drop {r["test"]["drop"]:.3f}</summary>'
                        '<table><thead><tr><th>Token</th><th>E0</th><th>E1</th><th>E2</th><th>E3</th><th>Choice</th><th>Accepted</th></tr></thead><tbody>'
                        + ''.join(rows) + '</tbody></table></details>')
    html = '<!doctype html><html lang="en"><meta charset="utf-8"><title>MoE actual routing traces</title>'
    html += '<style>body{font:16px system-ui;max-width:960px;margin:40px auto;padding:16px;color:#122c3d}table{border-collapse:collapse;margin:20px 0}td,th{padding:5px 14px;border:1px solid #ccd}summary{padding:16px;cursor:pointer}details{border-bottom:1px solid #ccd}</style>'
    html += '<h1>MoE actual routing traces</h1><p>First fixed test batch: 64 synthetic tokens, four experts. All runs use capacity 20 per expert. Cell opacity equals probability (0 to 1); numeric values remain visible. Choice is BEFORE capacity. False means expert branch skipped, residual retained. These are toy results, not paper benchmarks. Expand any run.</p>'
    (out/'routing_heatmap.html').write_text(html + ''.join(sections) + '</html>')


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', type=Path, default=Path(__file__).parent/'results')
    parser.add_argument('--check-only', action='store_true')
    args = parser.parse_args()
    check = checks()
    print(json.dumps(check, indent=2))
    if args.check_only:
        return
    out = args.output; out.mkdir(parents=True, exist_ok=True)
    data = make_data()
    (out/'data.json').write_text(json.dumps(data, indent=2))
    print('CPU', platform.python_version(), 'NumPy', np.__version__, 'X=[64,4], logits=[64,4], experts=[4,4], parameters=32')
    results = [train(data, seed, name, a, b, out) for seed in [70,71,72]
               for name, a, b in [('task', 0, 0), ('balance', 0.1, 0), ('balance_z', 0.1, 0.001)]]
    summary = {}
    for name in ['task', 'balance', 'balance_z']:
        subset = [r for r in results if r['name'] == name]
        summary[name] = {k: dict(mean=float(np.mean([r['test'][k] for r in subset])),
                                sd=float(np.std([r['test'][k] for r in subset], ddof=1)))
                         for k in ['mse','drop','max_share','load_cv','balance','z_loss','entropy']}
        print(name, json.dumps(summary[name]))
    (out/'summary.json').write_text(json.dumps(dict(checks=check, results=summary,
        environment=dict(python=platform.python_version(), numpy=np.__version__, device='CPU', dtype='float64'),
        protocol='Fixed600 steps; Adam0.03; batch64; capacity1.25; same data/initial weights/batch stream within seed; no dev/test selection'), indent=2))
    html_report(results, out)


if __name__ == '__main__':
    main()