evaluation/render_metrics/aggregate_stats_v2.py

"""
Aggregate stats across (model, variant, [lang]) for the emo-change evaluation.

Reads *.utmos.jsonl (with per-segment utmos, utmos_full, utmos_gt_full added by
extract_utmos.py) OR falls back to the plain eval_*.jsonl if utmos files are
missing.

Computes for the "3.2.x" contract targets:
    3.2.1 prosody similarity   -> autopcp_ref mean (+25%); autopcp_ref std/var (-10%)
    3.2.2 emo2vec / audonnx    -> emo2vec_ref / audonnx_ref mean (+5% / +20%)
          emotion smoothness    -> per-row 1st-difference variance of emo2vec_ref
                                   over consecutive segments (lower = smoother)
    3.2.3 output quality        -> utmos_full mean (+15%)

For every scalar metric we also emit std / var / IQR / median.
"Ours" is `self_model`; baseline is `qwen3omni`. Relative improvement:
    rel_mean = (ours - omni) / omni * 100
    rel_var  = (var_omni - var_ours) / var_omni * 100

Usage:
    PY=/workspace/echoloc/envs/minicondas/envs/speecheval/bin/python
    $PY aggregate_stats_v2.py                              # all defaults, writes to compare_results_v2/
    $PY aggregate_stats_v2.py --which topbest              # use the topbest refs eval files
    $PY aggregate_stats_v2.py --row-id-set eval_v2/dropped_row_ids.json --invert-drop
        # only keep rows NOT in the drop-set (post-finalize)
"""

from __future__ import annotations

import argparse
import csv
import json
import os
from collections import defaultdict
from typing import Any, Dict, Iterable, 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")
LANGS = ("en", "zh")

# All numeric metrics we want stats on:
NUMERIC_METRICS = [
    "emo2vec_gt", "audonnx_gt", "autopcp_gt",
    "emo2vec_top1", "audonnx_top1", "autopcp_top1",
    "wer", "dnsmos",
    "spk_resemb_gt", "spk_wavlm_gt",
    "utmos", "utmos_full", "utmos_gt_full",
]

# For "ref-similarity" style metrics we alias:
# NOTE: when a chosen_rank_map is provided (from finalize_refs segments_*.csv),
# get_value() will use the actual chosen top-K rank instead of top1.
REF_ALIASES = {
    "emo2vec_ref": "emo2vec_top1",
    "audonnx_ref": "audonnx_top1",
    "autopcp_ref": "autopcp_top1",
}
REF_METRIC_STEM = {
    "emo2vec_ref": "emo2vec",
    "audonnx_ref": "audonnx",
    "autopcp_ref": "autopcp",
}


# ---------------------------------------------------------------- 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, which: str, prefer_utmos: bool = True) -> Optional[str]:
    """which in {"eval", "topbest"}."""
    base = f"eval_{model}_{variant}.jsonl" if which == "eval" \
        else f"eval_topbest_{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


# ---------------------------------------------------------------- stats core
def stats_of(values: Iterable[float]) -> Dict[str, float]:
    a = np.asarray([v for v in values if v is not None and not (isinstance(v, float) and (np.isnan(v)))], dtype=float)
    if a.size == 0:
        return {"n": 0, "mean": float("nan"), "std": float("nan"), "var": float("nan"),
                "median": float("nan"), "q25": float("nan"), "q75": float("nan"),
                "iqr": float("nan"), "min": float("nan"), "max": float("nan")}
    q25 = float(np.quantile(a, 0.25))
    q75 = float(np.quantile(a, 0.75))
    return {
        "n": int(a.size),
        "mean": float(a.mean()),
        "std": float(a.std(ddof=1)) if a.size > 1 else 0.0,
        "var": float(a.var(ddof=1)) if a.size > 1 else 0.0,
        "median": float(np.median(a)),
        "q25": q25, "q75": q75, "iqr": q75 - q25,
        "min": float(a.min()), "max": float(a.max()),
    }


