analyze_speaker_pairs.py

from __future__ import annotations

import argparse
import logging
import re
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("pair_analysis")


plt.rcParams.update({
    "figure.dpi":           130,
    "savefig.dpi":          240,
    "savefig.bbox":         "tight",
    "savefig.facecolor":    "white",
    "figure.facecolor":     "white",
    "axes.facecolor":       "#FAFAFA",
    "font.family":          "DejaVu Sans",
    "font.size":            11,
    "axes.titlesize":       13,
    "axes.titleweight":     "semibold",
    "axes.titlepad":        14,
    "axes.labelsize":       11,
    "axes.spines.top":      False,
    "axes.spines.right":    False,
    "axes.edgecolor":       "#333333",
    "axes.linewidth":       1.0,
    "axes.grid":            True,
    "axes.axisbelow":       True,
    "grid.color":           "#CCCCCC",
    "grid.linestyle":       "-",
    "grid.linewidth":       0.6,
    "grid.alpha":           0.5,
    "legend.frameon":       True,
    "legend.framealpha":    0.95,
    "legend.edgecolor":     "#CCCCCC",
})

C_CROSS = "#009E73"
C_SAME  = "#CC79A7"


def parse_speakers_file(speakers_file: Path) -> Dict[str, str]:
    gender: Dict[str, str] = {}
    with open(speakers_file, "r", encoding="utf-8") as f:
        for line in f:
            if line.startswith(";") or not line.strip():
                continue
            parts = [p.strip() for p in line.split("|")]
            if len(parts) >= 2:
                gender[parts[0]] = parts[1].upper()
    return gender


SPEAKER_PAIR_RE = re.compile(r"(\d+)-(\d+)")


def extract_speaker_pair(file_id: str) -> Optional[Tuple[str, str]]:
    matches = SPEAKER_PAIR_RE.findall(file_id)
    if not matches:
        return None
    return matches[0]


def classify_pair(file_id: str, gender_map: Dict[str, str]) -> Optional[str]:
    pair = extract_speaker_pair(file_id)
    if pair is None:
        return None
    spk_a, spk_b = pair
    g_a = gender_map.get(spk_a)
    g_b = gender_map.get(spk_b)
    if g_a is None or g_b is None:
        return None
    return "cross" if g_a != g_b else "same"


def analyse_condition(csv_path: Path,
                      gender_map: Dict[str, str]) -> Optional[Dict]:
    if not csv_path.exists():
        return None

    df = pd.read_csv(csv_path)
    df = df[~df["file_id"].isin(["MEAN", "STD"])].copy()
    if df.empty:
        return None

    df["pair_type"] = df["file_id"].apply(lambda fid: classify_pair(fid, gender_map))

    n_unknown = df["pair_type"].isna().sum()
    if n_unknown > 0:
        log.warning(f"  · {csv_path.parent.name}: {n_unknown} file(s) had "
                    f"unknown speaker pairs (skipped from analysis).")

    df = df.dropna(subset=["pair_type"])
    if df.empty:
        return None

    out: Dict = {}
    for grp in ("cross", "same"):
        sub = df[df["pair_type"] == grp]
        if len(sub) > 0:
            out[f"{grp}_n"]    = int(len(sub))
            out[f"{grp}_mean"] = float(sub["der"].mean())
            out[f"{grp}_std"]  = float(sub["der"].std(ddof=0))
        else:
            out[f"{grp}_n"]    = 0
            out[f"{grp}_mean"] = float("nan")
            out[f"{grp}_std"]  = float("nan")
    out["delta"] = out["same_mean"] - out["cross_mean"]
    return out


def collect_results(baseline_dir: Path, nmf_dir: Path, spectral_dir: Path,
                    gender_map: Dict[str, str]) -> List[Dict]:
    rows: List[Dict] = []

    def add(label: str, csv_path: Path) -> None:
        result = analyse_condition(csv_path, gender_map)
        if result is None:
            log.info(f"  · Skipping {label}: no CSV at {csv_path}")
            return
        rows.append({"condition": label, **result})

    add("Set A — baseline (clean)",
        baseline_dir / "setA_clean" / "der_results.csv")

    for pct in (10, 20, 30):
        add(f"Set B baseline — {pct}% overlap",
            baseline_dir / "setB_overlap" / f"overlap_{pct:02d}pct" / "der_results.csv")

    for pct in (10, 20, 30):
        add(f"Set B NMF+VAD — {pct}% overlap",
            nmf_dir / "setB_overlap" / f"overlap_{pct:02d}pct" / "der_results.csv")

    for snr in (5, 10, 15):
        add(f"Set C baseline — {snr} dB",
            baseline_dir / "setC_noise" / f"snr_{snr:02d}dB" / "der_results.csv")

    for snr in (5, 10, 15):
        add(f"Set C Spectral Sub — {snr} dB",
            spectral_dir / "setC_noise" / f"snr_{snr:02d}dB" / "der_results.csv")

    return rows


