#!/usr/bin/env python3
"""
Slice 18: Sparse Lexical Action/Non-Action Gate
================================================

Answer: can Maven reliably distinguish executable requests from semantically
similar non-actions using lexical/local-order features alone?

Representations under test (all frozen-population, no e5):
    A. word 1-2 grams, TF-IDF
    B. character 3-5 grams, TF-IDF (Unicode, no transliteration)
    C. [word ; char] combined TF-IDF

Population reused from slices 15-17: development corpus v2 (dev_pool), router
labels, SplitGroup, cv_fold, tags. The e5 embedding vectors are ignored.
"""

import io
import json
import re
import sys
import time
import unicodedata
import warnings
from collections import Counter, defaultdict

import numpy as np
from sklearn.exceptions import ConvergenceWarning
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    average_precision_score,
    confusion_matrix,
    f1_score,
    precision_recall_fscore_support,
    roc_auc_score,
)
from sklearn.pipeline import make_pipeline
from scipy import sparse

warnings.filterwarnings("ignore", category=ConvergenceWarning)
warnings.filterwarnings("ignore", category=UserWarning)

EMBEDDING_PATH = "/tmp/mvn-experiment/embeddings.json"

# Present generator/contrast families in the v2 dev pool (for leave-family-out)
PRESENT_FAMILIES = [
    "polite_request",
    "modal_request",
    "first_person_request",
    "reordered_target",
    "capability_question",
    "question",
]

FAMILY_ALIASES = {
    "direct_imperative": "direct_imperative",
    "polite_request": "polite_request",
    "modal_request": "modal_request",
    "first_person_request": "first_person_request",
    "reordered_target": "reordered_target",
    "capability_question": "capability_question",
    "question": "question",
}


# ─── NormalizeMatchText (replicated from internal/router/matchtext.go) ─────

def normalize_match_text(s: str) -> str:
    """NFKC → lowercase → collapse Unicode whitespace. Keeps punctuation, ё."""
    out = unicodedata.normalize("NFKC", s).strip().lower()
    out = re.sub(r"\s+", " ", out)
    return out


# ─── Data Loading ───────────────────────────────────────────────────────────

def load_data():
    with open(EMBEDDING_PATH) as f:
        data = json.load(f)
    return data["meta"], data["examples"]


def filter_dev_pool(examples):
    return [e for e in examples if e["dev_pool"]]


def residual_only(examples):
    return [e for e in examples if not e["fast_path_resolved"]]


def fmt_pct(v, d=1):
    return f"{100*v:.{d}f}%"


def ff(v, d=3):
    return f"{v:.{d}f}"


def strip_punct(text: str) -> str:
    """Remove all punctuation (shared with slice 16/17 apply_voice_stress)."""
    t = re.sub(r"[?.!,;:]+$", "", text.strip())
    t = re.sub(r"[^\w\s]", "", t)
    t = t.lower()
    t = re.sub(r"\s+", " ", t).strip()
    return t


# ─── Feature Builders ───────────────────────────────────────────────────────

def build_features(texts, kind):
    """Build a TF-IDF matrix for the given representation kind.
    kind in {'word','char','both'}. Returns (X_sparse, vectorizer)."""
    if kind == "word":
        vec = TfidfVectorizer(
            ngram_range=(1, 2), analyzer="word",
            strip_accents=None, lowercase=False,
            min_df=2, sublinear_tf=True,
        )
    elif kind == "char":
        # preserve case (already lowered) and identity of missing diacritics;
        # token_pattern null => char analyzer
        vec = TfidfVectorizer(
            ngram_range=(3, 5), analyzer="char",
            strip_accents=None, lowercase=False,
            min_df=2, sublinear_tf=True,
        )
    elif kind == "both":
        vec_word = TfidfVectorizer(
            ngram_range=(1, 2), analyzer="word",
            strip_accents=None, lowercase=False, min_df=2, sublinear_tf=True,
        )
        vec_char = TfidfVectorizer(
            ngram_range=(3, 5), analyzer="char",
            strip_accents=None, lowercase=False, min_df=2, sublinear_tf=True,
        )
        Xw = vec_word.fit_transform(texts)
        Xc = vec_char.fit_transform(texts)
        X = sparse.hstack([Xw, Xc]).tocsr()
        return X, ("both", vec_word, vec_char)
    X = vec.fit_transform(texts)
    return X, vec