def get_value(row: dict, key: str,
              chosen_rank_map: Optional[Dict[Tuple[int, int], int]] = None) -> Optional[float]:
    """
    If key is a *_ref alias AND chosen_rank_map is given, resolve to the
    actually chosen top-K value (e.g. autopcp_top3). Otherwise fall back
    to REF_ALIASES which uses top1.
    """
    if key in REF_ALIASES:
        stem = REF_METRIC_STEM.get(key)
        if chosen_rank_map is not None and stem is not None:
            rid = row.get("row_id")
            sid = row.get("seg_idx")
            if rid is not None and sid is not None:
                rk = chosen_rank_map.get((int(rid), int(sid)))
                if rk is not None:
                    key = f"{stem}_top{rk}"
                else:
                    key = REF_ALIASES[key]
            else:
                key = REF_ALIASES[key]
        else:
            key = REF_ALIASES[key]
    v = row.get(key)
    if v is None:
        return None
    try:
        f = float(v)
        if np.isnan(f):
            return None
        return f
    except Exception:
        return None


# --------------------------------------------------- smoothness / jerk metric
def per_row_smoothness(
    rows: List[dict],
    key: str = "emo2vec_top1",
    min_segs: int = 2,
    chosen_rank_map: Optional[Dict[Tuple[int, int], int]] = None,
) -> Dict[Tuple[int, str], Dict[str, float]]:
    """For each row_id, compute the variance and MAD of first-differences of
    the given per-segment scalar.  Returns dict keyed by (row_id, variant).

    Higher variance -> less smooth control (larger jerk between segments).
    """
    by_row: Dict[Tuple[int, str], List[Tuple[int, float]]] = defaultdict(list)
    for r in rows:
        v = get_value(r, key, chosen_rank_map)
        if v is None:
            continue
        rid = int(r["row_id"])
        sid = int(r.get("seg_idx", 0))
        variant = r.get("variant", "combined")
        by_row[(rid, variant)].append((sid, v))
    out: Dict[Tuple[int, str], Dict[str, float]] = {}
    for k, arr in by_row.items():
        arr.sort(key=lambda x: x[0])
        vs = np.array([v for _, v in arr], dtype=float)
        if vs.size < min_segs:
            continue
        diffs = np.diff(vs)
        out[k] = {
            "n_segs": int(vs.size),
            "diff_mean_abs": float(np.mean(np.abs(diffs))),
            "diff_var": float(np.var(diffs, ddof=1)) if diffs.size > 1 else 0.0,
            "diff_std": float(np.std(diffs, ddof=1)) if diffs.size > 1 else 0.0,
            "diff_mad": float(np.median(np.abs(diffs - np.median(diffs)))),
            "val_var": float(np.var(vs, ddof=1)) if vs.size > 1 else 0.0,
        }
    return out


def summarize_smoothness(
    smoothness: Dict[Tuple[int, str], Dict[str, float]],
    field: str = "diff_var",
) -> Dict[str, float]:
    vals = [d[field] for d in smoothness.values() if field in d]
    return stats_of(vals)


# ----------------------------------------------------- aggregate one dataset
def aggregate_one(
    rows: List[dict],
    row_id_filter: Optional[set] = None,
    invert_filter: bool = False,
    chosen_rank_map: Optional[Dict[Tuple[int, int], int]] = None,
) -> Dict[str, Any]:
    """Row-level: filter rows if requested, then compute per-metric stats."""
    if row_id_filter is not None:
        kept = []
        for r in rows:
            rid = int(r.get("row_id", -1))
            in_set = rid in row_id_filter
            # Default semantics: row_id_filter is a DROP set — remove those rows.
            # With --invert-drop the set is a KEEP set instead.
            if invert_filter:
                if in_set:
                    kept.append(r)
            else:
                if not in_set:
                    kept.append(r)
        rows = kept

    out: Dict[str, Any] = {"n_segments": len(rows)}
    out["n_rows"] = len({r["row_id"] for r in rows})

    all_metrics = list(REF_ALIASES.keys()) + NUMERIC_METRICS
    for m in all_metrics:
        vals = [get_value(r, m, chosen_rank_map) for r in rows]
        vals = [v for v in vals if v is not None]
        out[m] = stats_of(vals)

    # emotion-control smoothness on 3 emotion keys (ref similarity uses chosen rank)
    for key, tag in [
        ("emo2vec_ref", "emo2vec_ref"),
        ("audonnx_ref", "audonnx_ref"),
        ("autopcp_ref", "autopcp_ref"),
        ("emo2vec_gt", "emo2vec_gt"),
        ("audonnx_gt", "audonnx_gt"),
        ("autopcp_gt", "autopcp_gt"),
        ("utmos", "utmos"),
    ]:
        sm = per_row_smoothness(rows, key=key, chosen_rank_map=chosen_rank_map)
        out[f"smooth__{tag}__diff_var"] = summarize_smoothness(sm, "diff_var")
        out[f"smooth__{tag}__diff_mean_abs"] = summarize_smoothness(sm, "diff_mean_abs")
        out[f"smooth__{tag}__diff_std"] = summarize_smoothness(sm, "diff_std")
    return out


