#!/usr/bin/env python3
"""
Slice 19: tiny sequence-sensitive pragmatics specialists (binary action gate)
=============================================================================

A narrow binary specialist: executable request vs semantically similar
non-executable utterance, generalizing across surface-generator families.

Commands
    grouped   — grouped semantic CV (existing cv_fold), all arch/sizes, saves
                per-config per-fold model checkpoints + OOF proba per variant
    metrics   — aggregate saved grouped-CV results into the report tables
                (binary metrics, threshold curves, pair ordering, stress)
    lfo       — leave-generator-out. All sizes on capability_question; the
                other present families for the leading config only.
    e5baseline— frozen-e5 logistic + MLP baselines: grouped, cap-Q LOFO, pair
                ordering (stress flagged NA — no re-embed on this box)
    runtime   — params / sizes / latency / tokenization for each candidate

Primary metrics (brief §2): cap-Q LOFO FA, pair ordering acc, pair margin,
grouped P/R, voice stress, fold variance. Aggregate accuracy is secondary.
"""

import argparse
import json
import os
import re
import resource
import subprocess
import sys
import time
import warnings
from collections import Counter

import numpy as np

warnings.filterwarnings("ignore")

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from slice19_bpe import BpeVocab, CharVocab, normalize_match_text, strip_punct

EMBEDDING_PATH = "/tmp/mvn-experiment/embeddings.json"
RESULTS_DIR = "/tmp/mvn-s19"
SEED = 42
MAX_CHAR = 64
MAX_BPE = 24
BPE_VOCAB = 2048

PRESENT_FAMILIES = [
    "polite_request",
    "modal_request",
    "first_person_request",
    "reordered_target",
    "capability_question",
    "question",
]
ABSENT_FAMILIES = ["negation", "reported_speech", "quotation", "hypothetical"]

FAMILY_PRIORITY = [
    "capability_question", "question", "first_person_request",
    "modal_request", "polite_request", "reordered_target", "direct_imperative",
]

DEVICES = [
    "свет", "люстру", "люстра", "жалюзи", "вытяжку", "вытяжка",
    "вентилятор", "кондиционер", "телевизор", "лампу", "лампа",
    "музыку", "музыка", "плеер", "колонку", "колонки", "чайник",
    "бойлер", "обогреватель", "пылесос", "пылесосом", "пол",
    "поливалки", "полив", "арка", "шторы", "штору",
    "динамики", "дверь", "двери", "замок", "гараж", "ворота",
    "кофе", "пасту", "зубы", "крючки", "лаймо", "куртку",
    "будильник", "таймер", "напоминание", "расписание",
]

ARCH_CONFIGS = {
    "char_cnn": ["tiny", "medium"],
    "bigru": ["tiny", "medium", "large"],
    "tiny_transformer": ["small", "medium"],
}

TRAIN_HYPER = {
    "char_cnn":          dict(epochs=20, lr=1e-3, bs=64, clip=None),
    "bigru":             dict(epochs=15, lr=5e-4, bs=64, clip=1.0),
    "tiny_transformer":  dict(epochs=25, lr=5e-4, bs=64, clip=1.0),
}

VARIANTS = ["orig", "nofinal", "strip"]


# ─── data ───────────────────────────────────────────────────────────────────

def load_dev():
    with open(EMBEDDING_PATH) as f:
        data = json.load(f)
    rows = [e for e in data["examples"] if e["dev_pool"]]
    out = []
    for e in rows:
        n_text = normalize_match_text(e["text"])
        out.append({
            "text_orig": n_text,
            "text_nofinal": re.sub(r"[?.!,;:]+$", "", n_text),
            "text_strip": strip_punct(n_text),
            "route": e["route"],
            "y": 1 if e["route"] == "action" else 0,
            "cv_fold": e["cv_fold"],
            "split_group": e["split_group"],
            "tags": set(e.get("tags", [])),
            "source_id": e["source_id"],
            "emb": np.asarray(e["embedding"], dtype=np.float32),
        })
    return out


def family_of(tags):
    for f in FAMILY_PRIORITY:
        if f in tags:
            return f
    return None