def vocab_size(vectorizer):
    if isinstance(vectorizer, tuple):
        _, vw, vc = vectorizer
        return vw.get_feature_names_out().shape[0] + vc.get_feature_names_out().shape[0]
    return vectorizer.get_feature_names_out().shape[0]


def transform_texts(texts, vectorizer):
    """Apply an already-fitted vectorizer (handles the 2-tuple 'both' case)."""
    if isinstance(vectorizer, tuple):
        _, vw, vc = vectorizer
        Xw = vw.transform(texts)
        Xc = vc.transform(texts)
        return sparse.hstack([Xw, Xc]).tocsr()
    return vectorizer.transform(texts)


# ─── Grouped CV ─────────────────────────────────────────────────────────────

def run_binary_grouped_cv(X, y, fold_ids, C=1.0):
    """Grouped CV for binary action vs not_action. Returns OOF rows + fold metrics."""
    yb = np.array([1 if t == "action" else 0 for t in y])
    fold_ids = np.asarray(fold_ids)
    unique_folds = sorted(set(fold_ids.tolist()))
    oof_rows = []
    fold_metrics = []

    for test_fold in unique_folds:
        tr = fold_ids != test_fold
        te = fold_ids == test_fold
        clf = LogisticRegression(C=C, max_iter=2000, solver="lbfgs", random_state=42)
        clf.fit(X[tr], yb[tr])
        proba = clf.predict_proba(X[te])[:, 1]
        pred = (proba >= 0.5).astype(int)
        yt = yb[te]

        fp = int(((yt == 0) & (pred == 1)).sum())
        fn = int(((yt == 1) & (pred == 0)).sum())
        tp = int(((yt == 1) & (pred == 1)).sum())
        tn = int(((yt == 0) & (pred == 0)).sum())

        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
        acc = accuracy_score(yt, pred)
        prec = tp / max(tp + fp, 1)
        rec = tp / max(tp + fn, 1)

        fold_metrics.append({
            "fold": int(test_fold), "n": int(len(yt)),
            "roc_auc": roc, "pr_auc": pr,
            "action_precision": prec, "action_recall": rec,
            "fp": fp, "fn": fn, "tp": tp, "tn": tn,
            "acc": acc,
        })
        te_idx = np.where(te)[0]
        for i in range(len(yt)):
            oof_rows.append({
                "fold": int(test_fold),
                "proba": float(proba[i]),
                "pred": int(pred[i]),
                "true": int(yt[i]),
            })

    return oof_rows, fold_metrics


# ─── Metrics from OOF ───────────────────────────────────────────────────────

def binary_metrics_from_oof(oof_rows):
    yt = np.array([r["true"] for r in oof_rows])
    yp = np.array([r["pred"] for r in oof_rows])
    proba = np.array([r["proba"] for r in oof_rows])
    n = len(yt)
    tp = int(((yt == 1) & (yp == 1)).sum())
    fp = int(((yt == 0) & (yp == 1)).sum())
    fn = int(((yt == 1) & (yp == 0)).sum())
    prec = tp / max(tp + fp, 1)
    rec = tp / max(tp + fn, 1)
    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": n, "tp": tp, "fp": fp, "fn": fn,
        "action_precision": prec, "action_recall": rec,
        "fa_rate": fp / max(n, 1),
        "roc_auc": roc, "pr_auc": pr,
    }


def threshold_curve(oof_rows, thresholds):
    yt = np.array([r["true"] for r in oof_rows])
    proba = np.array([r["proba"] for r in oof_rows])
    rows = []
    for thr in thresholds:
        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())
        rows.append({
            "threshold": round(float(thr), 4),
            "action_precision": round(tp / max(tp + fp, 1), 4),
            "action_recall": round(tp / max(tp + fn, 1), 4),
            "fa_count": fp,
            "fa_rate": round(fp / max(len(yt), 1), 4),
        })
    return rows


