plot_adaptive_correlations.py

#!/usr/bin/env python3
"""
Plot Figure 2-style correlations for adaptive testing.

This script computes Pearson correlations between estimated scores and
validated clinical scales over the number of questions asked.
It uses:
    - adaptive_outputs/<dir>/{Drule,random}/thetas_*.csv (item order)
    - model_outputs/qa_level_outputs_<strategy>_*/fold_*/qa_level_test_predictions.csv

Optional GPT-4 baseline can be provided via --gpt4-csv.
"""

from __future__ import annotations

import argparse
import ast
import csv
import json
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple
import warnings

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import pearsonr

from maqua_paths import (
    ADAPTIVE_OUTPUTS,
    ANALYSIS_OUTPUTS,
    FORMATTED_DATA,
    MODEL_OUTPUTS,
    POLYTOMIZED_DATA,
)
from maqua_mapping import get_questionnaire_mapping


OUTCOMES = [
    ("PHQ", "Depression"),
    ("GAD", "Anxiety"),
    ("MDQ", "Bipolar"),
    ("RAADS", "Autism"),
    ("DUDIT", "Drug use"),
    ("AUDIT", "Alcohol use"),
    ("BOCS", "OCD"),
    ("ASRS", "ADHD"),
    ("NSE", "PTSD"),
    ("EDE_QS", "Eating disorder"),
]

DISPLAY_TO_KEY = {display: key for key, display in OUTCOMES}
KEY_TO_DISPLAY = {key: display for key, display in OUTCOMES}

DEFAULT_QUESTION_ORDER = [
    "A1",
    "A3",
    "A4",
    "ADHD1",
    "ADHD2",
    "ASD2",
    "ASD3",
    "ASD4",
    "ASD5",
    "ASD6",
    "BD2",
    "BD3",
    "ED1",
    "ED2",
    "ED3",
    "ED4",
    "ED5",
    "ED6",
    "G1",
    "G10",
    "G12",
    "G2",
    "G3",
    "G4",
    "G5",
    "G6",
    "G7",
    "G8",
    "G9",
    "G91",
    "OCD1",
    "OCD2",
    "OCD3",
    "OMD1",
    "OMD2",
    "OMD3",
    "OMD4",
    "OMD5",
    "OMD6",
    "PTSD1",
    "PTSD2",
    "SUB1",
    "SUB2",
    "SUB3",
    "SUB4",
    "SUB5",
    "SUB6",
    "nse",
]

EXCLUDED_CODES = {"A2", "ASD1", "BD1", "G11", "SUB7", "PTSD3"}
STRICT_MODE = False
ITEM_LIST_COERCE_ERROR_COUNT = 0
ITEM_INDEX_PARSE_ERROR_COUNT = 0


def _handle_exception(context: str, exc: Exception, strict: bool | None = None) -> None:
    """Log contextual errors and optionally fail fast in strict mode."""
    use_strict = STRICT_MODE if strict is None else strict
    message = f"[{context}] {exc}"
    if use_strict:
        raise RuntimeError(message) from exc
    warnings.warn(message)


@dataclass
class MethodCurves:
    correlations: Dict[str, pd.Series]
    rolling_std: Dict[str, pd.Series]
    stability_points: Dict[str, int | None]


def _find_latest_dir(base: Path, prefix: str) -> Path | None:
    candidates = [d for d in base.iterdir() if d.is_dir() and d.name.startswith(prefix)]
    if not candidates:
        return None
    return max(candidates, key=lambda p: p.stat().st_mtime)


