plot_adaptive_results.py

#!/usr/bin/env python3
"""
Plot adaptive testing outputs (thetas + item selection frequency).

This script reads adaptive_outputs/<discretized>/<order>/thetas_*.csv
and generates summary plots and a CSV with convergence statistics.
It can also generate Figure 2-style plots (correlations + rolling std).
"""

from __future__ import annotations

import argparse
import csv
from collections import Counter
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 maqua_paths import ANALYSIS_OUTPUTS, FORMATTED_DATA


STRICT_MODE = False


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)


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"}


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 | None) -> List[str]:
    question_order = _load_question_order()
    if input_dir is None:
        return question_order

    dev_path = input_dir / "dev_0.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
            theta_vals = values[:current_theta_dim]
            item_vals = values[current_theta_dim:]
            rows.append((user_id, theta_vals, 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),
    }
    for i in range(theta_dim):
        data[f"theta_{i+1}"] = [r[1][i] for r in rows]
    data["items"] = [r[2] for r in rows]
    return pd.DataFrame(data), num_items


def _load_all_thetas(run_dir: Path) -> Tuple[pd.DataFrame, int]:
    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 = []
    max_items = 0
    for file_path in files:
        df, num_items = _read_thetas_file(file_path)
        if not df.empty:
            all_frames.append(df)
            max_items = max(max_items, num_items)

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

    combined = pd.concat(all_frames, ignore_index=True)
    return combined, max_items


def _summarize_thetas(thetas_df: pd.DataFrame) -> Tuple[pd.DataFrame, List[str], int]:
    theta_cols = [col for col in thetas_df.columns if col.startswith("theta_")]
    max_items = int(thetas_df["num_items"].max())
    final_df = thetas_df[thetas_df["num_items"] == max_items]
    final_df = final_df.groupby("user_id", as_index=False)[theta_cols].mean()

    summary_rows = []
    for num_items in sorted(thetas_df["num_items"].unique()):
        current = thetas_df[thetas_df["num_items"] == num_items]
        current = current.groupby("user_id", as_index=False)[theta_cols].mean()
        merged = current.merge(final_df, on="user_id", suffixes=("", "_final"))

        row = {"num_items": num_items}
        for col in theta_cols:
            diff = (merged[col] - merged[f"{col}_final"]).abs()
            row[f"{col}_mae"] = diff.mean()

            x = merged[col].to_numpy()
            y = merged[f"{col}_final"].to_numpy()
            if x.std() == 0 or y.std() == 0:
                row[f"{col}_corr"] = np.nan
            else:
                row[f"{col}_corr"] = np.corrcoef(x, y)[0, 1]
        summary_rows.append(row)

    summary_df = pd.DataFrame(summary_rows).sort_values("num_items")
    return summary_df, theta_cols, max_items


def _plot_theta_convergence(summary_df: pd.DataFrame, theta_cols: List[str], out_path: Path) -> None:
    plt.figure(figsize=(8, 5))
    for col in theta_cols:
        plt.plot(summary_df["num_items"], summary_df[f"{col}_mae"], label=f"{col} MAE")
    plt.xlabel("Number of items")
    plt.ylabel("Mean |theta_k - theta_final|")
    plt.title("Theta convergence to full-length estimate")
    plt.legend()
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close()


def _plot_theta_correlation(summary_df: pd.DataFrame, theta_cols: List[str], out_path: Path) -> None:
    plt.figure(figsize=(8, 5))
    for col in theta_cols:
        plt.plot(summary_df["num_items"], summary_df[f"{col}_corr"], label=f"{col} corr")
    plt.xlabel("Number of items")
    plt.ylabel("Correlation with full-length theta")
    plt.title("Theta correlation vs. full-length estimate")
    plt.legend()
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close()