def write_csv(rows: List[Dict], out_path: Path) -> None:
    out_path.parent.mkdir(parents=True, exist_ok=True)
    df = pd.DataFrame(rows)
    for col in ("cross_mean", "cross_std", "same_mean", "same_std", "delta"):
        if col in df.columns:
            df[col] = df[col].round(2)
    df.to_csv(out_path, index=False)
    log.info(f"  ✓ Summary CSV → {out_path}")


def print_table(rows: List[Dict]) -> None:
    from rich.console import Console
    from rich.table import Table

    console = Console()
    table = Table(title="Cross-gender vs Same-gender DER",
                  show_lines=False)
    table.add_column("Condition", overflow="fold", no_wrap=False)
    table.add_column("Cross n",        justify="right")
    table.add_column("Cross DER",      justify="right")
    table.add_column("Same n",         justify="right")
    table.add_column("Same DER",       justify="right")
    table.add_column("Δ (Same − Cross)", justify="right")

    for r in rows:
        cross = (f"{r['cross_mean']:.1f}% ± {r['cross_std']:.1f}"
                 if r["cross_n"] > 0 else "—")
        same  = (f"{r['same_mean']:.1f}% ± {r['same_std']:.1f}"
                 if r["same_n"]  > 0 else "—")
        delta = (f"{r['delta']:+.1f} pp"
                 if r["cross_n"] > 0 and r["same_n"] > 0 else "—")
        table.add_row(r["condition"], str(r["cross_n"]), cross,
                      str(r["same_n"]), same, delta)
    console.print(table)


def plot_grouped_bars(rows: List[Dict], out_path: Path) -> None:
    if not rows:
        return

    labels       = [r["condition"] for r in rows]
    cross_means  = [r["cross_mean"] for r in rows]
    cross_stds   = [r["cross_std"]  for r in rows]
    same_means   = [r["same_mean"]  for r in rows]
    same_stds    = [r["same_std"]   for r in rows]

    x = np.arange(len(labels))
    width = 0.38

    fig, ax = plt.subplots(figsize=(max(9, 0.8 * len(labels)), 5.5))

    ax.bar(x - width/2, cross_means, width,
           yerr=cross_stds, capsize=4,
           error_kw={"capthick": 1.2, "elinewidth": 1.2},
           label="Cross-gender (M/F)", color=C_CROSS,
           edgecolor="white", linewidth=0.8)
    ax.bar(x + width/2, same_means,  width,
           yerr=same_stds,  capsize=4,
           error_kw={"capthick": 1.2, "elinewidth": 1.2},
           label="Same-gender (M/M or F/F)", color=C_SAME,
           edgecolor="white", linewidth=0.8)

    ax.set_ylabel("Diarization Error Rate (%)")
    ax.set_title("DER by speaker-pair gender composition", loc="left")
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=30, ha="right")
    ax.grid(axis="x", visible=False)
    ax.legend(loc="upper left", title=None)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path)
    plt.close(fig)
    log.info(f"  ✓ Grouped bar chart → {out_path}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Cross-gender vs same-gender DER analysis.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--speakers_file", type=Path, required=True,
                        help="Path to LibriSpeech/SPEAKERS.TXT")
    parser.add_argument("--project_root", type=Path, default=Path("."),
                        help="Project root. Default: current directory.")
    parser.add_argument("--baseline_dir", type=Path, default=None,
                        help="Override baseline predictions dir. "
                             "Default: <project_root>/baseline_predictions")
    parser.add_argument("--nmf_dir",      type=Path, default=None,
                        help="Override NMF predictions dir. "
                             "Default: <project_root>/nmf_predictions")
    parser.add_argument("--spectral_dir", type=Path, default=None,
                        help="Override Spectral Subtraction predictions dir. "
                             "Default: <project_root>/spectral_predictions")
    parser.add_argument("--output_dir",   type=Path, default=None,
                        help="Where to write the figure and summary CSV. "
                             "Default: <project_root>/graphs")
    args = parser.parse_args()

    if not args.speakers_file.is_file():
        log.error(f"--speakers_file not found: {args.speakers_file}")
        return

    baseline_dir = args.baseline_dir or args.project_root / "baseline_predictions"
    nmf_dir      = args.nmf_dir      or args.project_root / "nmf_predictions"
    spectral_dir = args.spectral_dir or args.project_root / "spectral_predictions"
    output_dir   = args.output_dir   or args.project_root / "graphs"

    log.info("Loading speaker gender map from SPEAKERS.TXT...")
    gender_map = parse_speakers_file(args.speakers_file)
    log.info(f"  Found {len(gender_map)} speakers "
             f"({sum(1 for g in gender_map.values() if g == 'M')} M, "
             f"{sum(1 for g in gender_map.values() if g == 'F')} F)")

    log.info("")
    log.info("Walking prediction folders...")
    rows = collect_results(baseline_dir, nmf_dir, spectral_dir, gender_map)
    if not rows:
        log.error("No CSVs found. Has evaluate_der.py been run?")
        return

    log.info("")
    print_table(rows)

    write_csv(rows, output_dir / "speaker_pair_analysis.csv")
    plot_grouped_bars(rows, output_dir / "graph10_speaker_pair_analysis.png")

    log.info("")
    log.info("✓ Done.")


if __name__ == "__main__":
    main()