def build_pairs(rows, texts):
    """Capability-question vs action-by-shared-object pairs (eval only)."""
    def object_nouns(t):
        found = set()
        tl = t.lower()
        for d in DEVICES:
            if re.search(r"\b" + re.escape(d), tl):
                found.add(d)
        return found

    pairs = []
    for cidx, r in enumerate(rows):
        if "capability_question" not in r["tags"]:
            continue
        c_obj = object_nouns(texts[cidx])
        if not c_obj:
            continue
        dom = re.search(r"kq-cap-([^-]+)", r["source_id"])
        dom = dom.group(1) if dom else None
        for aidx, ra in enumerate(rows):
            if ra["y"] != 1:
                continue
            a_obj = object_nouns(texts[aidx])
            if not (c_obj & a_obj):
                continue
            a_dom = re.search(r"^([a-z]+)-", ra["source_id"])
            a_dom = a_dom.group(1) if a_dom else None
            if dom and a_dom and dom != a_dom:
                continue
            pairs.append((cidx, aidx))
    return pairs


def pair_metrics(pairs, proba):
    if not pairs:
        return {"pairs": 0}
    margins = []
    ties = 0
    order = 0
    rev = 0
    for cidx, aidx in pairs:
        pc, pa = proba[cidx], proba[aidx]
        margins.append(pa - pc)
        if pa > pc:
            order += 1
        elif pa == pc:
            ties += 1
        else:
            rev += 1
    m = np.array(margins)
    return {
        "pairs": len(pairs),
        "ordering_acc": order / len(pairs),
        "mean_margin": float(m.mean()),
        "median_margin": float(np.median(m)),
        "ties": ties,
        "reversed": rev,
    }


# ─── tokenizers ─────────────────────────────────────────────────────────────

def build_tokenizers(rows):
    nat_texts = [r["text_orig"] for r in rows]
    char_vocab = CharVocab(nat_texts)
    bpe = BpeVocab(nat_texts, vocab_size=BPE_VOCAB)
    return char_vocab, bpe


def encode_all(rows, tokenizer, kind):
    """Return dict variant -> (N, max_len) int64 array."""
    max_len = MAX_CHAR if kind == "char" else MAX_BPE
    out = {}
    for v in VARIANTS:
        arr = np.zeros((len(rows), max_len), dtype=np.int64)
        for i, r in enumerate(rows):
            if kind == "char":
                ids = tokenizer.encode(r[f"text_{v}"], max_len)
                arr[i, :len(ids)] = ids
            else:
                ids = tokenizer.encode(r[f"text_{v}"])[:max_len]
                arr[i, :len(ids)] = ids
        out[v] = arr
    return out


# ─── training ───────────────────────────────────────────────────────────────

def build_model(arch, size, vocab_size, max_len):
    import torch
    from slice19_models import CharCNN, BiGRU, TinyTransformer
    if arch == "char_cnn":
        c = (dict(embed_dim=32, filters=64, widths=[3, 4, 5]) if size == "tiny"
             else dict(embed_dim=64, filters=160, widths=[2, 3, 4, 5]))
        return CharCNN(vocab_size, c["embed_dim"], c["filters"], c["widths"])
    if arch == "bigru":
        c = (dict(embed_dim=64, hidden=64) if size == "tiny" else
             (dict(embed_dim=128, hidden=128) if size == "medium" else
              dict(embed_dim=256, hidden=256)))
        return BiGRU(vocab_size, c["embed_dim"], c["hidden"])
    c = (dict(d_model=128, n_layers=2, n_heads=4) if size == "small" else
         dict(d_model=192, n_layers=4, n_heads=4))
    return TinyTransformer(vocab_size, c["d_model"], c["n_layers"], c["n_heads"],
                           max_len=max_len)


def train_binary(X, y, arch, size, vocab_size, seed_offset=0, log=False):
    import torch
    torch.manual_seed(SEED + seed_offset)
    np.random.seed(SEED + seed_offset)
    Xt = torch.from_numpy(X)
    yt = torch.from_numpy(y.astype(np.float32))
    model = build_model(arch, size, vocab_size, X.shape[1])
    h = TRAIN_HYPER[arch]
    opt = torch.optim.AdamW(model.parameters(), lr=h["lr"], weight_decay=1e-4)
    lossf = torch.nn.BCEWithLogitsLoss()
    n = X.shape[0]
    model.train()
    t0 = time.time()
    for epoch in range(h["epochs"]):
        perm = torch.randperm(n)
        running = 0.0
        n_b = 0
        for start in range(0, n, h["bs"]):
            idx = perm[start:start + h["bs"]]
            xb = Xt[idx]
            if xb.dim() == 1:
                xb = xb.unsqueeze(0)
            logits = model(xb)
            loss = lossf(logits, yt[idx])
            opt.zero_grad()
            loss.backward()
            if h["clip"]:
                torch.nn.utils.clip_grad_norm_(model.parameters(), h["clip"])
            opt.step()
            running += float(loss)
            n_b += 1
        if log and (epoch + 1) % 5 == 0:
            print(f"      epoch {epoch+1}/{h['epochs']} loss {running/max(n_b,1):.4f}")
    return model, time.time() - t0


