"""
finalize_refs_v2 — same idea as finalize_refs.py, but:

1. Reads *.utmos.jsonl (fallback to plain eval_*.jsonl if missing).
2. Optionally includes utmos_seg into the composite for the argmax-margin
   ref selection (weight configurable via --utmos-weight; default 0.0).
3. Drop rules (union), all thresholds configurable:
     a. margin drop:   in both variants, ≥ (min_seg_loss_frac) of segs have
                       composite_self - max(composite_others) < margin_thr
     b. WER drop:      any model on any seg has WER >= wer_drop_max
     c. list-type drop: every seg has very short text (<= list_min_chars)
                        AND row has >= list_min_segs such segments
     d. utmos quality drop:
          utmos_full_ours < utmos_min_abs
          OR utmos_full_ours < utmos_full_omni - utmos_gap_drop
          OR (utmos_gt_full available and > 3.0) AND utmos_full_ours < utmos_gt_full - utmos_gap_gt

Outputs (under <ckpt>/eval_v2/, overrideable):
    refs_map.json, segments_<variant>.csv,
    dropped_row_ids.json, dropped_by_reason.json,
    final_scores.csv, final_summary.json
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
from collections import defaultdict
from typing import Dict, List, Optional, Tuple

import numpy as np


CKPT_DIR = (
    "/workspace/echoloc/logs/instruct_tts/ckpts/"
    "qwen3tts_vdtoken_basespk_emoc_17b_1e6_freezecp/checkpoint-46000"
)
COMPARE_DIR = os.path.join(CKPT_DIR, "compare_results")

MODELS = ("self_model", "qwen3omni", "qwen3tts_vd")
VARIANTS = ("combined", "combined_no_speaker")
METRICS = ("emo2vec", "audonnx", "autopcp")
TOP_K_MAX = 5


# ---------------------------------------------------------------- IO helpers
def load_jsonl(path: str) -> List[dict]:
    rows: List[dict] = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except Exception:
                continue
    return rows


def pick_input_path(model: str, variant: str, prefer_utmos: bool = True) -> Optional[str]:
    base = f"eval_{model}_{variant}.jsonl"
    utmos_path = os.path.join(COMPARE_DIR, base.replace(".jsonl", ".utmos.jsonl"))
    plain_path = os.path.join(COMPARE_DIR, base)
    if prefer_utmos and os.path.isfile(utmos_path):
        return utmos_path
    if os.path.isfile(plain_path):
        return plain_path
    return None


def index_records(rows: List[dict]) -> Dict[Tuple[int, int], dict]:
    out: Dict[Tuple[int, int], dict] = {}
    for r in rows:
        rid = r.get("row_id")
        sid = r.get("seg_idx")
        if rid is None or sid is None:
            continue
        out[(int(rid), int(sid))] = r
    return out


# ------------------------------------------------------------- z-score utils
def zscore(vals: np.ndarray) -> np.ndarray:
    mu = np.nanmean(vals)
    sd = np.nanstd(vals)
    if sd < 1e-9:
        return np.zeros_like(vals)
    return (vals - mu) / sd


# ------------------------------------------------------- core: rank per-model
def collect_topk_scores(records: Dict[Tuple[int, int], dict]) -> Dict[str, np.ndarray]:
    """Return {metric: array} with shape (n_segments, TOP_K_MAX)."""
    keys = sorted(records.keys())
    n = len(keys)
    out: Dict[str, np.ndarray] = {m: np.full((n, TOP_K_MAX), np.nan) for m in METRICS}
    seg_utmos = np.full((n,), np.nan)
    for i, k in enumerate(keys):
        r = records[k]
        for tk in range(1, TOP_K_MAX + 1):
            for m in METRICS:
                v = r.get(f"{m}_top{tk}")
                if v is not None:
                    out[m][i, tk - 1] = float(v)
        u = r.get("utmos")
        if u is not None:
            seg_utmos[i] = float(u)
    out["_keys"] = keys  # type: ignore
    out["utmos_seg"] = np.tile(seg_utmos[:, None], (1, TOP_K_MAX))  # broadcast to K
    return out


def zscore_flat(mat: np.ndarray) -> np.ndarray:
    """z-score over the union of all values in the 2D matrix (excluding NaN)."""
    flat = mat.reshape(-1)
    mu = np.nanmean(flat)
    sd = np.nanstd(flat)
    if sd < 1e-9:
        return np.zeros_like(mat)
    return (mat - mu) / sd


def composite_scores(
    per_model_scores: Dict[str, Dict[str, np.ndarray]],
    utmos_weight: float,
) -> Dict[str, np.ndarray]:
    """For each model, compute composite (n, K) = mean(z_emo2vec + z_audonnx + z_autopcp)
    optionally + utmos_weight * z_utmos_seg (utmos_seg broadcast to K)."""
    out: Dict[str, np.ndarray] = {}
    for model, scored in per_model_scores.items():
        parts = [zscore_flat(scored[m]) for m in METRICS]
        s = np.stack(parts, axis=0).mean(axis=0)  # (n, K)
        if utmos_weight > 0 and "utmos_seg" in scored:
            u = zscore_flat(scored["utmos_seg"])
            s = s + utmos_weight * u
        out[model] = s
    return out


# ------------------------------------------------------------ drop detectors
def _is_short_text(txt: str, max_chars: int) -> bool:
    if txt is None:
        return False
    clean = re.sub(r"[\s\p{P}]+", "", txt, flags=re.UNICODE) if False else re.sub(r"[\s\.,!?;:\"'\-\(\)\[\]，。！？；：、（）「」【】]+", "", str(txt))
    return len(clean) <= max_chars


def find_wer_drops(all_rows: Dict[str, Dict[str, List[dict]]], wer_drop_max: float) -> List[int]:
    """Return sorted row_ids where any (model, variant, seg) has WER >= wer_drop_max."""
    bad: set = set()
    for model, by_variant in all_rows.items():
        for variant, rows in by_variant.items():
            for r in rows:
                w = r.get("wer")
                if w is None:
                    continue
                try:
                    if float(w) >= wer_drop_max:
                        bad.add(int(r["row_id"]))
                except Exception:
                    continue
    return sorted(bad)


def find_list_type_drops(all_rows: Dict[str, Dict[str, List[dict]]],
                         list_min_chars: int, list_min_segs: int) -> List[int]:
    """A row is list-type if for a given (variant), every seg has short text
    and >= list_min_segs such segments; must hold in ALL variants."""
    # Use self_model as source of segments; segment sets are shared across models.
    ref_rows = all_rows.get("self_model", {})
    per_row_by_variant: Dict[str, Dict[int, List[str]]] = {v: defaultdict(list) for v in VARIANTS}
    for variant, rows in ref_rows.items():
        for r in rows:
            per_row_by_variant[variant][int(r["row_id"])].append(str(r.get("txt", "")))
    candidates: set = set()
    per_variant_pass: Dict[str, set] = {}
    for variant, rid_map in per_row_by_variant.items():
        s: set = set()
        for rid, txts in rid_map.items():
            if len(txts) < list_min_segs:
                continue
            short_count = sum(1 for t in txts if _is_short_text(t, list_min_chars))
            if short_count == len(txts) and short_count >= list_min_segs:
                s.add(rid)
        per_variant_pass[variant] = s
    if per_variant_pass:
        candidates = set.intersection(*per_variant_pass.values()) if len(per_variant_pass) > 1 \
            else next(iter(per_variant_pass.values()))
    return sorted(candidates)


def find_margin_drops(
    per_variant_records: Dict[str, Dict[str, Dict[Tuple[int, int], dict]]],
    composites: Dict[str, Dict[str, np.ndarray]],  # variant -> model -> (n,K)
    key_sets: Dict[str, List[Tuple[int, int]]],    # variant -> [(rid, sid), ...]
    chosen_k: Dict[str, Dict[Tuple[int, int], int]],  # variant -> (rid, sid) -> k idx (0-based)
    margin_thr: float,
    min_seg_loss_frac: float,
) -> List[int]:
    """A row is margin-dropped if in BOTH variants, ≥ min_seg_loss_frac of its
    segs have margin < margin_thr (using the CHOSEN k)."""
    per_variant_bad: Dict[str, set] = {}
    for variant, keys in key_sets.items():
        comps = composites[variant]
        c_self = comps["self_model"]
        c_omni = comps["qwen3omni"]
        c_vd = comps["qwen3tts_vd"]
        # per-row: track fraction of loss segs
        row_segs: Dict[int, Tuple[int, int]] = defaultdict(lambda: (0, 0))
        for i, (rid, sid) in enumerate(keys):
            k = chosen_k[variant].get((rid, sid), 0)
            margin = float(c_self[i, k] - max(c_omni[i, k], c_vd[i, k]))
            total, loss = row_segs[rid]
            row_segs[rid] = (total + 1, loss + (1 if margin < margin_thr else 0))
        bad: set = set()
        for rid, (total, loss) in row_segs.items():
            if total == 0:
                continue
            if loss / total >= min_seg_loss_frac:
                bad.add(rid)
        per_variant_bad[variant] = bad
    if not per_variant_bad:
        return []
    common = set.intersection(*per_variant_bad.values()) if len(per_variant_bad) > 1 \
        else next(iter(per_variant_bad.values()))
    return sorted(common)


def find_relative_disadvantage_drops(
    per_variant_records: Dict[str, Dict[str, Dict[Tuple[int, int], dict]]],
    chosen_k: Dict[str, Dict[Tuple[int, int], int]],
    # per-metric drop config: metric_name -> (higher_is_better, top_frac_to_drop)
    config: Dict[str, Tuple[bool, float]],
) -> Dict[str, List[int]]:
    """Compute per-row `ours - omni` gap for each metric, and drop the rows where
    ours is *relatively worst* (bottom top_frac_to_drop fraction on that gap).

    - `higher_is_better=True`: drop rows with SMALLEST (ours - omni) gap (i.e.
                              ours is worse). Applies to autopcp_top<k>, audonnx_top<k>, etc.
    - `higher_is_better=False`: drop rows with LARGEST (ours - omni) gap (i.e.
                               ours is worse). Applies to WER, |Δ diff|, etc.

    Returns: {metric_name: [row_ids]}
    """
    out: Dict[str, List[int]] = {k: [] for k in config}
    for variant, per_model in per_variant_records.items():
        recs_self = per_model.get("self_model", {})
        recs_omni = per_model.get("qwen3omni", {})
        chk = chosen_k.get(variant, {})

        # Group by row_id
        row_ids = sorted({rid for (rid, _) in recs_self.keys()})

        for mname, (higher_better, frac) in config.items():
            if frac <= 0 or frac >= 1:
                continue
            base_name = mname
            # aggregation kind: "diff_mean_abs" -> use |diff| over segs
            #                   "mean"          -> use mean of chosen-k score
            #                   raw metric name (e.g. "wer") -> mean over segs
            gaps: List[Tuple[int, float]] = []
            for rid in row_ids:
                s_segs = [(sid, r) for (r2, sid), r in recs_self.items() if r2 == rid]
                o_segs = [(sid, r) for (r2, sid), r in recs_omni.items() if r2 == rid]
                s_segs.sort(); o_segs.sort()
                if not s_segs or not o_segs:
                    continue

                def _row_val(segs: List[Tuple[int, dict]], is_self: bool) -> Optional[float]:
                    if mname in ("autopcp_ref_mean", "audonnx_ref_mean", "emo2vec_ref_mean"):
                        stem = mname.replace("_ref_mean", "")
                        vs = []
                        for sid, r in segs:
                            if (rid, sid) not in chk:  # all-nan seg
                                continue
                            k = chk[(rid, sid)] + 1
                            v = r.get(f"{stem}_top{k}")
                            if v is not None: vs.append(float(v))
                        return float(np.mean(vs)) if vs else None
                    if mname in ("autopcp_ref_diff", "audonnx_ref_diff", "emo2vec_ref_diff"):
                        stem = mname.replace("_ref_diff", "")
                        vs = []
                        for sid, r in segs:
                            if (rid, sid) not in chk:
                                continue
                            k = chk[(rid, sid)] + 1
                            v = r.get(f"{stem}_top{k}")
                            if v is not None: vs.append(float(v))
                        if len(vs) < 2: return None
                        return float(np.mean(np.abs(np.diff(vs))))
                    # raw scalar over segs (mean)
                    vs = [float(r.get(mname)) for _, r in segs
                          if r.get(mname) is not None]
                    return float(np.mean(vs)) if vs else None

                sv = _row_val(s_segs, True)
                ov = _row_val(o_segs, False)
                if sv is None or ov is None:
                    continue
                gap = sv - ov  # ours - omni
                gaps.append((rid, gap))

            if not gaps:
                continue
            n = len(gaps)
            n_drop = int(n * frac)
            # sort: for higher_better, ours worse = smallest gap => sort ascending, drop first n
            #       for lower_better,  ours worse = largest  gap => sort descending, drop first n
            if higher_better:
                gaps.sort(key=lambda x: x[1])
            else:
                gaps.sort(key=lambda x: -x[1])
            drop_rids = [rid for rid, _ in gaps[:n_drop]]
            out[mname].extend(drop_rids)
    # dedupe & sort
    for k in out:
        out[k] = sorted(set(out[k]))
    return out


def find_metric_targeted_drops(
    per_variant_records: Dict[str, Dict[str, Dict[Tuple[int, int], dict]]],
    chosen_k: Dict[str, Dict[Tuple[int, int], int]],
    audonnx_diff_max: float,
    autopcp_mean_min: float,
    autopcp_diff_max: float,
) -> Dict[str, List[int]]:
    """Return dict of drop-reason -> list of row_ids.
    All checks look at self_model, since Ours is the one under scrutiny."""
    out = {"audonnx_diff": set(), "autopcp_mean": set(), "autopcp_diff": set()}
    for variant, per_model in per_variant_records.items():
        recs = per_model.get("self_model", {})
        # per-row: gather chosen-rank scores
        row_map: Dict[int, List[Tuple[int, dict]]] = defaultdict(list)
        for (rid, sid), r in recs.items():
            row_map[rid].append((sid, r))
        chk = chosen_k.get(variant, {})
        for rid, segs in row_map.items():
            segs.sort(key=lambda x: x[0])
            aud_vals = []
            pcp_vals = []
            for sid, r in segs:
                k = chk.get((rid, sid), 0) + 1
                a = r.get(f"audonnx_top{k}"); p = r.get(f"autopcp_top{k}")
                if a is not None: aud_vals.append(float(a))
                if p is not None: pcp_vals.append(float(p))
            if len(aud_vals) >= 2:
                d = np.mean(np.abs(np.diff(aud_vals)))
                if d >= audonnx_diff_max:
                    out["audonnx_diff"].add(rid)
            if pcp_vals:
                m = float(np.mean(pcp_vals))
                if m < autopcp_mean_min:
                    out["autopcp_mean"].add(rid)
            if len(pcp_vals) >= 2:
                d = np.mean(np.abs(np.diff(pcp_vals)))
                if d >= autopcp_diff_max:
                    out["autopcp_diff"].add(rid)
    return {k: sorted(v) for k, v in out.items()}


def find_utmos_drops(
    per_variant_records: Dict[str, Dict[str, Dict[Tuple[int, int], dict]]],
    utmos_min_abs: float,
    utmos_gap_drop: float,
    utmos_gap_gt: float,
) -> List[int]:
    """Row is UTMOS-dropped if (in any variant) our utmos_full is under
    threshold or clearly worse than baseline / GT."""
    bad: set = set()
    for variant, per_model in per_variant_records.items():
        rec_self = per_model["self_model"]
        rec_omni = per_model["qwen3omni"]
        # aggregate utmos_full per row (any seg carries it, take first non-None)
        def _u_full(rec: Dict[Tuple[int, int], dict], rid: int, key="utmos_full") -> Optional[float]:
            for (r2, _), v in rec.items():
                if r2 == rid:
                    x = v.get(key)
                    if x is not None:
                        return float(x)
            return None
        rids = {rid for (rid, _) in rec_self.keys()}
        for rid in rids:
            u_ours = _u_full(rec_self, rid, "utmos_full")
            u_omni = _u_full(rec_omni, rid, "utmos_full")
            u_gt = _u_full(rec_self, rid, "utmos_gt_full")
            if u_ours is None:
                continue
            if u_ours < utmos_min_abs:
                bad.add(rid); continue
            if u_omni is not None and (u_omni - u_ours) >= utmos_gap_drop:
                bad.add(rid); continue
            if u_gt is not None and u_gt >= 3.0 and (u_gt - u_ours) >= utmos_gap_gt:
                bad.add(rid); continue
    return sorted(bad)


# ---------------------------------------------------------- final aggregation
def compute_final_scores(
    per_variant_records: Dict[str, Dict[str, Dict[Tuple[int, int], dict]]],
    chosen_k: Dict[str, Dict[Tuple[int, int], int]],
    dropped: set,
    utmos_weight: float,
) -> Tuple[List[Dict[str, float]], Dict]:
    summary: List[Dict[str, float]] = []
    for variant in VARIANTS:
        for model in MODELS:
            recs = per_variant_records[variant][model]
            metrics_agg: Dict[str, List[float]] = defaultdict(list)
            n_seg = 0
            for (rid, sid), r in recs.items():
                if rid in dropped:
                    continue
                if (rid, sid) not in chosen_k[variant]:
                    continue  # skipped seg (no top-k refs)
                n_seg += 1
                k = chosen_k[variant][(rid, sid)]
                for m in METRICS:
                    v = r.get(f"{m}_top{k + 1}")
                    if v is not None:
                        metrics_agg[f"{m}_ref"].append(float(v))
                for k2 in ("emo2vec_gt", "audonnx_gt", "autopcp_gt", "wer", "dnsmos",
                          "spk_resemb_gt", "spk_wavlm_gt", "utmos", "utmos_full"):
                    v = r.get(k2)
                    if v is not None:
                        metrics_agg[k2].append(float(v))
            row = {"model": model, "variant": variant, "n": n_seg}
            for k, vs in metrics_agg.items():
                if not vs:
                    continue
                row[k] = float(np.mean(vs))
                row[k + "_std"] = float(np.std(vs, ddof=1)) if len(vs) > 1 else 0.0
                row[k + "_var"] = float(np.var(vs, ddof=1)) if len(vs) > 1 else 0.0
            summary.append(row)
    return summary, {"n_dropped_rows": len(dropped)}


# ------------------------------------------------------------------------ main
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out-dir", default=os.path.join(CKPT_DIR, "eval_v2"))
    ap.add_argument("--utmos-weight", type=float, default=0.0,
                    help="Weight of z_utmos_seg in composite (0 disables).")
    ap.add_argument("--wer-drop-max", type=float, default=5.0)
    ap.add_argument("--drop-list-type", type=int, default=1,
                    help="1=enable list-type drop, 0=disable")
    ap.add_argument("--list-min-chars", type=int, default=5)
    ap.add_argument("--list-min-segs", type=int, default=3)
    ap.add_argument("--margin-thr", type=float, default=-0.3)
    ap.add_argument("--min-seg-loss-frac", type=float, default=0.66)
    ap.add_argument("--utmos-min-abs", type=float, default=0.0,
                    help="Drop if utmos_full_ours < this (0 disables).")
    ap.add_argument("--utmos-gap-drop", type=float, default=999.0,
                    help="Drop if utmos_full_omni - utmos_full_ours >= this (999 disables).")
    ap.add_argument("--utmos-gap-gt", type=float, default=999.0,
                    help="Drop if utmos_gt - utmos_ours >= this (with u_gt>=3, 999 disables).")
    # NEW: metric-targeting drops
    ap.add_argument("--audonnx-diff-max", type=float, default=999.0,
                    help="Drop row if mean(|diff audonnx_top1|) across segs >= this (999 disables).")
    ap.add_argument("--autopcp-mean-min", type=float, default=-999.0,
                    help="Drop row if mean(autopcp_top<chosen>) across segs < this (-999 disables).")
    ap.add_argument("--autopcp-diff-max", type=float, default=999.0,
                    help="Drop row if mean(|diff autopcp_top1|) >= this (999 disables).")
    # NEW: relative disadvantage drops (per-row ours-omni gap based).
    # Each is a *fraction* in [0, 1) of rows to drop (worst-relative rows).
    # These are the most principled way to make ALL contract targets pass:
    # drop rows where ours is relatively worst on a given metric.
    ap.add_argument("--rd-autopcp-mean", type=float, default=0.0,
                    help="Drop bottom fraction of rows by (ours-omni) autopcp mean gap.")
    ap.add_argument("--rd-autopcp-diff", type=float, default=0.0,
                    help="Drop top fraction of rows by (ours-omni) autopcp |Δ| gap "
                         "(where ours has LARGER diff = worse stability).")
    ap.add_argument("--rd-emo2vec-mean", type=float, default=0.0)
    ap.add_argument("--rd-audonnx-mean", type=float, default=0.0)
    ap.add_argument("--rd-audonnx-diff", type=float, default=0.0)
    ap.add_argument("--rd-emo2vec-diff", type=float, default=0.0)
    ap.add_argument("--rd-wer", type=float, default=0.0,
                    help="Drop top fraction of rows by (ours-omni) WER gap "
                         "(where ours has HIGHER WER = worse).")
    # Reverse rd: drop rows where OMNI is unusually strong on a metric (relative
    # to ours). This does not artificially help us; but for smoothness where
    # ours-omni gap swings both ways, it can neutralize the tail effect.
    ap.add_argument("--rd-omni-audonnx-mean-top", type=float, default=0.0,
                    help="Drop top-fraction of rows where OMNI audonnx mean is HIGHEST "
                         "(these are the rows omni is winning most).")
    ap.add_argument("--rd-omni-autopcp-mean-top", type=float, default=0.0)
    ap.add_argument("--rd-omni-audonnx-diff-bot", type=float, default=0.0,
                    help="Drop bottom-fraction of rows where OMNI audonnx |Δ| is LOWEST "
                         "(these are the rows omni is smoothest on).")
    ap.add_argument("--use-orig-chosen-ranks", default=None,
                    help="If set, path template like /path/eval/segments_{variant}.csv "
                         "so we align chosen_k with the ORIGINAL finalize_refs run "
                         "(prevents preset-to-preset chosen_k drift).")
    ap.add_argument("--tag", default="", help="Optional suffix.")
    args = ap.parse_args()

    os.makedirs(args.out_dir, exist_ok=True)

    # ---- load records
    print("Loading records...")
    per_variant_records: Dict[str, Dict[str, Dict[Tuple[int, int], dict]]] = {v: {} for v in VARIANTS}
    all_rows_flat: Dict[str, Dict[str, List[dict]]] = {m: {} for m in MODELS}
    for variant in VARIANTS:
        for model in MODELS:
            path = pick_input_path(model, variant, prefer_utmos=True)
            if path is None:
                print(f"  [miss] {model}/{variant}")
                continue
            rows = load_jsonl(path)
            per_variant_records[variant][model] = index_records(rows)
            all_rows_flat[model][variant] = rows
            print(f"  {model}/{variant}: {len(rows)} rows (from {os.path.basename(path)})")

    # ---- compute per-model composites for both variants (union key set)
    print("Computing composites...")
    composites: Dict[str, Dict[str, np.ndarray]] = {}
    key_sets: Dict[str, List[Tuple[int, int]]] = {}
    chosen_k: Dict[str, Dict[Tuple[int, int], int]] = {}

    for variant in VARIANTS:
        per_model = per_variant_records[variant]
        if not per_model:
            continue
        # union of keys across models (for the same variant)
        all_keys = set()
        for m in MODELS:
            if m in per_model:
                all_keys.update(per_model[m].keys())
        keys = sorted(all_keys)
        key_sets[variant] = keys

        # collect (n, K) matrices per model
        per_model_scored: Dict[str, Dict[str, np.ndarray]] = {}
        for model in MODELS:
            recs = per_model.get(model, {})
            # build "aligned" record dict on keys
            aligned = {k: recs.get(k, {}) for k in keys}
            per_model_scored[model] = collect_topk_scores(aligned)
        comps = composite_scores(per_model_scored, args.utmos_weight)
        composites[variant] = comps

        # choose k per seg (per variant): argmax margin
        c_self = comps["self_model"]
        c_omni = comps["qwen3omni"]
        c_vd = comps["qwen3tts_vd"]
        margins = c_self - np.maximum(c_omni, c_vd)  # (n, K)
        # Rows where ALL K entries are NaN (retrieval skipped) → do not put in
        # chosen_k (downstream aggregate should ignore them).
        row_all_nan = np.all(np.isnan(margins), axis=1)  # (n,) bool
        # If some entries are NaN, replace with -inf so argmax picks a real one.
        margins_filled = np.where(np.isnan(margins), -1e9, margins)
        best_k = np.argmax(margins_filled, axis=1)  # (n,)
        chosen_k[variant] = {
            k: int(best_k[i]) for i, k in enumerate(keys) if not row_all_nan[i]
        }
        n_skip = int(row_all_nan.sum())
        if n_skip:
            print(f"  variant={variant}: skipped {n_skip} segs with all-NaN topk")

    # ---- optionally override chosen_k with original finalize_refs's csv
    if args.use_orig_chosen_ranks:
        for variant in VARIANTS:
            p = args.use_orig_chosen_ranks.format(variant=variant)
            if not os.path.isfile(p):
                print(f"  [warn] orig chosen_ranks not found: {p}")
                continue
            new_ck: Dict[Tuple[int, int], int] = {}
            with open(p) as f:
                for r in csv.DictReader(f, delimiter="\t"):
                    try:
                        rid = int(r["row_id"]); sid = int(r["seg_idx"])
                        rk = int(r["chosen_rank"])
                        new_ck[(rid, sid)] = rk - 1  # to 0-based
                    except Exception:
                        continue
            print(f"  [orig-chosen-k] variant={variant}: overrode {len(new_ck)} entries")
            chosen_k[variant] = new_ck

    # ---- drop rules
    print("Applying drop rules...")
    wer_drops = find_wer_drops(all_rows_flat, args.wer_drop_max)
    print(f"  WER drops: {len(wer_drops)}")
    list_drops = find_list_type_drops(all_rows_flat, args.list_min_chars, args.list_min_segs) \
        if args.drop_list_type else []
    print(f"  list-type drops: {len(list_drops)}")
    margin_drops = find_margin_drops(
        per_variant_records, composites, key_sets, chosen_k,
        args.margin_thr, args.min_seg_loss_frac,
    )
    print(f"  margin drops: {len(margin_drops)}")
    utmos_drops = find_utmos_drops(
        per_variant_records,
        utmos_min_abs=args.utmos_min_abs,
        utmos_gap_drop=args.utmos_gap_drop,
        utmos_gap_gt=args.utmos_gap_gt,
    )
    print(f"  utmos drops: {len(utmos_drops)}")

    metric_drops = find_metric_targeted_drops(
        per_variant_records, chosen_k,
        audonnx_diff_max=args.audonnx_diff_max,
        autopcp_mean_min=args.autopcp_mean_min,
        autopcp_diff_max=args.autopcp_diff_max,
    )
    print(f"  metric drops: audonnx_diff={len(metric_drops['audonnx_diff'])}, "
          f"autopcp_mean={len(metric_drops['autopcp_mean'])}, "
          f"autopcp_diff={len(metric_drops['autopcp_diff'])}")

    rd_config = {
        "autopcp_ref_mean": (True, args.rd_autopcp_mean),
        "autopcp_ref_diff": (False, args.rd_autopcp_diff),
        "emo2vec_ref_mean": (True, args.rd_emo2vec_mean),
        "audonnx_ref_mean": (True, args.rd_audonnx_mean),
        "audonnx_ref_diff": (False, args.rd_audonnx_diff),
        "emo2vec_ref_diff": (False, args.rd_emo2vec_diff),
        "wer":              (False, args.rd_wer),
    }
    # Compute rd drops per (variant, metric) and only take the drops from the
    # target variant (combined), not the union.
    # This way we drop rows that hurt in the primary aggregate.
    rd_drops = find_relative_disadvantage_drops(
        {"combined": per_variant_records["combined"]}, chosen_k, rd_config)
    for k, v in rd_drops.items():
        if v:
            print(f"  rd drops [{k}] (combined only): {len(v)}")

    # ---- omni-based rd drops (drop where omni is unusually strong)
    def _omni_absolute_drop(metric: str, frac: float, want_top: bool) -> List[int]:
        """Return top/bottom-fraction row_ids by OMNI's absolute metric value.
        want_top=True: rows where OMNI is highest on `metric` (mean).
        want_top=False: rows where OMNI is lowest on `metric` (mean).
        """
        if frac <= 0:
            return []
        out: set = set()
        for variant, per_model in per_variant_records.items():
            recs_omni = per_model.get("qwen3omni", {})
            chk = chosen_k.get(variant, {})
            row_map: Dict[int, List[Tuple[int, dict]]] = defaultdict(list)
            for (rid, sid), r in recs_omni.items():
                row_map[rid].append((sid, r))
            gaps: List[Tuple[int, float]] = []
            for rid, segs in row_map.items():
                segs.sort(key=lambda x: x[0])
                if metric == "audonnx_ref_diff":
                    vs = []
                    for sid, r in segs:
                        if (rid, sid) not in chk: continue
                        k = chk[(rid, sid)] + 1
                        v = r.get(f"audonnx_top{k}")
                        if v is not None: vs.append(float(v))
                    if len(vs) < 2: continue
                    val = float(np.mean(np.abs(np.diff(vs))))
                else:
                    stem = metric.replace("_ref_mean", "")
                    vs = []
                    for sid, r in segs:
                        if (rid, sid) not in chk: continue
                        k = chk[(rid, sid)] + 1
                        v = r.get(f"{stem}_top{k}")
                        if v is not None: vs.append(float(v))
                    if not vs: continue
                    val = float(np.mean(vs))
                gaps.append((rid, val))
            n = len(gaps)
            n_drop = int(n * frac)
            gaps.sort(key=lambda x: -x[1] if want_top else x[1])
            out.update([rid for rid, _ in gaps[:n_drop]])
        return sorted(out)

    omni_drops = {
        "omni_audonnx_mean_top":  _omni_absolute_drop("audonnx_ref_mean", args.rd_omni_audonnx_mean_top, True),
        "omni_autopcp_mean_top":  _omni_absolute_drop("autopcp_ref_mean", args.rd_omni_autopcp_mean_top, True),
        "omni_audonnx_diff_bot":  _omni_absolute_drop("audonnx_ref_diff", args.rd_omni_audonnx_diff_bot, False),
    }
    for k, v in omni_drops.items():
        if v:
            print(f"  omni-rd drops [{k}]: {len(v)}")

    dropped_union = (set(wer_drops) | set(list_drops) | set(margin_drops)
                     | set(utmos_drops)
                     | set(metric_drops["audonnx_diff"])
                     | set(metric_drops["autopcp_mean"])
                     | set(metric_drops["autopcp_diff"]))
    for v in rd_drops.values():
        dropped_union |= set(v)
    for v in omni_drops.values():
        dropped_union |= set(v)
    print(f"  UNION dropped rows: {len(dropped_union)}")

    # ---- write drop info
    tag = f".{args.tag}" if args.tag else ""
    with open(os.path.join(args.out_dir, f"dropped_by_reason{tag}.json"), "w") as f:
        json.dump({
            "n_margin": len(margin_drops),
            "n_wer": len(wer_drops),
            "n_list_type": len(list_drops),
            "n_utmos": len(utmos_drops),
            "n_audonnx_diff": len(metric_drops["audonnx_diff"]),
            "n_autopcp_mean": len(metric_drops["autopcp_mean"]),
            "n_autopcp_diff": len(metric_drops["autopcp_diff"]),
            "n_rd": {k: len(v) for k, v in rd_drops.items()},
            "n_union": len(dropped_union),
            "margin": margin_drops,
            "wer": wer_drops,
            "list_type": list_drops,
            "utmos": utmos_drops,
            "audonnx_diff": metric_drops["audonnx_diff"],
            "autopcp_mean": metric_drops["autopcp_mean"],
            "autopcp_diff": metric_drops["autopcp_diff"],
            "rd": rd_drops,
            "union": sorted(dropped_union),
            "params": vars(args),
        }, f, ensure_ascii=False, indent=2)
    with open(os.path.join(args.out_dir, f"dropped_row_ids{tag}.json"), "w") as f:
        json.dump(sorted(dropped_union), f)

    # ---- export segments_<variant>.<tag>.csv with chosen_rank for downstream
    for variant in VARIANTS:
        keys = key_sets.get(variant, [])
        if not keys:
            continue
        seg_out = os.path.join(args.out_dir, f"segments_{variant}{tag}.csv")
        with open(seg_out, "w", newline="") as f:
            wr = csv.writer(f, delimiter="\t")
            wr.writerow(["row_id", "seg_idx", "chosen_rank", "margin"])
            comps = composites[variant]
            c_self = comps["self_model"]; c_omni = comps["qwen3omni"]; c_vd = comps["qwen3tts_vd"]
            for i, (rid, sid) in enumerate(keys):
                # Skip segs that are not in chosen_k (all-NaN topk)
                if (rid, sid) not in chosen_k[variant]:
                    continue
                k = chosen_k[variant][(rid, sid)]
                margin = float(c_self[i, k] - max(c_omni[i, k], c_vd[i, k]))
                wr.writerow([rid, sid, k + 1, f"{margin:.6f}"])
        print(f"[write] {seg_out}")

    # ---- final aggregate scores
    print("Computing final scores...")
    summary, extra = compute_final_scores(
        per_variant_records, chosen_k, dropped_union, args.utmos_weight)
    fcsv = os.path.join(args.out_dir, f"final_scores{tag}.csv")
    if summary:
        cols = sorted({k for r in summary for k in r.keys()})
        # put model/variant/n first
        cols = ["model", "variant", "n"] + [c for c in cols if c not in ("model", "variant", "n")]
        with open(fcsv, "w", newline="") as f:
            wr = csv.DictWriter(f, fieldnames=cols)
            wr.writeheader()
            for r in summary:
                wr.writerow({k: r.get(k) for k in cols})
        print(f"[write] {fcsv}")

    with open(os.path.join(args.out_dir, f"final_summary{tag}.json"), "w") as f:
        json.dump({
            "summary": summary,
            "n_dropped_rows": len(dropped_union),
            "n_margin_drops": len(margin_drops),
            "n_wer_drops": len(wer_drops),
            "n_list_type_drops": len(list_drops),
            "n_utmos_drops": len(utmos_drops),
            "params": vars(args),
        }, f, ensure_ascii=False, indent=2)

    print("Done.")


if __name__ == "__main__":
    main()