# ─── Leave-Family-Out ───────────────────────────────────────────────────────

def run_leave_family_out(texts, y, family_per_row, family, kind):
    """Train without `family`, evaluate on `family` only."""
    mask_members = family_per_row == family
    if mask_members.sum() == 0:
        return None
    only_family = (mask_members).astype(bool)
    train_idx = np.where(~only_family)[0]
    test_idx = np.where(only_family)[0]
    X, _ = build_features([texts[i] for i in train_idx], kind)
    # map test rows onto the full vocabulary
    tr_texts = [texts[i] for i in train_idx]
    te_texts = [texts[i] for i in test_idx]
    all_texts = tr_texts + te_texts
    Xall, _ = build_features(all_texts, kind)
    Xtr = Xall[:len(tr_texts)]
    Xte = Xall[len(tr_texts):]
    ytr = np.array([1 if y[i] == "action" else 0 for i in train_idx])
    yte = np.array([1 if y[i] == "action" else 0 for i in test_idx])
    clf = LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs", random_state=42)
    clf.fit(Xtr, ytr)
    pred = clf.predict(Xte)
    tp = int(((yte == 1) & (pred == 1)).sum())
    fp = int(((yte == 0) & (pred == 1)).sum())
    fn = int(((yte == 1) & (pred == 0)).sum())
    prec = tp / max(tp + fp, 1)
    rec = tp / max(tp + fn, 1)
    acc = accuracy_score(yte, pred)
    return {
        "family": family, "rows": int(len(yte)),
        "action_precision": prec, "action_recall": rec,
        "fa_count": fp, "acc": acc,
        "yticks": f"pos={int(yte.sum())} neg={int((yte==0).sum())}",
    }


# ─── Paired Action/Capability Test ──────────────────────────────────────────

def paired_action_capability(texts, y, tags_per_row_by_idx, kind, C=1.0):
    """For the best sparse representation, evaluate pairwise ordering."""
    idx = list(range(len(texts)))
    emb = build_features(texts, kind)[0]
    clf = LogisticRegression(C=C, max_iter=2000, solver="lbfgs", random_state=42)
    yb = np.array([1 if t == "action" else 0 for t in y])
    # use OOF-style: fit on full then derive? We report the pairwise score test;
    # to avoid leakage we use grouped OOF proba via grouped cv.
    return None


# ─── Main ──────────────────────────────────────────────────────────────────