def predict_proba(model, X, bs=256):
    import torch
    model.eval()
    out = []
    with torch.no_grad():
        Xt = torch.from_numpy(X)
        for start in range(0, X.shape[0], bs):
            xb = Xt[start:start + bs]
            if xb.dim() == 1:
                xb = xb.unsqueeze(0)
            logits = model(xb)
            out.append(torch.sigmoid(logits).numpy())
    return np.concatenate(out)


# ─── metrics helpers ────────────────────────────────────────────────────────

def binary_metrics(yt, proba, thr=0.5):
    yp = (proba >= thr).astype(int)
    tp = int(((yt == 1) & (yp == 1)).sum())
    fp = int(((yt == 0) & (yp == 1)).sum())
    fn = int(((yt == 1) & (yp == 0)).sum())
    from sklearn.metrics import roc_auc_score, average_precision_score
    roc = roc_auc_score(yt, proba) if len(np.unique(yt)) > 1 else 0.0
    pr = average_precision_score(yt, proba) if len(np.unique(yt)) > 1 else 0.0
    return {
        "n": int(len(yt)), "tp": tp, "fp": fp, "fn": fn,
        "P": tp / max(tp + fp, 1), "R": tp / max(tp + fn, 1),
        "FA": fp, "FA_rate": fp / max(len(yt), 1),
        "ROC_AUC": float(roc), "PR_AUC": float(pr),
    }


def threshold_curve(yt, proba, thr_grid):
    rows = []
    for thr in thr_grid:
        yp = (proba >= thr).astype(int)
        tp = int(((yt == 1) & (yp == 1)).sum())
        fp = int(((yt == 0) & (yp == 1)).sum())
        fn = int(((yt == 1) & (yp == 0)).sum())
        P = tp / max(tp + fp, 1)
        R = tp / max(tp + fn, 1)
        rows.append({
            "thr": round(float(thr), 4), "P": round(P, 4), "R": round(R, 4),
            "FA": fp, "FA_rate": round(fp / max(len(yt), 1), 4),
        })
    return rows


def operating_points(rows, pair_proba, pairs):
    """Report P>=0.95/0.98/0.99 points with pair separation at threshold."""
    res = {}
    for target in (0.95, 0.98, 0.99):
        pts = [r for r in rows if r["P"] >= target and r["R"] > 0.0]
        if not pts:
            res[str(target)] = None
            continue
        best = max(pts, key=lambda r: r["R"])
        # pair separation at that operating point
        sep = pair_sep_at(best["thr"], pair_proba, pairs)
        best = dict(best); best["pair_sep"] = round(sep, 4)
        res[str(target)] = best
    return res


def pair_sep_at(thr, pair_proba, pairs):
    """Fraction of pairs where action>=thr and cap<thr."""
    if not pairs:
        return 0.0
    ok = 0
    for cidx, aidx in pairs:
        if pair_proba[aidx] >= thr and pair_proba[cidx] < thr:
            ok += 1
    return ok / len(pairs)


def fold_variance(fold_rows):
    return {
        "folds": [
            {
                "fold": fr["fold"],
                "ROC_AUC": fr["ROC_AUC"], "PR_AUC": fr["PR_AUC"],
                "P": fr["P"], "R": fr["R"], "FA": fr["FA"], "n": fr["n"],
            }
            for fr in fold_rows
        ]
    }


# ─── subcommands ────────────────────────────────────────────────────────────

def _result_path():
    os.makedirs(RESULTS_DIR, exist_ok=True)
    return RESULTS_DIR