def _plot_item_frequency(
    item_counts: Counter,
    item_labels: Dict[int, str] | None,
    out_path: Path,
    top_k: int = 20,
) -> None:
    if not item_counts:
        return

    most_common = item_counts.most_common(top_k)
    item_ids = [item for item, _ in most_common]
    counts = [count for _, count in most_common]

    labels = []
    for item_id in item_ids:
        if item_labels and item_id in item_labels:
            labels.append(item_labels[item_id])
        else:
            labels.append(str(item_id))

    plt.figure(figsize=(10, 6))
    plt.bar(range(len(labels)), counts)
    plt.xticks(range(len(labels)), labels, rotation=45, ha="right")
    plt.xlabel("Item")
    plt.ylabel("Selection count")
    plt.title(f"Top {top_k} selected items (all users)")
    plt.tight_layout()
    plt.savefig(out_path, dpi=200)
    plt.close()


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 _plot_figure2(
    summaries: Dict[str, pd.DataFrame],
    theta_cols: List[str],
    output_dir: Path,
    window: int,
    threshold: float,
    include_mean: bool,
) -> pd.DataFrame:
    colors = {"Drule": "#1f77b4", "random": "#ff7f0e"}
    output_dir.mkdir(parents=True, exist_ok=True)
    stability_rows = []

    plot_cols = list(theta_cols)
    if include_mean and "theta_mean" not in plot_cols:
        plot_cols.append("theta_mean")

    for theta_col in plot_cols:
        fig, axes = plt.subplots(2, 1, figsize=(8, 6), sharex=True)
        ax_corr, ax_std = axes

        for run_name, summary_df in summaries.items():
            if theta_col == "theta_mean":
                available = [col for col in theta_cols if col in summary_df]
                if not available:
                    continue
                corr = summary_df[available].mean(axis=1)
            else:
                if theta_col not in summary_df:
                    continue
                corr = summary_df[theta_col]

            num_items = summary_df["num_items"]
            rolling_std = corr.rolling(window, min_periods=window).std()

            ax_corr.plot(num_items, corr, label=f"{run_name} corr", color=colors.get(run_name, None))
            ax_std.plot(num_items, rolling_std, label=f"{run_name} std", linestyle="--", color=colors.get(run_name, None))

            stability_point = _find_stability_point(num_items, rolling_std, threshold)
            if stability_point is not None:
                ax_corr.axvline(stability_point, color=colors.get(run_name, None), alpha=0.4)
                ax_std.axvline(stability_point, color=colors.get(run_name, None), alpha=0.4)
                ax_corr.text(
                    stability_point,
                    corr.max(),
                    f"{stability_point}",
                    rotation=90,
                    va="bottom",
                    ha="right",
                    fontsize=8,
                    color=colors.get(run_name, None),
                )

            stability_rows.append(
                {
                    "run": run_name,
                    "theta": theta_col,
                    "stability_item": stability_point,
                    "total_items": int(num_items.max()),
                }
            )

        ax_corr.set_ylabel("Pearson r vs full-length")
        ax_corr.set_title(f"{theta_col}: correlation and rolling std")
        ax_corr.legend(loc="lower right")

        ax_std.set_ylabel("Rolling std of r")
        ax_std.set_xlabel("Number of items")
        ax_std.axhline(threshold, color="gray", linestyle=":", linewidth=1)
        ax_std.legend(loc="upper right")

        plt.tight_layout()
        out_path = output_dir / f"figure2_{theta_col}.png"
        plt.savefig(out_path, dpi=200)
        plt.close(fig)

    stability_df = pd.DataFrame(stability_rows)
    return stability_df


def analyze_run(run_dir: Path, input_dir: Path | None, output_dir: Path) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    thetas_df, _ = _load_all_thetas(run_dir)
    summary_df, theta_cols, _ = _summarize_thetas(thetas_df)
    summary_df.to_csv(output_dir / "theta_convergence_summary.csv", index=False)

    _plot_theta_convergence(summary_df, theta_cols, output_dir / "theta_convergence.png")
    _plot_theta_correlation(summary_df, theta_cols, output_dir / "theta_correlation.png")

    item_counts = Counter()
    invalid_item_count = 0
    for row_items in thetas_df["items"]:
        for item in row_items:
            try:
                item_counts[int(item)] += 1
            except Exception as exc:
                invalid_item_count += 1
                if STRICT_MODE:
                    _handle_exception(f"parse selected item id in {run_dir}", exc)
                continue

    if invalid_item_count:
        warnings.warn(
            f"[{run_dir}] skipped {invalid_item_count} invalid item ids while summarizing item frequency"
        )

    item_labels = None
    if input_dir is not None:
        active_order = _active_question_order(input_dir)
        item_labels = {idx + 1: code for idx, code in enumerate(active_order)}

    _plot_item_frequency(item_counts, item_labels, output_dir / "item_frequency_top20.png")