def main():
    meta, examples = load_data()
    dev = filter_dev_pool(examples)
    dev_res = residual_only(dev)

    # normalized texts
    for e in dev:
        e["n_text"] = normalize_match_text(e["text"])

    texts = [e["n_text"] for e in dev]
    y = [e["route"] for e in dev]
    folds = [e["cv_fold"] for e in dev]
    tags = [set(e.get("tags", [])) for e in dev]
    split_groups = [e["split_group"] for e in dev]
    source_ids = [e["source_id"] for e in dev]

    n_action = sum(1 for r in y if r == "action")
    n_not = len(y) - n_action
    print(f"Dev pool: {len(dev)}  action={n_action}  not_action={n_not}")
    print()

    results = {}

    # ── 1. Build features and run grouped CV for each representation ──────
    for kind in ["word", "char", "both"]:
        print(f"\n=== {kind} TF-IDF ===")
        t0 = time.time()
        X, vec = build_features(texts, kind)
        build_t = time.time() - t0
        vs = vocab_size(vec)
        print(f"  vocab={vs}  X.shape={X.shape}  build={build_t:.2f}s")

        oof, folds_m = run_binary_grouped_cv(X, y, folds, C=1.0)
        m = binary_metrics_from_oof(oof)
        results[kind] = {
            "vs": vs, "build_t": build_t, "oof": oof, "fold_metrics": folds_m,
            "metrics": m, "X": X, "vec": vec,
            "texts": texts, "y": y, "folds": folds,
        }
        print(f"  ROC={ff(m['roc_auc'])} PR={ff(m['pr_auc'])} P={ff(m['action_precision'])} "
              f"R={ff(m['action_recall'])} FA={m['fp']} ({fmt_pct(m['fa_rate'])})")

        # fold-level
        for fm in folds_m:
            print(f"    fold {fm['fold']}: ROC={ff(fm['roc_auc'])} PR={ff(fm['pr_auc'])} "
                  f"P={ff(fm['action_precision'])} R={ff(fm['action_recall'])} "
                  f"FP={fm['fp']} FN={fm['fn']} n={fm['n']}")

    # pick best by PR-AUC
    best_kind = max(["word", "char", "both"], key=lambda k: results[k]["metrics"]["pr_auc"])
    print(f"\nBest representation by PR-AUC: {best_kind}")

    # ── 2. Leave-family-out for best kind ─────────────────────────────────
    print(f"\n=== Leave-family-out ({best_kind}) ===")
    fam_per_row = []
    for tg in tags:
        fam = None
        # prefer the more specific contrast/request families first (a single
        # utterance may carry several generator tags, e.g. capability_question
        # plus direct_imperative). Check the informative ones before the
        # generic direct_imperative fallback.
        for f in ["capability_question", "question", "first_person_request",
                  "modal_request", "polite_request", "reordered_target",
                  "direct_imperative"]:
            if f in tg:
                fam = f
                break
        fam_per_row.append(fam)
    fam_per_row = np.array(fam_per_row, dtype=object)
    lfo = {}
    for fam in PRESENT_FAMILIES:
        r = run_leave_family_out(texts, y, fam_per_row, fam, best_kind)
        if r is None:
            print(f"  {fam}: (no rows)")
            continue
        lfo[fam] = r
        print(f"  {fam}: rows={r['rows']} ({r['yticks']}) P={ff(r['action_precision'])} "
              f"R={ff(r['action_recall'])} FA={r['fa_count']} acc={fmt_pct(r['acc'])}")

    # ── 3. Action/capability paired test ─────────────────────────────────
    print(f"\n=== Paired action/capability test ({best_kind}) ===")
    # Use grouped-CV OOF proba for ordering (no leakage)
    oof = results[best_kind]["oof"]
    # map oof rows back by source_id order
    # oof rows are appended per fold in dev order; reconstruct
    # We'll re-embed and get proba via grouped CV with proba recorded per row.
    # Re-run grouped CV capturing per-row proba aligned to dev indices.
    X = results[best_kind]["X"]
    yb = np.array([1 if r == "action" else 0 for r in y])
    folds_arr = np.array(folds)
    dev_proba = np.zeros(len(dev))
    for te_fold in sorted(set(folds)):
        tr = folds_arr != te_fold
        te = folds_arr == te_fold
        clf = LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs", random_state=42)
        clf.fit(X[tr], yb[tr])
        dev_proba[te] = clf.predict_proba(X[te])[:, 1]

    # group action seeds: pair each capability-question row with action rows
    # sharing the same target object noun and the same act (turn_on/turn_off/…).
    # Both were generated over a common entity+event inventory, so the object
    # lexeme is the semantic link between a knowledge question and its
    # executable sibling.
    DEVICES = [
        "свет", "люстру", "люстра", "жалюзи", "вытяжку", "вытяжка",
        "вентилятор", "кондиционер", "телевизор", "лампу", "лампа",
        "музыку", "музыка", "плеер", "колонку", "колонки", "чайник",
        "бойлер", "обогреватель", "пылесос", "пылесосом", "пол",
        "поливалки", "полив", "арка", "шторы", "штору", "штору",
        "динамики", "дверь", "двери", "замок", "гараж", "ворота",
        "кофе", "пасту", "зубы", "крючки", "лаймо", "куртку",
        "будильник", "таймер", "напоминание", "расписание",
    ]

    def object_nouns(t):
        found = set()
        tl = t.lower()
        for d in DEVICES:
            # match as standalone word (handle Russian case endings loosely via prefix)
            if re.search(r"\b" + re.escape(d), tl):
                found.add(d)
        return found

    # verb/act family per row from source_id (e.g. ha-light-off -> off)
    def act_family(src):
        # pull the 3rd token-ish: ha-light-off => 'off'; ha-light-on => 'on'
        m = re.search(r"^(\w+)-([a-z_]+)-([a-zA-Z_]+)", src)
        if m:
            return f"{m.group(1)}-{m.group(2)}-{m.group(3)}"
        # generic fallback
        return src.split("-")[0]

    cap_rows = [i for i in range(len(dev)) if "capability_question" in tags[i]]
    act_idxs = [i for i in range(len(dev)) if y[i] == "action"]

    # For each capability row, candidate sibling actions: same object noun
    # AND same act (turn_on vs turn_off), i.e. same domain+object. We accept
    # any action row sharing an object and matching the on/off sense if present.
    pairs = []
    for cidx in cap_rows:
        c_obj = object_nouns(texts[cidx])
        if not c_obj:
            continue
        # cap rows are kq-cap-<domain>-<n>; domain token after 'kq-cap-'
        dom = re.search(r"kq-cap-([^-]+)", source_ids[cidx])
        dom = dom.group(1) if dom else None
        for aidx in act_idxs:
            a_obj = object_nouns(texts[aidx])
            if not (c_obj & a_obj):
                continue
            # require same home domain when both carry one
            a_dom = re.search(r"^([a-z]+)-", source_ids[aidx])
            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))

    # keep it bounded: cap row pairs with many actions (one per room); that's fine
    order_ok = 0
    margins = []
    reversed_pairs = []
    for cidx, aidx in pairs:
        pc = dev_proba[cidx]
        pa = dev_proba[aidx]
        margins.append(pa - pc)
        if pa > pc:
            order_ok += 1
        else:
            reversed_pairs.append((texts[cidx][:40], pc, texts[aidx][:40], pa))
    if pairs:
        order_acc = order_ok / len(pairs)
        margins_arr = np.array(margins)
        print(f"  pairs={len(pairs)}  order_acc={ff(order_acc)}  mean_margin={ff(margins_arr.mean())}  "
              f"median_margin={ff(np.median(margins_arr))}  reversed={len(reversed_pairs)}")
        for rev in reversed_pairs[:12]:
            print(f"    REV: cap '{rev[0]}' P={rev[1]:.3f} < act '{rev[2]}' P={rev[3]:.3f}")
    else:
        order_acc = None
        print("  (no matched pairs)")

    paired = {
        "pairs": len(pairs), "order_acc": order_acc,
        "mean_margin": float(np.mean(margins)) if margins else None,
        "median_margin": float(np.median(margins)) if margins else None,
        "reversed": len(reversed_pairs),
    }

    # ── 4. Voice stress scoring ───────────────────────────────────────────
    print(f"\n=== Voice-like stress ({best_kind}) ===")
    # Build stress variants, embed, score with a model trained on normal text
    # Train one model on full dev (normal punctuation); score stressed variants.
    clf_full = LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs", random_state=42)
    clf_full.fit(X, yb)

    stress_texts = []
    stress_meta = []  # (orig_index, mode)
    for i in range(len(dev)):
        orig = texts[i]
        s_nofinal = re.sub(r"[?.!,;:]+$", "", orig)
        s_all = strip_punct(dev[i]["text"])
        if s_nofinal != orig:
            stress_texts.append(s_nofinal); stress_meta.append((i, "no_final"))
        if s_all != orig:
            stress_texts.append(s_all); stress_meta.append((i, "all"))
    if stress_texts:
        # reuse the already-fitted best-kind vectorizer to keep the feature space
        # identical to what the full model was trained on.
        X_stress = transform_texts(stress_texts, results[best_kind]["vec"])
        proba_stress = clf_full.predict_proba(X_stress)[:, 1]
    else:
        proba_stress = []
    # classify ability: capability-question false-action rate and modal-action recall
    stress_by_mode = defaultdict(list)
    for (i, mode), p in zip(stress_meta, proba_stress):
        stress_by_mode[mode].append((i, p, y[i], tags[i]))
    voice = {}
    for mode, rows in stress_by_mode.items():
        cap_q = [p for (i, p, lbl, tg) in rows if "capability_question" in tg]
        cap_fa = sum(1 for p in cap_q if p >= 0.5)
        mod_act = [p for (i, p, lbl, tg) in rows if lbl == "action" and ("polite_request" in tg or "modal_request" in tg)]
        mod_rec = sum(1 for p in mod_act if p >= 0.5) / max(len(mod_act), 1)
        voice[mode] = {
            "n": len(rows),
            "cap_q_fa": cap_fa / max(len(cap_q), 1), "cap_q_n": len(cap_q),
            "modal_action_recall": mod_rec, "modal_n": len(mod_act),
        }
        print(f"  {mode}: n={len(rows)} cap_q_FA={fmt_pct(voice[mode]['cap_q_fa'])} ({voice[mode]['cap_q_n']}) "
              f"modal_recall={ff(voice[mode]['modal_action_recall'])} ({voice[mode]['modal_n']})")

    # ── 5. Punctuation ablation ───────────────────────────────────────────
    print(f"\n=== Punctuation ablation ({best_kind}) ===")
    texts_stripped = [strip_punct(dev[i]["text"]) for i in range(len(dev))]
    X_stripped_train, vec_stripped = build_features(texts_stripped, best_kind)
    clf_stripped = LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs", random_state=42)
    clf_stripped.fit(X_stripped_train, yb)
    # eval on stripped (in-train) and punctuated (out-of-train) via same vectorizer
    X_orig = transform_texts(texts, vec_stripped)
    p_orig = clf_stripped.predict_proba(X_orig)[:, 1]
    p_strip = clf_stripped.predict_proba(X_stripped_train)[:, 1]
    def report_ablation(proba, name):
        yp = (proba >= 0.5).astype(int)
        tp = int(((yb == 1) & (yp == 1)).sum())
        fp = int(((yb == 0) & (yp == 1)).sum())
        fn = int(((yb == 1) & (yp == 0)).sum())
        prec = tp / max(tp + fp, 1); rec = tp / max(tp + fn, 1)
        print(f"  trained-stripped, eval {name}: P={ff(prec)} R={ff(rec)} FA={fp} ({fmt_pct(fp/len(yb))})")
    report_ablation(p_orig, "punctuated")
    report_ablation(p_strip, "stripped")

    # ── 6. Threshold curve for best kind ──────────────────────────────────
    print(f"\n=== Threshold curve ({best_kind}) ===")
    oof = results[best_kind]["oof"]
    thresh = np.arange(0.10, 0.995, 0.015).tolist()
    tcurve = threshold_curve(oof, thresh)
    print(f"{'thr':>6} {'P':>6} {'R':>6} {'FA':>5} {'FArate':>8}")
    for t in tcurve:
        marker = ""
        if t["action_precision"] >= 0.95:
            marker = " ← P>=0.95"
        print(f"{t['threshold']:>6.3f} {t['action_precision']:>6.3f} {t['action_recall']:>6.3f} "
              f"{t['fa_count']:>5} {t['fa_rate']:>8.4f}{marker}")
    p95 = [t for t in tcurve if t["action_precision"] >= 0.95 and t["action_recall"] > 0.01]
    print(f"\nP>=0.95 region: {len(p95)} points; best recall there = "
          f"{max((t['action_recall'] for t in p95), default=0.0):.4f}")

    # threshold curve using the (no_leak) dev proba instead of .5-threshold OOF
    # OOF pred used fixed 0.5; curve above re-derives from proba. Good.

    # ── 7. False-action decomposition ────────────────────────────────────
    print(f"\n=== False-action decomposition ({best_kind}) ===")
    # recompute OOF predictions at 0.5 from stored 'pred'
    fa_by_family = Counter()
    fa_by_group = Counter()
    # need oof aligned to source_ids — oof stored without source id; rebuild
    # Re-do grouped cv capturing source_id + tag + split_group
    fake_oof = []
    for te_fold in sorted(set(folds_arr)):
        tr = folds_arr != te_fold; te = folds_arr == te_fold
        clf = LogisticRegression(C=1.0, max_iter=2000, solver="lbfgs", random_state=42)
        clf.fit(X[tr], yb[tr])
        p = clf.predict_proba(X[te])[:, 1]
        pr = (p >= 0.5).astype(int)
        te_idx = np.where(te)[0]
        for k, i in enumerate(te_idx):
            fake_oof.append({
                "source_id": source_ids[i], "split_group": split_groups[i],
                "tags": tags[i], "true": y[i], "proba": float(p[k]),
                "pred": int(pr[k]), "fold": int(te_fold),
            })

    def classify_semantic_family(row):
        tg = row["tags"]
        if "capability_question" in tg: return "capability_question"
        if "question" in tg: return "ordinary_question"
        if "remember" in tg or "note" in tg or "idea" in tg or "free_form" in tg: return "memory_write"
        if "version" in tg or "system" in tg or "health" in tg or "status" in tg: return "system"
        if "recall" in tg or "world" in tg or "definition" in tg or "aggregate" in tg: return "knowledge_general"
        if "greeting" in tg or "goodbye" in tg or "thanks" in tg: return "conversation"
        if row["true"].startswith("system"): return "system"
        return "other"

    for r in fake_oof:
        if r["pred"] == 1 and r["true"] != "action":
            fam = classify_semantic_family(r)
            fa_by_family[fam] += 1
            fa_by_group[r["split_group"]] += 1

    print("By family:")
    for k, v in fa_by_family.most_common():
        print(f"  {k}: {v}")
    print("By SplitGroup (top 15):")
    for k, v in fa_by_group.most_common(15):
        print(f"  {k}: {v}")

    # ── 8. Six-way probe on best representation ──────────────────────────
    print(f"\n=== Six-way probe ({best_kind}) ===")
    routes = ["action", "conversation", "knowledge", "memory_write", "system", "uncertain"]
    y6 = np.array(y)
    # fit grouped cv 6-way
    oof6 = []
    fold6 = []
    for te_fold in sorted(set(folds_arr)):
        tr = folds_arr != te_fold; te = folds_arr == te_fold
        clf = LogisticRegression(C=1.0, max_iter=3000, solver="lbfgs", random_state=42)
        clf.fit(X[tr], y6[tr])
        pr = clf.predict(X[te])
        p6 = clf.predict_proba(X[te])
        classes = clf.classes_
        te_idx = np.where(te)[0]
        for k, i in enumerate(te_idx):
            oof6.append({
                "true": y6[i], "pred": pr[k], "proba": {c: float(p6[k][j]) for j, c in enumerate(classes)},
                "tags": tags[i], "split_group": split_groups[i], "source_id": source_ids[i],
            })
    yt6 = [r["true"] for r in oof6]; yp6 = [r["pred"] for r in oof6]
    acc6 = accuracy_score(yt6, yp6)
    macro6 = f1_score(yt6, yp6, average="macro", zero_division=0)
    prec6, rec6, f16, sup6 = precision_recall_fscore_support(yt6, yp6, labels=routes, zero_division=0)
    fa6 = sum(1 for t, p in zip(yt6, yp6) if t != "action" and p == "action")
    ap6 = sum(1 for t, p in zip(yt6, yp6) if t == "action" and p == "action") / max(sum(1 for p in yp6 if p == "action"), 1)
    ar6 = sum(1 for t, p in zip(yt6, yp6) if t == "action" and p == "action") / max(sum(1 for t in yt6 if t == "action"), 1)
    print(f"  acc={fmt_pct(acc6)} macroF1={ff(macro6)} actionP={ff(ap6)} actionR={ff(ar6)} FA={fa6} ({fmt_pct(fa6/len(yt6))})")
    for i, r in enumerate(routes):
        print(f"    {r:<12} P={ff(prec6[i])} R={ff(rec6[i])} F1={ff(f16[i])} n={int(sup6[i])}")

    # ── 9. Artifact size / runtime ───────────────────────────────────────
    print(f"\n=== Artifact size / runtime ===")
    for kind in ["word", "char", "both"]:
        rr = results[kind]
        vec = rr["vec"]
        ncoef = rr["metrics"]["n"]
        # non-zero coefficients = vocab (TF-IDF), logistic has 1 weight per vocab
        print(f"  {kind}: vocab={rr['vs']}  fp32 model bytes={rr['vs']*4}  "
              f"build={rr['build_t']:.3f}s")

    # done
    print("\nDone.")


if __name__ == "__main__":
    main()