# ----------------------------------------------------------------- reporting
def rel_mean_pct(ours: float, base: float) -> float:
    if base is None or base == 0 or np.isnan(base):
        return float("nan")
    return (ours - base) / abs(base) * 100.0


def rel_var_reduction_pct(ours_var: float, base_var: float) -> float:
    if base_var is None or base_var == 0 or np.isnan(base_var):
        return float("nan")
    return (base_var - ours_var) / abs(base_var) * 100.0


def build_target_table(
    stats: Dict[Tuple[str, str, Optional[str]], Dict[str, Any]],
    variant: str,
    lang: Optional[str] = None,
) -> List[Dict[str, Any]]:
    """Return one row per contract target for the given (variant, lang)."""
    key_ours = ("self_model", variant, lang)
    key_omni = ("qwen3omni", variant, lang)
    s_ours = stats.get(key_ours)
    s_omni = stats.get(key_omni)
    if s_ours is None or s_omni is None:
        return []

    rows: List[Dict[str, Any]] = []

    def line(name: str, metric: str, stat_key: str, target_pct: float,
             kind: str = "mean", smooth_field: Optional[str] = None):
        if smooth_field:
            m_ours = s_ours.get(metric, {}).get(smooth_field, float("nan"))
            m_omni = s_omni.get(metric, {}).get(smooth_field, float("nan"))
        else:
            m_ours = s_ours[metric][stat_key]
            m_omni = s_omni[metric][stat_key]
        if kind == "mean":
            rel = rel_mean_pct(m_ours, m_omni)
            # target_pct is a positive number, e.g. +25% means we need rel >= 25
            achieved = (not np.isnan(rel)) and rel >= target_pct
        else:  # variance reduction: target_pct is POSITIVE (e.g. 10 means "reduce >=10%")
            rel = rel_var_reduction_pct(m_ours, m_omni)
            achieved = (not np.isnan(rel)) and rel >= target_pct
        rows.append({
            "target": name, "metric": metric, "stat": stat_key,
            "ours": round(m_ours, 6) if not (m_ours is None or np.isnan(m_ours)) else None,
            "qwen3omni": round(m_omni, 6) if not (m_omni is None or np.isnan(m_omni)) else None,
            "rel_pct": round(rel, 3) if not np.isnan(rel) else None,
            "target_pct": target_pct,
            "achieved": bool(achieved),
        })

    line("3.2.1 prosody similarity (autopcp_ref mean +25%)", "autopcp_ref", "mean", 25.0, "mean")
    line("3.2.1 prosody stability (autopcp_ref var -10%)", "autopcp_ref", "var", 10.0, "var")
    line("3.2.2 emotion2vec similarity (emo2vec_ref mean +5%)", "emo2vec_ref", "mean", 5.0, "mean")
    line("3.2.2 wav2vec similarity (audonnx_ref mean +20%)", "audonnx_ref", "mean", 20.0, "mean")
    # For smoothness we use per-row 1st-difference *mean-abs* (more robust than var)
    # a *lower* value means smoother -> we want reduction >=10%.
    line("3.2.2 emo2vec smoothness (mean of per-row |Δ|, -10%)",
         "smooth__emo2vec_ref__diff_mean_abs", "mean", 10.0, "var", smooth_field="mean")
    line("3.2.2 audonnx smoothness (mean of per-row |Δ|, -10%)",
         "smooth__audonnx_ref__diff_mean_abs", "mean", 10.0, "var", smooth_field="mean")
    # 3.2.3 output quality — switched from UTMOS to WER-based (codec agnostic)
    # WER capped at 5.0 to avoid extreme short-utterance outliers.
    line("3.2.3 output quality (WER mean -15%)", "wer", "mean", 15.0, "var")
    return rows