def main():
    global STRICT_MODE
    parser = argparse.ArgumentParser(description="Plot adaptive testing outputs.")
    parser.add_argument("--adaptive-dir", type=str, required=True, help="Adaptive output dir or its parent.")
    parser.add_argument("--input-dir", type=str, default=None, help="Discretized input dir (for item labels).")
    parser.add_argument("--out-dir", type=str, default=None, help="Output directory for plots.")
    parser.add_argument("--figure2", action="store_true", help="Generate Figure 2-style correlation/std plots.")
    parser.add_argument("--rolling-window", type=int, default=5, help="Rolling window size for std.")
    parser.add_argument("--std-threshold", type=float, default=0.01, help="Std threshold for stability line.")
    parser.add_argument("--include-mean", action="store_true", help="Include mean correlation across thetas.")
    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)
    input_dir = Path(args.input_dir) if args.input_dir else None

    run_dirs = []
    if (adaptive_dir / "thetas_1.csv").exists():
        run_dirs = [adaptive_dir]
    else:
        for name in ("random", "Drule"):
            candidate = adaptive_dir / name
            if candidate.exists():
                run_dirs.append(candidate)

    if not run_dirs:
        raise FileNotFoundError("No adaptive testing output directories found.")

    summaries = {}
    for run_dir in run_dirs:
        fold_dirs = [d for d in run_dir.iterdir() if d.is_dir() and d.name.startswith("fold_")]
        if fold_dirs:
            for fold_dir in sorted(fold_dirs):
                if args.out_dir:
                    out_dir = Path(args.out_dir) / run_dir.name / fold_dir.name
                else:
                    out_dir = ANALYSIS_OUTPUTS / "adaptive_plots" / adaptive_dir.name / run_dir.name / fold_dir.name
                analyze_run(fold_dir, input_dir, out_dir)
                print(f"Saved plots to: {out_dir}")
            continue

        if args.out_dir:
            out_dir = Path(args.out_dir) / run_dir.name
        else:
            out_dir = ANALYSIS_OUTPUTS / "adaptive_plots" / adaptive_dir.name / run_dir.name
        analyze_run(run_dir, input_dir, out_dir)
        try:
            thetas_df, _ = _load_all_thetas(run_dir)
            summary_df, theta_cols, _ = _summarize_thetas(thetas_df)
            corr_cols = {f"{theta}_corr": theta for theta in theta_cols}
            corr_df = summary_df[["num_items"] + list(corr_cols.keys())].rename(columns=corr_cols)
            summaries[run_dir.name] = corr_df
        except Exception as exc:
            _handle_exception(f"build figure2 summary for {run_dir}", exc)
        print(f"Saved plots to: {out_dir}")

    if args.figure2 and summaries:
        theta_cols = [
            col for col in next(iter(summaries.values())).columns if col.startswith("theta_")
        ]
        base_out = Path(args.out_dir) if args.out_dir else ANALYSIS_OUTPUTS / "adaptive_plots" / adaptive_dir.name
        stability_df = _plot_figure2(
            summaries,
            theta_cols,
            base_out / "figure2",
            window=args.rolling_window,
            threshold=args.std_threshold,
            include_mean=args.include_mean,
        )
        stability_df["reduction_pct"] = (
            1.0 - stability_df["stability_item"] / stability_df["total_items"]
        ) * 100
        stability_df.to_csv(base_out / "figure2" / "stabilization_points.csv", index=False)
        print(f"Saved Figure 2-style plots to: {base_out / 'figure2'}")


if __name__ == "__main__":
    main()