def cmd_grouped(args):
    rows = load_dev()
    char_vocab, bpe = build_tokenizers(rows)
    y = np.array([r["y"] for r in rows])
    folds = np.array([r["cv_fold"] for r in rows])
    X_char = encode_all(rows, char_vocab, "char")
    X_bpe = encode_all(rows, bpe, "bpe")
    tokenizers = {"char": char_vocab, "bpe": bpe}

    os.makedirs(RESULTS_DIR, exist_ok=True)
    # save tokenizer metadata for reproducibility
    meta = {
        "char_vocab": char_vocab.size,
        "char_vocab_sample": char_vocab.id_to_char[:50],
        "bpe_vocab": bpe.size,
        "bpe_serialized_bytes": bpe.serialized_bytes(),
        "max_char": MAX_CHAR, "max_bpe": MAX_BPE,
        "n": len(rows),
    }
    with open(os.path.join(RESULTS_DIR, "corpus_meta.json"), "w") as f:
        json.dump(meta, f)

    kind_of = {"char_cnn": "char", "bigru": "bpe", "tiny_transformer": "bpe"}
    for arch, sizes in ARCH_CONFIGS.items():
        kind = kind_of[arch]
        X = X_char if kind == "char" else X_bpe
        for size in sizes:
            name = f"{arch}_{size}"
            os.makedirs(os.path.join(RESULTS_DIR, "models"), exist_ok=True)
            probs = {v: np.zeros(len(rows)) for v in VARIANTS}
            fold_rows = []
            for fold in range(5):
                tr = folds != fold
                te = folds == fold
                Xtr = X["strip"][tr]
                model, secs = train_binary(Xtr, y[tr], arch, size,
                                           vocab_size=tokenizers[kind].size,
                                           seed_offset=fold, log=args.verbose)
                torch_models_dir = os.path.join(RESULTS_DIR, "models")
                import torch
                torch.save(model.state_dict(), os.path.join(torch_models_dir, f"{name}_fold{fold}.pt"))
                for v in VARIANTS:
                    probs[v][te] = predict_proba(model, X[v][te])
                fold_m = binary_metrics(y[te], probs["strip"][te])
                fold_m["fold"] = fold
                fold_rows.append(fold_m)
                print(f"  {name} fold {fold}: ROC={fold_m['ROC_AUC']:.3f} "
                      f"PR={fold_m['PR_AUC']:.3f} P={fold_m['P']:.3f} R={fold_m['R']:.3f} "
                      f"FA={fold_m['FA']} n={fold_m['n']} ({secs:.1f}s)")
            np.savez(os.path.join(RESULTS_DIR, f"{name}_probs.npz"),
                     var_orig=probs["orig"], var_nofinal=probs["nofinal"],
                     var_strip=probs["strip"])
            summary = binary_metrics(y, probs["strip"])
            print(f"  {name} OOF: ROC={summary['ROC_AUC']:.3f} PR={summary['PR_AUC']:.3f} "
                  f"P={summary['P']:.3f} R={summary['R']:.3f} FA={summary['FA']}")
    print("grouped done")


def cmd_metrics(args):
    rows = load_dev()
    y = np.array([r["y"] for r in rows])
    tags = [r["tags"] for r in rows]
    pairs = build_pairs(rows, [r["text_strip"] for r in rows])
    print(f"pairs={len(pairs)}")

    out = {}
    for arch, sizes in ARCH_CONFIGS.items():
        for size in sizes:
            name = f"{arch}_{size}"
            fp = os.path.join(RESULTS_DIR, f"{name}_probs.npz")
            if not os.path.exists(fp):
                continue
            z = np.load(fp)
            entry = {"name": name, "arch": arch, "size": size}
            # OOF binary on strip variant (primary training input)
            entry["strip"] = binary_metrics(y, z["var_strip"])
            entry["pairs"] = {}
            entry["pairs"]["strip"] = pair_metrics(pairs, z["var_strip"])
            entry["pairs"]["orig"] = pair_metrics(pairs, z["var_orig"])
            entry["pairs"]["nofinal"] = pair_metrics(pairs, z["var_nofinal"])
            # stress: same OOF models, per-variant metrics
            entry["stress"] = {}
            for v in VARIANTS:
                p = z[f"var_{v}"]
                entry["stress"][v] = {
                    "all_FA": binary_metrics(y, p)["FA_rate"],
                    "capQ_FA": capq_fa(tags, y, p),
                    "modal_recall": modal_recall(tags, y, p),
                }
            # threshold curve + operating points on strip
            entry["curve"] = threshold_curve(y, z["var_strip"],
                                             np.arange(0.30, 1.0, 0.02))
            entry["ops"] = operating_points(entry["curve"], z["var_strip"], pairs)
            out[name] = entry

    with open(os.path.join(RESULTS_DIR, "metrics.json"), "w") as f:
        json.dump(out, f, indent=2, default=str)
    print(json.dumps(out, indent=2, default=str))