def _find_latest_qa_dir(base: Path, wanted_strategy: str = "all_questions") -> Path | None:
    def _is_complete(d: Path, n_folds: int) -> bool:
        if not (d / "qa_level_test_outputs.csv").exists():
            return False
        for i in range(n_folds):
            if not (d / f"fold_{i}" / "qa_level_test_predictions.csv").exists():
                return False
        return True

    candidates = []
    for d in base.iterdir():
        if not d.is_dir():
            continue
        if not (d.name.startswith("qa_level_outputs_") or d.name.startswith("all_questions_qa_level_outputs_")):
            continue
        summary = d / "cross_validation_summary.json"
        if not summary.exists():
            summary = d / "overall_summary.json"
        if not summary.exists():
            summary = d / "config.json"
        if not summary.exists():
            continue
        try:
            with summary.open("r") as f:
                data = json.load(f)
        except Exception:
            continue
        if data.get("question_strategy") != wanted_strategy:
            continue
        try:
            n_folds = int(data.get("n_folds", 9))
        except Exception:
            continue
        if n_folds <= 0:
            continue
        if not _is_complete(d, n_folds):
            continue
        candidates.append(d)
    if not candidates:
        return None
    return max(candidates, key=lambda p: p.stat().st_mtime)


def _load_question_order() -> List[str]:
    questions_path = FORMATTED_DATA / "questions.csv"
    if questions_path.exists():
        try:
            df = pd.read_csv(questions_path)
            code_col = "code" if "code" in df.columns else df.columns[0]
            codes = df[code_col].astype(str).tolist()
            ordered = [code for code in codes if code not in EXCLUDED_CODES]
            if len(ordered) == len(DEFAULT_QUESTION_ORDER):
                return ordered
        except Exception as exc:
            _handle_exception(f"load question order from {questions_path}", exc)
    return DEFAULT_QUESTION_ORDER


def _prepare_irt_matrix(df: pd.DataFrame, question_order: List[str]) -> pd.DataFrame:
    if "user_id" not in df.columns:
        raise ValueError("Input data missing required column: user_id")

    if {"question_code", "discrete_score"}.issubset(df.columns):
        work = df.copy()
        work["question_code"] = work["question_code"].astype(str)
        wide = (
            work.pivot_table(
                index="user_id",
                columns="question_code",
                values="discrete_score",
                aggfunc="mean",
            )
            .reset_index()
        )
    else:
        wide = df.copy()

    response_cols = [col for col in wide.columns if col != "user_id"]
    missing = [code for code in question_order if code not in wide.columns]
    for code in missing:
        wide[code] = np.nan
    extra = [col for col in response_cols if col not in question_order]
    if extra:
        wide = wide.drop(columns=extra)
    wide = wide[["user_id"] + question_order]

    for col in question_order:
        wide[col] = pd.to_numeric(wide[col], errors="coerce")

    wide[question_order] = wide[question_order].fillna(0)
    wide[question_order] = np.rint(wide[question_order].to_numpy()).astype(int)
    return wide


def _active_question_order(input_dir: Path, fold_idx: int) -> List[str]:
    question_order = _load_question_order()
    dev_path = input_dir / f"dev_{fold_idx}.csv"
    if not dev_path.exists():
        return question_order

    dev_df = pd.read_csv(dev_path)
    dev_df = dev_df[[col for col in dev_df.columns if not col.startswith("pred_score_")]]
    dev_wide = _prepare_irt_matrix(dev_df, question_order)
    constant_cols = [
        col for col in question_order
        if dev_wide[col].nunique(dropna=False) <= 1
    ]
    if constant_cols:
        return [code for code in question_order if code not in constant_cols]
    return question_order


def _read_thetas_file(path: Path) -> Tuple[pd.DataFrame, int]:
    num_items = int(path.stem.split("_")[1])
    rows = []
    theta_dim = None
    with path.open(newline="") as csvfile:
        reader = csv.reader(csvfile)
        for row in reader:
            if not row:
                continue
            user_id = row[0]
            values = [float(x) for x in row[1:] if x != ""]
            if len(values) <= num_items:
                continue
            current_theta_dim = len(values) - num_items
            if theta_dim is None:
                theta_dim = current_theta_dim
            item_vals = values[current_theta_dim:]
            rows.append((user_id, item_vals))

    if theta_dim is None:
        return pd.DataFrame(), num_items

    data = {
        "user_id": [r[0] for r in rows],
        "num_items": [num_items] * len(rows),
        "items": [r[1] for r in rows],
    }
    return pd.DataFrame(data), num_items