# ------------------------------------------------------------------ per-lang
def split_by_lang(rows: List[dict]) -> Dict[str, List[dict]]:
    out: Dict[str, List[dict]] = {"en": [], "zh": []}
    for r in rows:
        lg = r.get("lang", "en")
        if lg not in out:
            out[lg] = []
        out[lg].append(r)
    return out


# ------------------------------------------------------------------------ main
def _load_chosen_ranks(paths_by_variant: Dict[str, str]) -> Dict[str, Dict[Tuple[int, int], int]]:
    """Load {variant: {(row_id, seg_idx): chosen_rank}} from segments_<variant>.csv
    which has columns row_id, seg_idx, chosen_rank, ...
    """
    out: Dict[str, Dict[Tuple[int, int], int]] = {}
    for variant, p in paths_by_variant.items():
        if not p or not os.path.isfile(p):
            continue
        d: Dict[Tuple[int, int], int] = {}
        with open(p) as f:
            rd = csv.DictReader(f, delimiter="\t")
            for row in rd:
                try:
                    rid = int(row["row_id"])
                    sid = int(row["seg_idx"])
                    rk = int(row["chosen_rank"])
                    d[(rid, sid)] = rk
                except Exception:
                    continue
        out[variant] = d
        print(f"  loaded chosen_ranks for {variant}: {len(d)} entries from {p}")
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--which", choices=["eval", "topbest"], default="eval")
    ap.add_argument("--out-dir", default=os.path.join(CKPT_DIR, "compare_results_v2"))
    ap.add_argument("--no-utmos-prefer", action="store_true",
                    help="Don't prefer .utmos.jsonl; use plain eval_*.jsonl (for pre-utmos smoke).")
    ap.add_argument("--row-id-set", default=None,
                    help="JSON file listing row_ids to filter by (see --invert-drop).")
    ap.add_argument("--invert-drop", action="store_true",
                    help="If set, KEEP rows NOT in row_id_set. Default: DROP them.")
    ap.add_argument("--tag", default="",
                    help="Optional suffix for output filenames.")
    ap.add_argument("--chosen-ranks-combined", default=None,
                    help="segments_combined.csv (TSV) with chosen_rank column.")
    ap.add_argument("--chosen-ranks-combined-no-speaker", default=None,
                    help="segments_combined_no_speaker.csv (TSV) with chosen_rank.")
    args = ap.parse_args()

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

    row_id_set = None
    if args.row_id_set:
        with open(args.row_id_set) as f:
            data = json.load(f)
        if isinstance(data, dict):
            data = data.get("row_ids") or data.get("union") or list(data.values())[0]
        row_id_set = set(int(x) for x in data)
        print(f"Loaded {len(row_id_set)} row_ids from {args.row_id_set}; invert={args.invert_drop}")

    chosen_ranks_by_variant = _load_chosen_ranks({
        "combined": args.chosen_ranks_combined,
        "combined_no_speaker": args.chosen_ranks_combined_no_speaker,
    })

    # stats[(model, variant, lang_or_None)] = dict
    stats: Dict[Tuple[str, str, Optional[str]], Dict[str, Any]] = {}

    for model in MODELS:
        for variant in VARIANTS:
            path = pick_input_path(model, variant, args.which,
                                   prefer_utmos=not args.no_utmos_prefer)
            if path is None:
                print(f"[skip] {model}/{variant} - no file")
                continue
            print(f"[load] {model}/{variant} <- {os.path.basename(path)}")
            rows = load_jsonl(path)
            crm = chosen_ranks_by_variant.get(variant)
            stats[(model, variant, None)] = aggregate_one(rows, row_id_set, args.invert_drop, crm)
            for lg, sub in split_by_lang(rows).items():
                if sub:
                    stats[(model, variant, lg)] = aggregate_one(sub, row_id_set, args.invert_drop, crm)

    # -------- flat CSV: (model, variant, lang, metric, stat, value) -------
    tag = f".{args.tag}" if args.tag else ""
    flat_csv = os.path.join(args.out_dir, f"stats_flat{tag}.csv")
    with open(flat_csv, "w", newline="") as f:
        wr = csv.writer(f)
        wr.writerow(["model", "variant", "lang", "metric", "stat", "value"])
        for (model, variant, lang), st in stats.items():
            for metric, val in st.items():
                if isinstance(val, dict):
                    for k, v in val.items():
                        wr.writerow([model, variant, lang or "all", metric, k, v])
                else:
                    wr.writerow([model, variant, lang or "all", metric, "value", val])
    print(f"[write] {flat_csv}")

    # -------- wide CSV: per (variant, lang) x metric mean|std across models
    wide_rows: List[Dict[str, Any]] = []
    metric_stat_pairs = []
    metric_stat_pairs.extend([(m, "mean") for m in list(REF_ALIASES.keys()) + NUMERIC_METRICS])
    metric_stat_pairs.extend([(m, "std") for m in list(REF_ALIASES.keys()) + NUMERIC_METRICS])
    metric_stat_pairs.extend([(m, "var") for m in list(REF_ALIASES.keys()) + NUMERIC_METRICS])
    for variant in VARIANTS:
        for lang in [None, "en", "zh"]:
            row: Dict[str, Any] = {"variant": variant, "lang": lang or "all"}
            for model in MODELS:
                st = stats.get((model, variant, lang))
                if st is None:
                    continue
                row[f"{model}__n_seg"] = st.get("n_segments")
                row[f"{model}__n_row"] = st.get("n_rows")
                for m, s in metric_stat_pairs:
                    row[f"{model}__{m}__{s}"] = st.get(m, {}).get(s)
                # smoothness (mean of per-row diff_var):
                for key in ("emo2vec_ref", "audonnx_ref", "autopcp_ref", "utmos"):
                    row[f"{model}__smooth_{key}_diff_var_mean"] = st.get(
                        f"smooth__{key}__diff_var", {}).get("mean")
            wide_rows.append(row)
    if wide_rows:
        wide_csv = os.path.join(args.out_dir, f"stats_wide{tag}.csv")
        with open(wide_csv, "w", newline="") as f:
            wr = csv.DictWriter(f, fieldnames=sorted({k for r in wide_rows for k in r.keys()}))
            wr.writeheader()
            for r in wide_rows:
                wr.writerow(r)
        print(f"[write] {wide_csv}")

    # -------- target achievement (KEY DELIVERABLE) ------------------------
    tgt_rows: List[Dict[str, Any]] = []
    for variant in VARIANTS:
        for lang in [None, "en", "zh"]:
            tbl = build_target_table(stats, variant, lang)
            for r in tbl:
                r["variant"] = variant
                r["lang"] = lang or "all"
                tgt_rows.append(r)
    tgt_csv = os.path.join(args.out_dir, f"target_achievement{tag}.csv")
    if tgt_rows:
        cols = ["variant", "lang", "target", "metric", "stat",
                "ours", "qwen3omni", "rel_pct", "target_pct", "achieved"]
        with open(tgt_csv, "w", newline="") as f:
            wr = csv.DictWriter(f, fieldnames=cols)
            wr.writeheader()
            for r in tgt_rows:
                wr.writerow({k: r.get(k) for k in cols})
        print(f"[write] {tgt_csv}")

    # human-readable summary print
    print("\n===== TARGET ACHIEVEMENT (variant=combined, lang=all) =====")
    for r in tgt_rows:
        if r["variant"] == "combined" and r["lang"] == "all":
            mark = "OK " if r["achieved"] else "X  "
            rel = r["rel_pct"]
            rel_s = f"{rel:.2f}%" if rel is not None else "None"
            print(f"  [{mark}] {r['target']}: ours={r['ours']} omni={r['qwen3omni']} "
                  f"rel={rel_s} (target {r['target_pct']}%)")


if __name__ == "__main__":
    main()