def capq_fa(tags, y, proba):
    mask = np.array(["capability_question" in t for t in tags])
    if mask.sum() == 0:
        return 0.0
    sub = proba[mask]
    return float((sub >= 0.5).sum() / mask.sum())


def modal_recall(tags, y, proba):
    mask = np.array([
        (y[i] == 1 and ("polite_request" in tags[i] or "modal_request" in tags[i]))
        for i in range(len(y))
    ])
    if mask.sum() == 0:
        return 0.0
    sub = proba[mask]
    return float((sub >= 0.5).sum() / mask.sum())


def cmd_lfo(args):
    rows = load_dev()
    char_vocab, bpe = build_tokenizers(rows)
    y = np.array([r["y"] for r in rows])
    tags = [r["tags"] for r in rows]
    X_char = encode_all(rows, char_vocab, "char")["strip"]
    X_bpe = encode_all(rows, bpe, "bpe")["strip"]
    kind_of = {"char_cnn": "char", "bigru": "bpe", "tiny_transformer": "bpe"}

    results = {}
    tokenizers = {"char": char_vocab, "bpe": bpe}
    # capability_question LOFO for every config/size (the critical case)
    for arch, sizes in ARCH_CONFIGS.items():
        X = X_char if kind_of[arch] == "char" else X_bpe
        src_idx = np.array(["capability_question" not in t for t in tags])
        tgt_idx = np.array(["capability_question" in t for t in tags])
        for size in sizes:
            model, _ = train_binary(X[src_idx], y[src_idx], arch, size,
                                    vocab_size=tokenizers[kind_of[arch]].size,
                                    seed_offset=17)
            p = predict_proba(model, X[tgt_idx])
            m = binary_metrics(y[tgt_idx], p)
            results[f"{arch}_{size}:capability_question"] = m
            print(f"LFO ability {arch}_{size}: cap rows={m['n']} "
                  f"pos={int(y[tgt_idx].sum())} FA={m['FA']} FA_rate={m['FA_rate']:.3f} "
                  f"P={m['P']:.3f} R={m['R']:.3f} acc={1-m['FA_rate']:.3f}")

    # full family LOFO for the leading config per architecture
    leading = {"char_cnn": "char_cnn_medium", "bigru": "bigru_tiny",
               "tiny_transformer": "tiny_transformer_small"}
    for arch, name in leading.items():
        X = X_char if kind_of[arch] == "char" else X_bpe
        for fam in PRESENT_FAMILIES:
            src = np.array([fam not in t for t in tags])
            tgt = np.array([fam in t for t in tags])
            model, _ = train_binary(X[src], y[src], arch, name.split("_")[-1],
                                    vocab_size=tokenizers[kind_of[arch]].size,
                                    seed_offset=41)
            p = predict_proba(model, X[tgt])
            m = binary_metrics(y[tgt], p)
            results[f"{name}:{fam}"] = m
            print(f"LFO {fam}: {name} rows={m['n']} pos={int(y[tgt].sum())} "
                  f"P={m['P']:.3f} R={m['R']:.3f} FA={m['FA']} acc={1-m['FA_rate']:.3f}")

    with open(os.path.join(RESULTS_DIR, "lfo.json"), "w") as f:
        json.dump(results, f, indent=2, default=str)
    print("lfo done")