def _load_all_items(run_dir: Path) -> pd.DataFrame:
    files = sorted(run_dir.glob("thetas_*.csv"), key=lambda p: int(p.stem.split("_")[1]))
    if not files:
        raise FileNotFoundError(f"No thetas_*.csv files found in {run_dir}")

    all_frames = []
    for file_path in files:
        df, _ = _read_thetas_file(file_path)
        if not df.empty:
            all_frames.append(df)

    if not all_frames:
        raise ValueError("No valid thetas data found.")

    return pd.concat(all_frames, ignore_index=True)


def _load_qa_outputs_by_fold(qa_dir: Path) -> Dict[int, pd.DataFrame]:
    fold_dirs = sorted([d for d in qa_dir.iterdir() if d.is_dir() and d.name.startswith("fold_")])
    if not fold_dirs:
        raise FileNotFoundError(f"No fold_* directories found in {qa_dir}")

    outputs = {}
    for fold_dir in fold_dirs:
        fold_idx = int(fold_dir.name.split("_")[1])
        path = fold_dir / "qa_level_test_predictions.csv"
        if not path.exists():
            continue
        df = pd.read_csv(path)
        df["question_code"] = df["question_code"].astype(str)
        df["user_id"] = df["user_id"].astype(str)
        outputs[fold_idx] = df

    if not outputs:
        raise FileNotFoundError("No qa_level_test_predictions.csv files found.")

    return outputs


def _merge_qa_outputs(qa_by_fold: Dict[int, pd.DataFrame]) -> pd.DataFrame:
    merged = pd.concat(list(qa_by_fold.values()), ignore_index=True)
    merged["question_code"] = merged["question_code"].astype(str)
    merged["user_id"] = merged["user_id"].astype(str)
    return merged


def _build_true_scores(qa_df: pd.DataFrame) -> pd.DataFrame:
    true_cols = [col for col in qa_df.columns if col.startswith("true_")]
    if true_cols:
        true_df = qa_df.groupby("user_id", as_index=False)[true_cols].mean()
        return true_df
    if "output_embeddings" in qa_df.columns:
        rows = []
        for user_id, user_data in qa_df.groupby("user_id"):
            vec = user_data["output_embeddings"].iloc[0]
            if isinstance(vec, str):
                try:
                    vec = ast.literal_eval(vec)
                except Exception as exc:
                    _handle_exception("parse output_embeddings literal", exc)
                    vec = None
            if isinstance(vec, (list, np.ndarray)) and len(vec) >= len(OUTCOMES):
                row = {"user_id": user_id}
                for idx, (key, _) in enumerate(OUTCOMES):
                    row[f"true_{key}"] = vec[idx]
                rows.append(row)
        return pd.DataFrame(rows)
    raise ValueError("QA outputs missing true_* columns or output_embeddings.")


def _build_question_mapping() -> Dict[str, set]:
    mapping_data = get_questionnaire_mapping()
    outcome_map = mapping_data["questionnaire_specific"]
    return {outcome: set(codes) for outcome, codes in outcome_map.items()}


def _aggregate_predictions(
    qa_df: pd.DataFrame,
    asked_questions: Dict[str, List[str]],
    outcome_map: Dict[str, set],
) -> pd.DataFrame:
    pred_cols = [col for col in qa_df.columns if col.startswith("pred_")]
    if not pred_cols:
        raise ValueError("QA outputs missing pred_* columns.")

    qa_df = qa_df.copy()
    qa_df["user_id"] = qa_df["user_id"].astype(str)
    asked_questions = {str(k): v for k, v in asked_questions.items()}

    user_groups = qa_df.groupby("user_id")
    rows = []
    for user_id, user_data in user_groups:
        if user_id not in asked_questions:
            continue
        asked = set(asked_questions[user_id])
        asked_data = user_data[user_data["question_code"].isin(asked)]
        if asked_data.empty:
            continue

        row = {"user_id": user_id}
        for key, _ in OUTCOMES:
            pred_col = f"pred_{key}"
            if pred_col not in asked_data.columns:
                row[pred_col] = np.nan
                continue
            if key in outcome_map:
                target = asked_data[asked_data["question_code"].isin(outcome_map[key])]
                if not target.empty:
                    row[pred_col] = target[pred_col].mean()
                else:
                    row[pred_col] = asked_data[pred_col].mean()
            else:
                row[pred_col] = asked_data[pred_col].mean()
        rows.append(row)

    if not rows:
        columns = ["user_id"] + [f"pred_{key}" for key, _ in OUTCOMES]
        return pd.DataFrame(columns=columns)

    return pd.DataFrame(rows)


def _compute_correlations(
    pred_df: pd.DataFrame,
    true_df: pd.DataFrame,
) -> Dict[str, float]:
    merged = pred_df.merge(true_df, on="user_id", how="inner")
    results = {}
    for key, _ in OUTCOMES:
        pred_col = f"pred_{key}"
        true_col = f"true_{key}"
        if pred_col not in merged.columns or true_col not in merged.columns:
            results[key] = np.nan
            continue
        x = merged[pred_col].to_numpy()
        y = merged[true_col].to_numpy()
        if len(x) < 2 or np.std(x) == 0 or np.std(y) == 0:
            results[key] = np.nan
            continue
        results[key] = float(pearsonr(x, y)[0])
    return results


def _find_stability_point(num_items: pd.Series, rolling_std: pd.Series, threshold: float) -> int | None:
    for idx in range(len(num_items)):
        tail = rolling_std.iloc[idx:]
        tail = tail[~tail.isna()]
        if tail.empty:
            continue
        if (tail <= threshold).all():
            return int(num_items.iloc[idx])
    return None


def _load_gpt4_csv(path: Path) -> pd.DataFrame:
    df = pd.read_csv(path)
    required = {"user_id", "num_items"}
    if not required.issubset(df.columns):
        raise ValueError("GPT-4 CSV must include user_id and num_items columns.")
    return df


def _plot_outcome(
    outcome_display: str,
    methods: Dict[str, Tuple[MethodCurves, pd.Series]],
    threshold: float,
    output_path: Path,
) -> None:
    fig, axes = plt.subplots(2, 1, figsize=(8, 6), sharex=True)
    ax_corr, ax_std = axes

    colors = {"MAQuA": "#1f77b4", "Random IRT": "#ff7f0e", "GPT-4": "#2ca02c"}

    for method_name, (curves, series) in methods.items():
        corr = curves.correlations[outcome_display]
        std = curves.rolling_std[outcome_display]
        ax_corr.plot(series, corr, label=method_name, color=colors.get(method_name))
        ax_std.plot(series, std, label=method_name, linestyle="--", color=colors.get(method_name))

        stability = curves.stability_points[outcome_display]
        if stability is not None:
            ax_corr.axvline(stability, color=colors.get(method_name), alpha=0.4)
            ax_std.axvline(stability, color=colors.get(method_name), alpha=0.4)

    ax_corr.set_ylabel("Correlation")
    ax_corr.set_title(outcome_display)
    ax_corr.legend(loc="lower right")

    ax_std.set_ylabel("Rolling Std. Dev.")
    ax_std.set_xlabel("Question Number")
    ax_std.axhline(threshold, color="gray", linestyle=":", linewidth=1)
    ax_std.legend(loc="upper right")

    plt.tight_layout()
    plt.savefig(output_path, dpi=200)
    plt.close(fig)


def _compute_method_curves(
    correlations: pd.DataFrame,
    num_items: pd.Series,
    window: int,
    threshold: float,
) -> MethodCurves:
    outcome_series = {}
    rolling_std = {}
    stability_points = {}

    for _, display in OUTCOMES:
        series = correlations[display]
        outcome_series[display] = series
        rolling_std[display] = series.rolling(window, min_periods=window).std()
        stability_points[display] = _find_stability_point(num_items, rolling_std[display], threshold)

    return MethodCurves(outcome_series, rolling_std, stability_points)