def cmd_e5baseline(args):
    from sklearn.linear_model import LogisticRegression
    from sklearn.neural_network import MLPClassifier
    rows = load_dev()
    y = np.array([r["y"] for r in rows])
    folds = np.array([r["cv_fold"] for r in rows])
    X = np.vstack([r["emb"] for r in rows])
    tags = [r["tags"] for r in rows]
    pairs = build_pairs(rows, [r["text_strip"] for r in rows])

    out = {}
    for model_name, model, extra in [
        ("e5_linear", LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs", random_state=42), {}),
        ("e5_mlp_h32", MLPClassifier(hidden_layer_sizes=(32,), activation="relu",
                                     solver="adam", alpha=0.01, random_state=42,
                                     max_iter=800, early_stopping=True,
                                     validation_fraction=0.15, n_iter_no_change=10), {}),
    ]:
        proba = np.zeros(len(rows))
        for fold in range(5):
            tr = folds != fold
            te = folds == fold
            m2 = type(model)(**{k: v for k, v in model.get_params().items()})
            m2.fit(X[tr], y[tr])
            proba[te] = m2.predict_proba(X[te])[:, 1]
        entry = {
            "grouped": binary_metrics(y, proba),
            "pairs": pair_metrics(pairs, proba),
        }
        # cap-Q leave-generator-out (train without the family)
        src = np.array(["capability_question" not in t for t in tags])
        tgt = np.array(["capability_question" in t for t in tags])
        m3 = LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs", random_state=42) \
            if model_name == "e5_linear" else \
            MLPClassifier(hidden_layer_sizes=(32,), alpha=0.01, random_state=42, max_iter=800)
        m3.fit(X[src], y[src])
        p3 = m3.predict_proba(X[tgt])[:, 1]
        entry["capq_lofo"] = binary_metrics(y[tgt], p3)
        entry["stress"] = "NA (no re-embed on this box)"
        out[model_name] = entry
        print(f"{model_name}: grouped FA_rate={entry['grouped']['FA_rate']:.4f} "
              f"PR={entry['grouped']['PR_AUC']:.3f} capQ_LOFO_FA_rate={entry['capq_lofo']['FA_rate']:.4f} "
              f"pairs={entry['pairs']['ordering_acc']:.3f}")
    with open(os.path.join(RESULTS_DIR, "e5baseline.json"), "w") as f:
        json.dump(out, f, indent=2, default=str)
    print("e5baseline done")


def cmd_runtime(args):
    import torch
    rows = load_dev()
    char_vocab, bpe = build_tokenizers(rows)
    X_char = encode_all(rows, char_vocab, "char")["strip"]
    X_bpe = encode_all(rows, bpe, "bpe")["strip"]
    kind_of = {"char_cnn": "char", "bigru": "bpe", "tiny_transformer": "bpe"}
    report = {}
    for arch, sizes in ARCH_CONFIGS.items():
        X = X_char if kind_of[arch] == "char" else X_bpe
        for size in sizes:
            name = f"{arch}_{size}"
            model = build_model(arch, size, int(X.max()) + 1, X.shape[1])
            n_params = sum(p.numel() for p in model.parameters())
            fp32 = n_params * 4
            int8 = n_params
            model.eval()
            # warmup + latency (batch-1, eval mode)
            xb = torch.from_numpy(X[:1])
            with torch.no_grad():
                for _ in range(20):
                    model(xb)
                # tokenization latency
                if kind_of[arch] == "char":
                    t0 = time.perf_counter()
                    for r in rows[:1000]:
                        char_vocab.encode(r["text_strip"], MAX_CHAR)
                    tl = (time.perf_counter() - t0) / 1000
                else:
                    t0 = time.perf_counter()
                    for r in rows[:1000]:
                        bpe.encode(r["text_strip"])
                    tl = (time.perf_counter() - t0) / 1000
                lat = []
                for _ in range(300):
                    t0 = time.perf_counter()
                    model(xb)
                    lat.append(time.perf_counter() - t0)
            lat = np.array(lat) * 1e6
            report[name] = {
                "params": n_params, "fp32_bytes": fp32, "int8_bytes": int8,
                "latency_us_mean": float(lat.mean()), "latency_us_p50": float(np.median(lat)),
                "latency_us_p95": float(np.percentile(lat, 95)),
                "throughput_b1": round(1e6 / float(lat.mean()), 1),
                "tok_us": round(tl * 1e6, 1),
                "tokenizer": "char" if kind_of[arch] == "char" else "bpe",
            }
            print(f"{name}: {n_params} params fp32={fp32/1024:.0f}KiB "
                  f"lat={lat.mean():.0f}us tok={tl*1e6:.1f}us")
    with open(os.path.join(RESULTS_DIR, "runtime.json"), "w") as f:
        json.dump(report, f, indent=2)
    print("runtime done")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["grouped", "metrics", "lfo", "e5baseline", "runtime"])
    ap.add_argument("--verbose", action="store_true")
    args = ap.parse_args()
    t0 = time.time()
    globals()[f"cmd_{args.cmd}"](args)
    print(f"elapsed {time.time()-t0:.1f}s")


if __name__ == "__main__":
    main()