def _add_average(curves: MethodCurves, num_items: pd.Series, window: int, threshold: float) -> None:
    all_series = pd.concat(curves.correlations.values(), axis=1)
    avg = all_series.mean(axis=1)
    curves.correlations["Average"] = avg
    curves.rolling_std["Average"] = avg.rolling(window, min_periods=window).std()
    curves.stability_points["Average"] = _find_stability_point(num_items, curves.rolling_std["Average"], threshold)


def main():
    global STRICT_MODE
    global ITEM_LIST_COERCE_ERROR_COUNT, ITEM_INDEX_PARSE_ERROR_COUNT
    parser = argparse.ArgumentParser(description="Plot Figure 2 adaptive correlations.")
    parser.add_argument("--adaptive-dir", type=str, default=None, help="Adaptive output dir (or parent).")
    parser.add_argument("--input-dir", type=str, default=None, help="Discretized input dir for question mapping.")
    parser.add_argument("--qa-dir", type=str, default=None, help="QA-level outputs dir.")
    parser.add_argument("--gpt4-csv", type=str, default=None, help="Optional GPT-4 baseline CSV.")
    parser.add_argument("--rolling-window", type=int, default=5, help="Rolling window size.")
    parser.add_argument("--std-threshold", type=float, default=0.01, help="Std threshold for stability.")
    parser.add_argument("--out-dir", type=str, default=None, help="Output directory.")
    parser.add_argument(
        "--strict",
        action="store_true",
        help="Fail fast on parse/summary errors instead of warning and continuing.",
    )
    args = parser.parse_args()
    STRICT_MODE = args.strict

    adaptive_dir = Path(args.adaptive_dir) if args.adaptive_dir else _find_latest_dir(ADAPTIVE_OUTPUTS, "")
    if adaptive_dir is None:
        raise FileNotFoundError("Adaptive output directory not found.")

    input_dir = Path(args.input_dir) if args.input_dir else _find_latest_dir(POLYTOMIZED_DATA, "")
    if input_dir is None:
        raise FileNotFoundError("Polytomized input directory not found.")

    qa_dir = Path(args.qa_dir) if args.qa_dir else _find_latest_qa_dir(MODEL_OUTPUTS, "all_questions")
    if qa_dir is None:
        raise FileNotFoundError("QA output directory not found.")

    output_dir = Path(args.out_dir) if args.out_dir else ANALYSIS_OUTPUTS / "figure2"
    output_dir.mkdir(parents=True, exist_ok=True)

    qa_by_fold = _load_qa_outputs_by_fold(qa_dir)
    outcome_map = _build_question_mapping()

    run_dirs = {}
    if (adaptive_dir / "thetas_1.csv").exists():
        run_dirs["MAQuA"] = adaptive_dir
    else:
        for name in ("Drule", "random"):
            candidate = adaptive_dir / name
            if candidate.exists():
                label = "MAQuA" if name == "Drule" else "Random IRT"
                run_dirs[label] = candidate

    method_results = {}

    for label, run_dir in run_dirs.items():
        fold_dirs = [d for d in run_dir.iterdir() if d.is_dir() and d.name.startswith("fold_")]
        if not fold_dirs and (run_dir / "thetas_1.csv").exists():
            # Legacy flat output: compute correlations on merged QA data.
            qa_df = _merge_qa_outputs(qa_by_fold)
            true_df = _build_true_scores(qa_df)
            items_df = _load_all_items(run_dir)
            num_items = sorted(items_df["num_items"].unique())

            correlations = defaultdict(list)
            for n in num_items:
                subset = items_df[items_df["num_items"] == n]
                asked_questions = {}
                fold_order = _active_question_order(input_dir, 0)
                for _, row in subset.iterrows():
                    user_id = row["user_id"]
                    items = row["items"]
                    if not isinstance(items, list):
                        try:
                            items = list(items)
                        except Exception as exc:
                            if STRICT_MODE:
                                _handle_exception(f"coerce adaptive item list for user {user_id}", exc, strict=True)
                            ITEM_LIST_COERCE_ERROR_COUNT += 1
                            items = []
                    question_codes = []
                    for item_idx in items:
                        try:
                            idx = int(item_idx)
                        except Exception as exc:
                            if STRICT_MODE:
                                _handle_exception(
                                    f"parse adaptive item index '{item_idx}' for user {user_id}",
                                    exc,
                                    strict=True,
                                )
                            ITEM_INDEX_PARSE_ERROR_COUNT += 1
                            continue
                        if 1 <= idx <= len(fold_order):
                            question_codes.append(fold_order[idx - 1])
                    asked_questions[user_id] = question_codes

                pred_df = _aggregate_predictions(qa_df, asked_questions, outcome_map)
                corr = _compute_correlations(pred_df, true_df)
                for key, display in OUTCOMES:
                    correlations[display].append(corr.get(key, np.nan))

            corr_df = pd.DataFrame({"num_items": num_items})
            for _, display in OUTCOMES:
                corr_df[display] = correlations[display]

            curves = _compute_method_curves(
                corr_df.drop(columns=["num_items"]),
                corr_df["num_items"],
                window=args.rolling_window,
                threshold=args.std_threshold,
            )
            _add_average(curves, corr_df["num_items"], args.rolling_window, args.std_threshold)
            method_results[label] = (curves, corr_df["num_items"])
            continue

        fold_corrs = defaultdict(lambda: defaultdict(list))
        all_num_items = set()

        for fold_idx, qa_df in qa_by_fold.items():
            fold_run_dir = run_dir / f"fold_{fold_idx}"
            if not fold_run_dir.exists():
                continue

            true_df = _build_true_scores(qa_df)
            items_df = _load_all_items(fold_run_dir)
            fold_order = _active_question_order(input_dir, fold_idx)

            for n in sorted(items_df["num_items"].unique()):
                subset = items_df[items_df["num_items"] == n]
                asked_questions = {}
                for _, row in subset.iterrows():
                    user_id = row["user_id"]
                    items = row["items"]
                    if not isinstance(items, list):
                        try:
                            items = list(items)
                        except Exception as exc:
                            if STRICT_MODE:
                                _handle_exception(
                                    f"coerce fold item list for user {user_id} fold {fold_idx}",
                                    exc,
                                    strict=True,
                                )
                            ITEM_LIST_COERCE_ERROR_COUNT += 1
                            items = []
                    question_codes = []
                    for item_idx in items:
                        try:
                            idx = int(item_idx)
                        except Exception as exc:
                            if STRICT_MODE:
                                _handle_exception(
                                    f"parse fold item index '{item_idx}' for user {user_id} fold {fold_idx}",
                                    exc,
                                    strict=True,
                                )
                            ITEM_INDEX_PARSE_ERROR_COUNT += 1
                            continue
                        if 1 <= idx <= len(fold_order):
                            question_codes.append(fold_order[idx - 1])
                    asked_questions[user_id] = question_codes

                pred_df = _aggregate_predictions(qa_df, asked_questions, outcome_map)
                corr = _compute_correlations(pred_df, true_df)
                all_num_items.add(n)
                for key, display in OUTCOMES:
                    fold_corrs[display][n].append(corr.get(key, np.nan))

        num_items = sorted(all_num_items)
        corr_df = pd.DataFrame({"num_items": num_items})
        for _, display in OUTCOMES:
            corr_values = []
            for n in num_items:
                vals = [v for v in fold_corrs[display].get(n, []) if not np.isnan(v)]
                corr_values.append(np.nan if not vals else float(np.mean(vals)))
            corr_df[display] = corr_values

        curves = _compute_method_curves(
            corr_df.drop(columns=["num_items"]),
            corr_df["num_items"],
            window=args.rolling_window,
            threshold=args.std_threshold,
        )
        _add_average(curves, corr_df["num_items"], args.rolling_window, args.std_threshold)
        method_results[label] = (curves, corr_df["num_items"])

    if args.gpt4_csv:
        gpt_df = _load_gpt4_csv(Path(args.gpt4_csv))
        true_df = _build_true_scores(_merge_qa_outputs(qa_by_fold))
        num_items = sorted(gpt_df["num_items"].unique())
        corr_rows = {display: [] for _, display in OUTCOMES}
        for n in num_items:
            subset = gpt_df[gpt_df["num_items"] == n]
            merged = subset.merge(true_df, on="user_id", how="inner")
            for key, display in OUTCOMES:
                pred_col = display if display in merged.columns else key
                true_col = f"true_{key}"
                if pred_col not in merged.columns or true_col not in merged.columns:
                    corr_rows[display].append(np.nan)
                    continue
                x = merged[pred_col].to_numpy()
                y = merged[true_col].to_numpy()
                if len(x) < 2 or np.std(x) == 0 or np.std(y) == 0:
                    corr_rows[display].append(np.nan)
                else:
                    corr_rows[display].append(float(pearsonr(x, y)[0]))
        corr_df = pd.DataFrame({"num_items": num_items})
        for _, display in OUTCOMES:
            corr_df[display] = corr_rows[display]

        curves = _compute_method_curves(
            corr_df.drop(columns=["num_items"]),
            corr_df["num_items"],
            window=args.rolling_window,
            threshold=args.std_threshold,
        )
        _add_average(curves, corr_df["num_items"], args.rolling_window, args.std_threshold)
        method_results["GPT-4"] = (curves, corr_df["num_items"])

    outcomes_to_plot = [display for _, display in OUTCOMES] + ["Average"]
    summary_rows = []
    for outcome in outcomes_to_plot:
        methods = {}
        for method_name, (curves, series) in method_results.items():
            methods[method_name] = (curves, series)

            for idx, n in enumerate(series):
                summary_rows.append(
                    {
                        "method": method_name,
                        "outcome": outcome,
                        "num_items": int(n),
                        "correlation": float(curves.correlations[outcome].iloc[idx]),
                        "rolling_std": float(curves.rolling_std[outcome].iloc[idx])
                        if not np.isnan(curves.rolling_std[outcome].iloc[idx])
                        else np.nan,
                    }
                )

        if not methods:
            continue

        output_path = output_dir / f"figure2_{outcome.replace(' ', '_')}.png"
        _plot_outcome(outcome, methods, args.std_threshold, output_path)

    summary_df = pd.DataFrame(summary_rows)
    summary_df.to_csv(output_dir / "figure2_summary.csv", index=False)

    stability_rows = []
    for method_name, (curves, series) in method_results.items():
        for outcome, stability in curves.stability_points.items():
            if stability is None:
                reduction = np.nan
            else:
                reduction = (1.0 - stability / series.max()) * 100
            stability_rows.append(
                {
                    "method": method_name,
                    "outcome": outcome,
                    "stability_item": stability,
                    "total_items": int(series.max()),
                    "reduction_pct": reduction,
                }
            )

    pd.DataFrame(stability_rows).to_csv(output_dir / "stabilization_points.csv", index=False)
    if ITEM_LIST_COERCE_ERROR_COUNT:
        warnings.warn(
            f"Failed to coerce adaptive item lists {ITEM_LIST_COERCE_ERROR_COUNT} times; "
            "set --strict to fail on first occurrence."
        )
    if ITEM_INDEX_PARSE_ERROR_COUNT:
        warnings.warn(
            f"Failed to parse adaptive item indices {ITEM_INDEX_PARSE_ERROR_COUNT} times; "
            "set --strict to fail on first occurrence."
        )
    print(f"Saved Figure 2 outputs to: {output_dir}")


if __name__ == "__main__":
    main()