#!/usr/bin/env python3
"""
Script to plot dummy evaluation results with aggregation across experiments.

Handles directory structure like:
    dummy04_coinrun_exp0/
    dummy04_coinrun_exp1/
    dummy04_maze_exp0/
    dummy04_robust400_maze_exp0/
    ...

Aggregates results for the same method across different experiment IDs,
plotting mean with shaded standard error regions.
Robust eval groups are plotted as distinct method variants.
"""

from __future__ import annotations

import argparse
import re
import textwrap
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Optional, Tuple

# Width at which titles get wrapped onto multiple lines. ~80 characters fits
# the default 8" figure width at the default title font size.
TITLE_WRAP_WIDTH = 80

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns  # type: ignore
from matplotlib.lines import Line2D
from scipy import interpolate, integrate

# Methods whose curves should be rendered in the same blue used for the
# Expert reference line. Keyed by canonical base method (no robust suffix).
SPECIAL_METHOD_COLORS = {
    "oracle-lb-random": "blue",
}
EXPERT_REFERENCE_COLOR = "blue"
NOVICE_REFERENCE_COLOR = "red"

from analyzing.utils import extract_x_and_y_values
from analyzing.plotting_common import (
    METHOD_NAMES,
    setup_plot_style,
    get_line_styles,
    style_plot_for_publication,
)

matplotlib.use("TkAgg")

SUPPORTED_ENVS = [
    "coinrun",
    "coinrun_proxy_fail",
    "maze",
    "maze_afh",
    "maze_proxy_fail",
    "heist",
]
ENV_PATTERN = "|".join(
    re.escape(env) for env in sorted(SUPPORTED_ENVS, key=len, reverse=True)
)
ROBUST_VARIANTS = ("robust200", "robust400")
ROBUST_LABELS = {
    "robust200": "Robust 200M",
    "robust400": "Robust 400M",
}
ROBUST_PATTERN = "|".join(re.escape(variant) for variant in ROBUST_VARIANTS)

# Data key display names mapping
DATA_KEY_NAMES = {
    "step_afhp": "Ask-For-Help Percentage (AFHP, per timestep)",
    "level_afhp": "Ask-For-Help Percentage (AFHP)",
    "performance": "Mean Return",
    "performance_asked": "Mean Reward (Asked for Help)",
    "performance_not_asked": "Mean Reward (Did Not Ask)",
    "performance_asked_correctly": "Mean Reward (True Positive)",
    "performance_not_asked_correctly": "Mean Reward (True Negative)",
    "performance_not_asked_incorrectly": "Mean Reward (False Negative)",
    "ood_accuracy": "OOD Accuracy",
    "true_positive": "True Positive Rate",
    "false_positive": "False Positive Rate",
    "true_negative": "True Negative Rate",
    "false_negative": "False Negative Rate",
    "episode_length_mean": "Mean Episode Length",
    "episode_length_success_mean": "Mean Episode Length (Success)",
    "first_ood_timestep_mean": "Mean First OOD Timestep",
}

# Keys that filter by asking behavior (need special handling for reference lines)
FILTERED_PERFORMANCE_KEYS = {
    "performance_asked",
    "performance_not_asked",
    "performance_asked_correctly",
    "performance_not_asked_correctly",
    "performance_not_asked_incorrectly",
}


def parse_robust_experiment_dir(
    dir_name: str,
) -> Optional[Tuple[str, str, str, int]]:
    """
    Parse robust experiment directory names.

    Expected format: {prefix}_robust{200,400}_{env}_exp{id}
    Example: dummy04_robust400_maze_exp0

    Returns:
        Tuple of (base_prefix, robust_variant, env, exp_id), or None.
    """
    pattern = rf"^(.+)_({ROBUST_PATTERN})_({ENV_PATTERN})_exp(\d+)$"
    match = re.match(pattern, dir_name)
    if match:
        prefix = match.group(1)
        robust_variant = match.group(2)
        env = match.group(3)
        exp_id = int(match.group(4))
        return prefix, robust_variant, env, exp_id
    return None


def parse_experiment_dir(dir_name: str) -> Optional[Tuple[str, str, int]]:
    """
    Parse experiment directory name to extract prefix, env, and experiment ID.

    Expected format: {prefix}_{env}_exp{id}
    Examples: dummy04_coinrun_exp0, dummy04_maze_proxy_fail_exp1

    Returns:
        Tuple of (prefix, env, exp_id) or None if pattern doesn't match
    """
    robust_parsed = parse_robust_experiment_dir(dir_name)
    if robust_parsed is not None:
        prefix, _, env, exp_id = robust_parsed
        return prefix, env, exp_id

    pattern = rf"^(.+)_({ENV_PATTERN})_exp(\d+)$"
    match = re.match(pattern, dir_name)
    if match:
        prefix = match.group(1)
        env = match.group(2)
        exp_id = int(match.group(3))
        return prefix, env, exp_id
    return None


def split_robust_method(method: str) -> Tuple[str, Optional[str]]:
    """Split method names like max-prob_robust400 into base method and variant."""
    for robust_variant in ROBUST_VARIANTS:
        suffix = f"_{robust_variant}"
        if method.endswith(suffix):
            return method[: -len(suffix)], robust_variant
    return method, None


def canonicalize_method(method: str) -> str:
    """Normalize method names to kebab-case (matching run_eval.py).

    Robust suffixes like `_robust400` are preserved verbatim; only the
    base-method portion is hyphenated.
    """
    base, robust_variant = split_robust_method(method)
    base = base.replace("_", "-")
    return add_robust_suffix(base, robust_variant)


def method_has_robust_variant(method: str) -> bool:
    """True if the method name carries one of the known robust suffixes."""
    _, robust_variant = split_robust_method(method)
    return robust_variant is not None


_PARTIAL_ORACLE_LABEL_MARKER = "PartialOracle"
_NOVICE_LABEL = r"\textsc{Novice}"
_EXPERT_LABEL = r"\textsc{Expert}"
_RANDOM_LABEL = r"\textsc{Random}"

# Canonical (kebab-case) base method names, in the order they should appear
# above the PartialOracle divider when the user does not pass `--method_order`.
# Any methods present in `results` that aren't listed here are appended at the
# end in alphabetical order. Robust variants follow their base method.
DEFAULT_METHOD_ORDER = [
    "ts-random",       # Heuristic
    "ensemble-single",  # Ensemble (single weak)
    "max-prob",
    "max-logit",
    "svdd-image",
    "svdd-latent",
]


def _build_legend_entries(ax):
    """Reorder the auto-discovered legend so PartialOracle sits just above the
    Novice/Expert/Random reference block, with a thin gray rule above it.

    Returns:
        (handles, labels, separator_indices) — the third element lists the row
        positions in `handles` where the marker-column placeholder should be
        expanded into a full-legend-width horizontal rule.
    """
    raw_handles, raw_labels = ax.get_legend_handles_labels()

    method_entries = []
    partial_oracle_entries = []
    reference_entries: Dict[str, Tuple[object, str]] = {}

    for handle, label in zip(raw_handles, raw_labels):
        if _PARTIAL_ORACLE_LABEL_MARKER in label:
            partial_oracle_entries.append((handle, label))
        elif label == _NOVICE_LABEL:
            reference_entries["novice"] = (handle, label)
        elif label == _EXPERT_LABEL:
            reference_entries["expert"] = (handle, label)
        elif label == _RANDOM_LABEL:
            reference_entries["random"] = (handle, label)
        else:
            method_entries.append((handle, label))

    ordered: List[Tuple[object, str]] = list(method_entries)
    separator_indices: List[int] = []

    if partial_oracle_entries:
        separator = Line2D([], [], color="gray", linewidth=0.7)
        ordered.append((separator, ""))
        separator_indices.append(len(ordered) - 1)
        ordered.extend(partial_oracle_entries)

    # Reference lines always come last in this fixed order.
    for key in ("novice", "expert", "random"):
        if key in reference_entries:
            ordered.append(reference_entries[key])

    handles = [h for h, _ in ordered]
    labels = [l for _, l in ordered]
    return handles, labels, separator_indices


def _draw_legend_separator(legend, indices):
    """Replace the marker-column placeholders at `indices` in `legend` with
    thin horizontal rules that span the full legend width.

    Computed in figure-fraction coordinates from the legend's bbox after a
    forced draw. Suitable for batch save workflows; the rules do not update
    if the figure is later resized interactively.
    """
    if not indices:
        return

    fig = legend.figure
    if fig is None:
        return

    fig.canvas.draw()  # populate legend / row bboxes

    handles_attr = getattr(legend, "legend_handles", None)
    if handles_attr is None:
        handles_attr = getattr(legend, "legendHandles", None)
    if not handles_attr:
        return

    legend_bbox = legend.get_window_extent()
    if legend_bbox.width <= 0 or legend_bbox.height <= 0:
        return
    inv = fig.transFigure.inverted()
    pad_disp = 4  # display-coord inset from each legend frame edge

    for idx in indices:
        if idx >= len(handles_attr):
            continue
        placeholder = handles_attr[idx]
        placeholder_bbox = placeholder.get_window_extent()
        y_mid_disp = (placeholder_bbox.y0 + placeholder_bbox.y1) / 2
        x0_disp = legend_bbox.x0 + pad_disp
        x1_disp = legend_bbox.x1 - pad_disp
        p0 = inv.transform((x0_disp, y_mid_disp))
        p1 = inv.transform((x1_disp, y_mid_disp))
        rule = Line2D(
            [p0[0], p1[0]],
            [p0[1], p1[1]],
            color="gray",
            linewidth=0.7,
            transform=fig.transFigure,
            zorder=legend.get_zorder() + 1,
        )
        fig.add_artist(rule)
        # Hide the marker-column placeholder so we don't render two lines on
        # the same row.
        placeholder.set_visible(False)


def filter_results_by_robust(
    results: Dict[str, Dict[int, Path]], robust_filter: str
) -> Dict[str, Dict[int, Path]]:
    """Filter the results dict by robust-variant membership.

    robust_filter:
        "all" -> identity
        "robust" -> keep only methods with a robust suffix
        "non-robust" -> keep only methods without a robust suffix
    """
    if robust_filter == "all":
        return results
    if robust_filter == "robust":
        return {m: d for m, d in results.items() if method_has_robust_variant(m)}
    if robust_filter == "non-robust":
        return {m: d for m, d in results.items() if not method_has_robust_variant(m)}
    raise ValueError(f"Unknown robust_filter: {robust_filter!r}")


def add_robust_suffix(method: str, robust_variant: Optional[str]) -> str:
    """Attach robust variant to method name unless it is already present."""
    if robust_variant is None:
        return method
    _, existing_variant = split_robust_method(method)
    if existing_variant is not None:
        return method
    return f"{method}_{robust_variant}"


def method_display_name(method: str) -> str:
    """Return a display name, accepting both hyphen and underscore method ids."""
    return METHOD_NAMES.get(method, METHOD_NAMES.get(method.replace("-", "_"), method))


def format_plot_label(
    method: str,
    paper_mode: bool,
    n_experiments: Optional[int] = None,
    hide_robust_suffix: bool = False,
) -> str:
    """Format method labels, including robust strong-policy variants.

    The `(Robust ...)` qualifier is omitted when either `hide_robust_suffix`
    is True or `paper_mode` is True. In paper figures we expect the robust
    variant to be communicated via the caption, so leaving it in the per-
    method legend just clutters the label.
    """
    base_method, robust_variant = split_robust_method(method)
    label = method_display_name(base_method)

    if (
        robust_variant is not None
        and not hide_robust_suffix
        and not paper_mode
    ):
        label = f"{label} ({ROBUST_LABELS[robust_variant]})"

    if not paper_mode and n_experiments is not None:
        label = f"{label} (n={n_experiments})"

    return label


def prefix_matches_filter(
    prefix: str,
    robust_variant: Optional[str],
    prefix_filter: Optional[List[str]],
) -> bool:
    """Match both base prefixes and robust output prefixes."""
    if prefix_filter is None:
        return True

    candidates = [prefix]
    if robust_variant is not None:
        candidates.append(f"{prefix}_{robust_variant}")

    return any(candidate in prefix_filter for candidate in candidates)


def expand_method_order_for_robust_variants(
    method_order: List[str], available_methods: Dict[str, Dict[int, Path]]
) -> List[str]:
    """Place robust variants after their base method when a base order is provided."""
    expanded_order = []
    seen = set()

    for method in method_order:
        candidate_methods = [method]
        candidate_methods.extend(
            f"{method}_{robust_variant}" for robust_variant in ROBUST_VARIANTS
        )

        for candidate_method in candidate_methods:
            if candidate_method in available_methods and candidate_method not in seen:
                expanded_order.append(candidate_method)
                seen.add(candidate_method)

    return expanded_order


def method_is_filtered(method: str, method_filter: List[str]) -> bool:
    """Filter robust method variants when either full or base method is filtered."""
    base_method, _ = split_robust_method(method)
    return method in method_filter or base_method in method_filter


def method_is_included(method: str, method_include_filter: List[str]) -> bool:
    """Include robust method variants when either full or base method is included."""
    base_method, _ = split_robust_method(method)
    return method in method_include_filter or base_method in method_include_filter


def parse_method_dir(dir_name: str) -> Optional[Tuple[str, str, int]]:
    """
    Parse method directory name to extract env, method, and experiment ID.

    Expected format: {env}_{method}_exp{id}
    Examples: coinrun_max_prob_exp0, maze_proxy_fail_max_prob_exp1

    Returns:
        Tuple of (env, method, exp_id) or None if pattern doesn't match
    """
    pattern = rf"^({ENV_PATTERN})_(.+)_exp(\d+)$"
    match = re.match(pattern, dir_name)
    if match:
        env = match.group(1)
        method = match.group(2)
        exp_id = int(match.group(3))
        return env, method, exp_id
    return None


def extract_dummy_results(
    eval_dir: Path,
    prefix_filter: Optional[List[str]] = None,
    env_filter: Optional[str] = None,
) -> Dict[str, Dict[int, Path]]:
    """
    Extract evaluation results from dummy directory structure.

    Args:
        eval_dir: Directory containing evaluation results
        prefix_filter: Only include runs with these prefixes (list)
        env_filter: Only include runs for this environment

    Returns:
        Dictionary mapping method names to dict of {exp_id: result_file_path}
    """
    results: Dict[str, Dict[int, Path]] = defaultdict(dict)
    # Track the run-dir name we picked per (canonical_method, exp_id) so that
    # kebab/snake siblings collapse onto the most recent run.
    latest_run_dir_name: Dict[Tuple[str, int], str] = {}

    for child in eval_dir.iterdir():
        if not child.is_dir():
            continue

        robust_parsed = parse_robust_experiment_dir(child.name)
        robust_variant = None
        if robust_parsed is not None:
            prefix, robust_variant, env, exp_id = robust_parsed
        else:
            parsed = parse_experiment_dir(child.name)
            if parsed is None:
                continue
            prefix, env, exp_id = parsed

        if not prefix_matches_filter(prefix, robust_variant, prefix_filter):
            continue

        # Apply filters
        if env_filter is not None and env != env_filter:
            continue

        # Find method directories within this experiment
        for method_dir in child.iterdir():
            if not method_dir.is_dir():
                continue

            # Parse method directory name (format: {env}_{method}_exp{id})
            parsed_method = parse_method_dir(method_dir.name)
            if parsed_method is None:
                # Fallback to using the directory name as-is
                method_name = method_dir.name
            else:
                method_env, method_name, method_exp_id = parsed_method
                # Verify consistency with parent directory
                if method_exp_id != exp_id:
                    continue  # Skip mismatched experiment IDs
            method_name = add_robust_suffix(method_name, robust_variant)
            method_name = canonicalize_method(method_name)

            # Find the most recent run by directory-name timestamp
            # (YYYYMMDD_HHMMSS sorts lexicographically).
            latest_run = None
            latest_name = None

            for run_dir in sorted(method_dir.iterdir(), key=lambda p: p.name):
                if not run_dir.is_dir():
                    continue

                for run_file in run_dir.iterdir():
                    if run_file.is_file() and run_file.suffix == ".npz":
                        if latest_name is None or run_dir.name > latest_name:
                            latest_name = run_dir.name
                            latest_run = run_file
                        break

            if latest_run is not None:
                key = (method_name, exp_id)
                prev_name = latest_run_dir_name.get(key)
                if prev_name is None or latest_name > prev_name:
                    latest_run_dir_name[key] = latest_name
                    results[method_name][exp_id] = latest_run

    return dict(results)


def interpolate_to_common_x(
    x_arrays: List[np.ndarray],
    y_arrays: List[np.ndarray],
    num_points: int = 100,
) -> Tuple[np.ndarray, List[np.ndarray]]:
    """
    Interpolate multiple curves to a common x-axis for aggregation.

    Args:
        x_arrays: List of x-value arrays
        y_arrays: List of y-value arrays
        num_points: Number of points for interpolation

    Returns:
        Tuple of (common_x, list of interpolated y arrays)
    """
    # Find common x range
    x_min = max(arr.min() for arr in x_arrays)
    x_max = min(arr.max() for arr in x_arrays)

    common_x = np.linspace(x_min, x_max, num_points)
    interpolated_y = []

    for x, y in zip(x_arrays, y_arrays):
        # Sort by x for interpolation
        sort_idx = np.argsort(x)
        x_sorted = x[sort_idx]
        y_sorted = y[sort_idx]

        # Use linear interpolation
        f = interpolate.interp1d(
            x_sorted, y_sorted, kind="linear", fill_value="extrapolate"
        )
        interpolated_y.append(f(common_x))

    return common_x, interpolated_y


def calculate_minmax_bands(
    x_arrays: List[np.ndarray],
    y_arrays: List[np.ndarray],
    quantiles: Tuple[float, float] = (0.25, 0.75),
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Calculate min-max quantile bands from multiple curves.
    
    For each x value, interpolate all curves to that x, then calculate
    the median and specified quantiles across experiments.
    
    Args:
        x_arrays: List of x-value arrays for each experiment
        y_arrays: List of y-value arrays for each experiment
        quantiles: Tuple of (lower_quantile, upper_quantile), default (0.25, 0.75) for IQR
        
    Returns:
        Tuple of (common_x, y_median, y_lower_quantile, y_upper_quantile)
    """
    # Collect all unique x values
    all_x_values = set()
    for x_arr in x_arrays:
        all_x_values.update(x_arr.tolist())
    common_x = np.array(sorted(all_x_values))
    
    # Create interpolation functions for each experiment
    interp_funcs = []
    for x, y in zip(x_arrays, y_arrays):
        sort_idx = np.argsort(x)
        x_sorted = x[sort_idx]
        y_sorted = y[sort_idx]
        f = interpolate.interp1d(
            x_sorted, y_sorted, kind="linear", 
            bounds_error=False, fill_value=np.nan
        )
        interp_funcs.append(f)
    
    # Calculate statistics at each x value
    y_medians = []
    y_lower_quantiles = []
    y_upper_quantiles = []
    
    for x_val in common_x:
        # Collect y values from all experiments at this x
        y_values = []
        
        for f in interp_funcs:
            y_val = f(x_val)
            if not np.isnan(y_val):
                y_values.append(y_val)
        
        if len(y_values) > 0:
            # Calculate median and quantiles
            y_values = np.array(y_values)
            median = np.median(y_values)
            lower_q = np.quantile(y_values, quantiles[0])
            upper_q = np.quantile(y_values, quantiles[1])
            
            y_medians.append(median)
            y_lower_quantiles.append(lower_q)
            y_upper_quantiles.append(upper_q)
        else:
            # No values available at this x
            y_medians.append(np.nan)
            y_lower_quantiles.append(np.nan)
            y_upper_quantiles.append(np.nan)
    
    return common_x, np.array(y_medians), np.array(y_lower_quantiles), np.array(y_upper_quantiles)


def calculate_auc_with_bands(
    x: np.ndarray, 
    y_median: np.ndarray, 
    y_lower: np.ndarray, 
    y_upper: np.ndarray
) -> Tuple[float, float, float]:
    """
    Calculate area under the curve (AUC) for median and quantile bands.
    
    Args:
        x: Common x values
        y_median: Median y values
        y_lower: Lower quantile y values  
        y_upper: Upper quantile y values
        
    Returns:
        Tuple of (auc_median, auc_lower, auc_upper)
    """
    # Filter out NaN values
    valid_mask = ~(np.isnan(y_median) | np.isnan(y_lower) | np.isnan(y_upper))
    if np.sum(valid_mask) < 2:
        return np.nan, np.nan, np.nan
    
    x_valid = x[valid_mask]
    y_median_valid = y_median[valid_mask]
    y_lower_valid = y_lower[valid_mask]
    y_upper_valid = y_upper[valid_mask]
    
    # Calculate AUC using trapezoidal rule
    auc_median = integrate.trapezoid(y_median_valid, x_valid)
    auc_lower = integrate.trapezoid(y_lower_valid, x_valid)
    auc_upper = integrate.trapezoid(y_upper_valid, x_valid)
    
    return auc_median, auc_lower, auc_upper


def _print_endpoint_baselines(
    results: Dict[str, Dict[int, Path]],
    valid_methods: List[str],
) -> None:
    """Print mean novice and expert return broken out by ID / OOD / Average.

    Aggregates across the `(method, exp_id)` pairs that survived the user's
    method/robust/include/exclude filters (i.e. the same set the plot uses).
    Values are always in raw return units — they are not affected by
    `--normalize_y`, which only rescales the plot's y-axis.
    """
    split_name = "test"  # AFHP curve is evaluated on the test split

    novice = {"id": [], "ood": [], "all": []}
    expert = {"id": [], "ood": [], "all": []}

    def _resolve_summary(idx: int, meta_arr) -> Optional[Dict]:
        try:
            entry = meta_arr[idx]
        except (IndexError, TypeError):
            return None
        if hasattr(entry, "item"):
            entry = entry.item()
        if not isinstance(entry, dict):
            return None
        summary = entry.get("summary", {})
        if not isinstance(summary, dict):
            return None
        if split_name in summary:
            s = summary[split_name]
        else:
            s = next(iter(summary.values()), None)
        if not isinstance(s, dict):
            return None
        return s

    def _validate_role(s: Dict, role: str, file_label: str) -> bool:
        """Behavioral check: novice asked for help on no timestep; expert
        asked for help on the very first timestep of every episode."""
        if role == "novice":
            level_pred = s.get("level_ood_pred", None)
            if level_pred is None:
                return True  # nothing to check against; accept
            level_pred = list(level_pred)
            if any(bool(v) for v in level_pred):
                helped = int(sum(1 for v in level_pred if v))
                print(
                    f"Warning: novice baseline at {file_label} had help "
                    f"asked on {helped}/{len(level_pred)} episodes; "
                    "skipping for endpoint baseline aggregation."
                )
                return False
            return True
        if role == "expert":
            # In always-ask mode the policy must ask at its earliest feasible
            # timestep on every episode. The exact value of that earliest
            # timestep is policy-dependent: standard threshold policies record
            # 1 (1-indexed first action), while OracleLevelBasedRandomPolicy
            # structurally records 2 because it can't ask before the env's
            # first info dict surfaces the OOD label.
            #
            # We therefore check that `first_ood_timestep` is (a) non-None on
            # every episode and (b) identical across episodes — i.e. asks are
            # deterministic and as early as the policy can fire them.
            first_ood = s.get("first_ood_timestep", None)
            if first_ood is None:
                return True  # nothing to check against; accept
            first_ood = list(first_ood)
            if not first_ood:
                return True
            none_count = sum(1 for t in first_ood if t is None)
            unique = {t for t in first_ood if t is not None}
            if none_count == 0 and len(unique) == 1:
                return True
            bits = []
            if none_count > 0:
                bits.append(f"{none_count}/{len(first_ood)} episodes never asked")
            if len(unique) > 1:
                sample = sorted(unique)[:5]
                bits.append(
                    f"first_ood_timestep varied across episodes (e.g. {sample})"
                )
            print(
                f"Warning: expert baseline at {file_label}: {'; '.join(bits)}; "
                "skipping for endpoint baseline aggregation."
            )
            return False
        return True

    def _bucket(idx: int, meta_arr, role: str, bucket, file_label: str):
        s = _resolve_summary(idx, meta_arr)
        if s is None:
            return
        if not _validate_role(s, role, file_label):
            return
        raw = np.asarray(s.get("raw_returns", []), dtype=float)
        ood = np.asarray(s.get("level_ood_gt", []), dtype=bool)
        if raw.size == 0 or raw.size != ood.size:
            return
        bucket["id"].extend(raw[~ood].tolist())
        bucket["ood"].extend(raw[ood].tolist())
        bucket["all"].extend(raw.tolist())

    for method in valid_methods:
        exp_dict = results.get(method, {})
        for exp_id, npz_path in exp_dict.items():
            try:
                data = np.load(npz_path, allow_pickle=True)
            except Exception:
                continue
            if "afhps" not in data.files or "meta" not in data.files:
                continue
            afhps = np.asarray(data["afhps"], dtype=float)
            meta = data["meta"]
            if afhps.size == 0:
                continue
            file_label = f"{method}/exp{exp_id}"
            _bucket(int(np.argmin(afhps)), meta, "novice", novice, file_label)
            _bucket(int(np.argmax(afhps)), meta, "expert", expert, file_label)

    def _mean(values: List[float]) -> float:
        if not values:
            return float("nan")
        return float(np.mean(values))

    if not novice["all"] and not expert["all"]:
        # Nothing to print (e.g. .npz files lacked the expected meta keys).
        return

    print("\n" + "=" * 60)
    print("ENDPOINT BASELINES (raw mean return, filtered method x exp)")
    print("=" * 60)
    print(f"{'':<10}  {'ID':>10}  {'OOD':>10}  {'Average':>10}")
    print("-" * 60)
    print(
        f"{'Novice':<10}  "
        f"{_mean(novice['id']):>10.4f}  "
        f"{_mean(novice['ood']):>10.4f}  "
        f"{_mean(novice['all']):>10.4f}"
    )
    print(
        f"{'Expert':<10}  "
        f"{_mean(expert['id']):>10.4f}  "
        f"{_mean(expert['ood']):>10.4f}  "
        f"{_mean(expert['all']):>10.4f}"
    )
    print("=" * 60)


def _compute_accuracy_at_afhp(
    results: Dict[str, Dict[int, Path]],
    valid_methods: List[str],
    x_data_key: str,
    target_afhp: float = 0.5,
) -> Dict[str, Optional[Tuple[float, float, float]]]:
    """For each method, linearly interpolate `ood_accuracy` at `target_afhp`
    on every experiment's curve, then aggregate across experiments as
    (median, Q1, Q3).

    Returns:
        Mapping `method -> (median, q1, q3)`, or `method -> None` when no
        experiment had a curve that bracketed `target_afhp`.
    """
    out: Dict[str, Optional[Tuple[float, float, float]]] = {}
    for method in valid_methods:
        per_exp: List[float] = []
        for exp_id in sorted(results.get(method, {}).keys()):
            npz = results[method][exp_id]
            try:
                ed = np.load(npz, allow_pickle=True)
                xs, accs = extract_x_and_y_values(ed, x_data_key, "ood_accuracy")
            except Exception:
                continue
            xs = np.asarray(xs, dtype=float)
            accs = np.asarray(accs, dtype=float)
            if xs.size < 2 or accs.size < 2 or xs.size != accs.size:
                continue
            sort_idx = np.argsort(xs)
            xs = xs[sort_idx]
            accs = accs[sort_idx]
            if target_afhp < xs[0] or target_afhp > xs[-1]:
                continue
            f = interpolate.interp1d(
                xs, accs, kind="linear", bounds_error=False, fill_value=np.nan
            )
            val = float(f(target_afhp))
            if not np.isnan(val):
                per_exp.append(val)
        if per_exp:
            arr = np.asarray(per_exp)
            out[method] = (
                float(np.median(arr)),
                float(np.quantile(arr, 0.25)),
                float(np.quantile(arr, 0.75)),
            )
        else:
            out[method] = None
    return out


def print_auc_latex_table(method_auc_data: Dict[str, Tuple[float, float, float]],
                         x_label: str, y_label: str,
                         normalize_by_range: bool = True,
                         x_range: Optional[Tuple[float, float]] = None,
                         weak_performance: Optional[float] = None,
                         strong_performance: Optional[float] = None,
                         hide_robust_suffix: bool = False,
                         accuracy_data: Optional[
                             Dict[str, Optional[Tuple[float, float, float]]]
                         ] = None,
                         show_iqr: bool = True):
    """
    Print AUC results as a LaTeX-compatible table.

    Method ordering: PartialOracle row(s) first, then a separator, then the
    figure's default ordering (`DEFAULT_METHOD_ORDER` + alphabetical leftover).
    Bolding: always bolds the method with the highest AUC among non-oracle
    methods. PartialOracle is bolded only when it is the best of *all*
    methods (i.e. beats every non-oracle method on AUC); otherwise it is
    rendered in normal weight.
    """
    print("\n" + "="*80)
    print("AREA UNDER THE CURVE (AUC) RESULTS")
    print("="*80)

    # ---- Per-row normalization (applied once, used by both tables) -----------
    def _normalize_triplet(
        triplet: Tuple[float, float, float],
    ) -> Optional[Tuple[float, float, float]]:
        auc_median, auc_lower, auc_upper = triplet
        if np.isnan(auc_median):
            return None
        if weak_performance is not None and strong_performance is not None:
            perf_range = strong_performance - weak_performance
            if abs(perf_range) > 1e-6 and x_range is not None:
                baseline_area = weak_performance * (x_range[1] - x_range[0])
                auc_median = (auc_median - baseline_area) / perf_range
                auc_lower = (auc_lower - baseline_area) / perf_range
                auc_upper = (auc_upper - baseline_area) / perf_range
        if normalize_by_range and x_range is not None:
            range_width = x_range[1] - x_range[0]
            if range_width > 0:
                auc_median /= range_width
                auc_lower /= range_width
                auc_upper /= range_width
        return auc_median, auc_lower, auc_upper

    normalized: Dict[str, Optional[Tuple[float, float, float]]] = {
        m: _normalize_triplet(v) for m, v in method_auc_data.items()
    }

    # ---- Row ordering --------------------------------------------------------
    all_keys = list(method_auc_data.keys())
    po_keys = [
        m for m in all_keys
        if split_robust_method(m)[0] == "oracle-lb-random"
    ]
    non_po_keys_set = [m for m in all_keys if m not in po_keys]

    # Apply DEFAULT_METHOD_ORDER (with robust expansion) + alphabetical leftover.
    # expand_method_order_for_robust_variants only does `in` membership checks,
    # so an empty value-dict is sufficient as a presence map.
    non_po_lookup = {m: {} for m in non_po_keys_set}
    ordered_non_po = expand_method_order_for_robust_variants(
        DEFAULT_METHOD_ORDER, non_po_lookup  # type: ignore[arg-type]
    )
    leftover = sorted(m for m in non_po_keys_set if m not in ordered_non_po)
    ordered_non_po.extend(leftover)

    final_order: List[str] = list(po_keys) + ordered_non_po
    separator_after_idx = len(po_keys) - 1 if po_keys else -1

    # ---- Bold sets -----------------------------------------------------------
    def _best_key(
        candidates: List[str],
        data: Dict[str, Optional[Tuple[float, float, float]]],
    ) -> Optional[str]:
        best_k = None
        best_v = -float("inf")
        for k in candidates:
            triplet = data.get(k)
            if triplet is None:
                continue
            m = triplet[0]
            if isinstance(m, float) and np.isnan(m):
                continue
            if m > best_v:
                best_v = m
                best_k = k
        return best_k

    def _bold_winners(
        data: Dict[str, Optional[Tuple[float, float, float]]],
    ) -> "set[str]":
        """Best non-oracle row is always bolded. PartialOracle is bolded only
        when it tops every non-oracle row on the same metric."""
        winners: set[str] = set()
        best_non_po_k = _best_key(non_po_keys_set, data)
        best_overall_k = _best_key(all_keys, data)
        if best_non_po_k is not None:
            winners.add(best_non_po_k)
        if (
            best_overall_k is not None
            and split_robust_method(best_overall_k)[0] == "oracle-lb-random"
        ):
            winners.add(best_overall_k)
        return winners

    auc_bold_keys = _bold_winners(normalized)
    acc_bold_keys: "set[str]" = (
        _bold_winners(accuracy_data) if accuracy_data is not None else set()
    )

    # ---- LaTeX output --------------------------------------------------------
    has_accuracy = accuracy_data is not None
    auc_header = "AUC (Median [IQR])" if show_iqr else "Median AUC"
    acc_header = "Accuracy (Median [IQR])" if show_iqr else "Median Accuracy"
    print("\n% LaTeX Table")
    print("\\begin{table}[h]")
    print("\\centering")
    tab_spec = "lll" if has_accuracy else "ll"
    print(f"\\begin{{tabular}}{{{tab_spec}}}")
    print("\\toprule")
    if has_accuracy:
        print(f"Method & {auc_header} & {acc_header} \\\\")
    else:
        print(f"Method & {auc_header} \\\\")
    print("\\midrule")

    def _fmt_triplet(triplet: Optional[Tuple[float, float, float]]) -> str:
        if triplet is None:
            return "--"
        a, b, c = triplet
        if np.isnan(a):
            return "--"
        if show_iqr:
            return f"{a:.3f} [{b:.3f}, {c:.3f}]"
        return f"{a:.3f}"

    for i, method in enumerate(final_order):
        display_name = format_plot_label(
            method, paper_mode=True, hide_robust_suffix=hide_robust_suffix
        )
        value_str = _fmt_triplet(normalized[method])
        auc_is_bold = method in auc_bold_keys
        acc_is_bold = method in acc_bold_keys
        # Method-name bolding follows the AUC-column bolding, so the column
        # headers ("AUC" / "Accuracy") both highlight winners in their own
        # column without the method name flipping between rows.
        if auc_is_bold:
            display_name = f"\\textbf{{{display_name}}}"
            if value_str != "--":
                value_str = f"\\textbf{{{value_str}}}"
        if has_accuracy:
            acc_str = _fmt_triplet(accuracy_data.get(method))
            if acc_is_bold and acc_str != "--":
                acc_str = f"\\textbf{{{acc_str}}}"
            print(f"{display_name} & {value_str} & {acc_str} \\\\")
        else:
            print(f"{display_name} & {value_str} \\\\")
        if i == separator_after_idx and i < len(final_order) - 1:
            print("\\midrule")

    print("\\bottomrule")
    print("\\end{tabular}")
    print(f"\\caption{{Area Under the Curve (AUC) for {y_label} vs {x_label}.")
    if weak_performance is not None and strong_performance is not None:
        print(" Values are normalized such that 0 = novice performance and 1 = expert performance.")
    if normalize_by_range:
        print(" Values are also normalized by the x-axis range.")
    print(" IQR shows 25th-75th percentile range.}")
    print("\\label{tab:auc_results}")
    print("\\end{table}")

    # ---- Human-readable output ----------------------------------------------
    ANSI_BOLD = "\033[1m"
    ANSI_RESET = "\033[0m"

    if show_iqr:
        sep_width = 90 if has_accuracy else 60
    else:
        sep_width = 55 if has_accuracy else 45
    print("\n" + "-" * sep_width)
    print("Human-readable version:")
    print("-" * sep_width)
    if has_accuracy and show_iqr:
        print(
            f"{'Method':<30} {'AUC Median':<12} {'AUC IQR':<22} "
            f"{'Acc Median':<12} {'Acc IQR':<22}"
        )
    elif has_accuracy and not show_iqr:
        print(f"{'Method':<30} {'Median AUC':<12} {'Median Accuracy':<18}")
    elif not has_accuracy and show_iqr:
        print(f"{'Method':<30} {'AUC Median':<15} {'AUC IQR':<25}")
    else:
        print(f"{'Method':<30} {'Median AUC':<15}")
    print("-" * sep_width)

    def _fmt_human_triplet(
        triplet: Optional[Tuple[float, float, float]], median_width: int
    ) -> str:
        if triplet is None or np.isnan(triplet[0]):
            return f"{'N/A':<{median_width}}"
        a, b, c = triplet
        if show_iqr:
            return f"{a:<{median_width}.3f} [{b:.3f}, {c:.3f}]"
        return f"{a:<{median_width}.3f}"

    def _ansi(text: str, bold: bool) -> str:
        return f"{ANSI_BOLD}{text}{ANSI_RESET}" if bold else text

    for i, method in enumerate(final_order):
        display_name = (
            format_plot_label(
                method, paper_mode=True, hide_robust_suffix=hide_robust_suffix
            )
            .replace(r"\textsc{", "")
            .replace("}", "")
        )
        auc_is_bold = method in auc_bold_keys
        acc_is_bold = method in acc_bold_keys

        auc_width = 12 if (has_accuracy and show_iqr) else 15
        # Pad to width *before* ANSI-wrapping so the escape codes don't
        # consume column space.
        name_padded = f"{display_name:<30}"
        auc_padded = _fmt_human_triplet(normalized[method], auc_width)
        row = f"{_ansi(name_padded, auc_is_bold)} {_ansi(auc_padded, auc_is_bold)}"

        if has_accuracy:
            acc_width = 12 if show_iqr else 18
            acc_padded = _fmt_human_triplet(
                (accuracy_data or {}).get(method), acc_width
            )
            row = f"{row}  {_ansi(acc_padded, acc_is_bold)}"

        print(row)
        if i == separator_after_idx and i < len(final_order) - 1:
            print("-" * sep_width)

    print("-" * sep_width)


def plot_dummy_results(
    eval_dir: Path,
    prefix_filter: Optional[List[str]],
    env_filter: Optional[str],
    x_data_key: str,
    y_data_key: str,
    method_order: Optional[List[str]] = None,
    method_include_filter: Optional[List[str]] = None,
    method_filter: Optional[List[str]] = None,
    use_stderr: bool = True,
    disable_horizontal_lines: bool = False,
    disable_random_line: bool = False,
    save_path: Optional[str] = None,
    title: Optional[str] = None,
    no_aggregate: bool = False,
    paper_mode: bool = False,
    calculate_auc: bool = False,
    robust_filter: str = "all",
    normalize_y: bool = False,
    paper_app: bool = False,
    show_iqr: bool = True,
    show_accuracy: bool = False,
):
    """
    Plot dummy results with aggregation across experiments.

    Args:
        eval_dir: Directory containing evaluation results
        prefix_filter: Only include runs with these prefixes (list)
        env_filter: Only include runs for this environment
        x_data_key: Key for x-axis data
        y_data_key: Key for y-axis data
        method_order: Order of methods to plot
        method_include_filter: Methods to include
        method_filter: Methods to exclude
        use_stderr: If True, use standard error; otherwise use standard deviation
        disable_horizontal_lines: Disable weak/oracle reference lines
        disable_random_line: Disable random baseline diagonal line
        save_path: Path to save the figure
        title: Custom title for the plot (overrides auto-generated title)
        no_aggregate: Plot experiments separately instead of aggregating
        paper_mode: If True, remove title and n=X from labels for paper figures
        calculate_auc: If True, calculate and display AUC for each method
    """
    results = extract_dummy_results(eval_dir, prefix_filter, env_filter)

    if not results:
        print("No results found matching the filters.")
        return

    results = filter_results_by_robust(results, robust_filter)
    if not results:
        print(f"No results left after robust_filter={robust_filter!r}.")
        return

    # Canonicalize user-supplied method filters so kebab and snake CLI inputs
    # both match the kebab-canonical keys stored in `results`.
    if method_order is not None:
        method_order = [canonicalize_method(m) for m in method_order]
    if method_include_filter is not None:
        method_include_filter = [canonicalize_method(m) for m in method_include_filter]
    if method_filter is not None:
        method_filter = [canonicalize_method(m) for m in method_filter]

    # Determine method order
    if method_order is None:
        method_order = expand_method_order_for_robust_variants(
            DEFAULT_METHOD_ORDER, results
        )
        leftover = sorted(m for m in results if m not in method_order)
        method_order.extend(leftover)
    else:
        method_order = expand_method_order_for_robust_variants(method_order, results)

    if method_include_filter is not None:
        method_order = [
            m
            for m in method_order
            if method_is_included(m, method_include_filter)
        ]

    if method_filter is not None:
        method_order = [
            m for m in method_order if not method_is_filtered(m, method_filter)
        ]

    # Filter to valid methods
    valid_methods = [m for m in method_order if m in results]

    if not valid_methods:
        print("No valid methods found to plot.")
        return

    # If filtering to robust-only and a single robust variant survives, move
    # the "(Robust ...)" qualifier from per-method labels into the title.
    robust_variants_in_plot: set[str] = {
        v
        for v in (split_robust_method(m)[1] for m in valid_methods)
        if v is not None
    }
    hide_robust_label = (
        robust_filter == "robust" and len(robust_variants_in_plot) == 1
    )

    # Optional y-axis normalization: rescale every curve so that the mean
    # weak-agent return (curve y[0] across all method × exp) maps to 0 and the
    # mean expert return (curve y[-1] across all method × exp) maps to 1. We
    # need these constants before the plotting loop, so do a one-time
    # endpoint pre-scan over the .npz files.
    do_normalize_y = False
    weak_for_norm: Optional[float] = None
    perf_range_for_norm: Optional[float] = None
    if normalize_y:
        _first_y: List[float] = []
        _last_y: List[float] = []
        _weak_y: List[float] = []  # only populated for FILTERED_PERFORMANCE_KEYS
        for _m in valid_methods:
            for _exp_id in sorted(results[_m].keys()):
                try:
                    _ed = np.load(results[_m][_exp_id], allow_pickle=True)
                    _, _y_arr = extract_x_and_y_values(_ed, x_data_key, y_data_key)
                except Exception:
                    continue
                if len(_y_arr) == 0:
                    continue
                _first_y.append(float(_y_arr[0]))
                _last_y.append(float(_y_arr[-1]))
                if y_data_key in FILTERED_PERFORMANCE_KEYS:
                    try:
                        _, _y_unf = extract_x_and_y_values(
                            _ed, x_data_key, "performance"
                        )
                    except Exception:
                        _y_unf = []
                    if len(_y_unf) > 0:
                        _weak_y.append(float(_y_unf[0]))
        if _first_y and _last_y:
            if y_data_key in FILTERED_PERFORMANCE_KEYS and _weak_y:
                weak_for_norm = float(np.mean(_weak_y))
            else:
                weak_for_norm = float(np.mean(_first_y))
            _strong_for_norm = float(np.mean(_last_y))
            perf_range_for_norm = _strong_for_norm - weak_for_norm
            if abs(perf_range_for_norm) > 1e-6:
                do_normalize_y = True
            else:
                print(
                    "Warning: --normalize_y requested but mean weak ~= mean "
                    "expert; skipping y normalization."
                )

    # Set up plot style
    setup_plot_style(paper_mode=paper_mode, use_latex=True)
    
    # `--paper-app` swaps to a taller-than-default aspect ratio so the plot
    # sits well as an appendix figure alongside its (longer) legend.
    figsize = (8, 6) if paper_app else (8, 4.5)
    plt.figure(figsize=figsize)
    colors = sns.color_palette("husl", len(valid_methods))
    
    # Get line styles for paper mode
    line_styles = get_line_styles(len(valid_methods), paper_mode, valid_methods)

    # Store weak/oracle performance for reference lines
    all_first_performances = []
    all_last_performances = []
    # Track unfiltered weak agent performance (for performance_asked)
    all_weak_performances = []
    # Track x range for random baseline line
    all_x_min = []
    all_x_max = []
    
    # Store AUC data for each method
    method_auc_data = {}
    overall_x_min = float('inf')
    overall_x_max = float('-inf')

    for method_idx, method in enumerate(valid_methods):
        exp_data = results[method]
        exp_ids = sorted(exp_data.keys())

        if len(exp_ids) == 0:
            print(f"Warning: No experiments found for {method}, skipping...")
            continue

        # Override the auto palette color for special methods (e.g. PartialOracle).
        base_method_key, _ = split_robust_method(method)
        method_color = SPECIAL_METHOD_COLORS.get(base_method_key, colors[method_idx])

        # Load all experiment data
        x_arrays = []
        y_arrays = []
        meta_arrays = []  # Store meta data for confidence interval calculation

        for exp_id in exp_ids:
            data_path = exp_data[exp_id]
            try:
                eval_data = np.load(data_path, allow_pickle=True)
                x, y = extract_x_and_y_values(eval_data, x_data_key, y_data_key)

                if len(x) > 0:
                    # Apply optional y normalization in-place; downstream
                    # collectors and plot/AUC code then operate uniformly in
                    # the chosen y-space.
                    if do_normalize_y:
                        assert weak_for_norm is not None
                        assert perf_range_for_norm is not None
                        y = (y - weak_for_norm) / perf_range_for_norm

                    x_arrays.append(x)
                    y_arrays.append(y)

                    # Store meta data if available
                    if 'meta' in eval_data:
                        meta_arrays.append(eval_data['meta'])
                    else:
                        meta_arrays.append([])

                    # Track first/last for reference lines (same y-space as
                    # the plotted curves)
                    all_first_performances.append(y[0])
                    all_last_performances.append(y[-1])
                    # Track x range
                    all_x_min.append(x.min())
                    all_x_max.append(x.max())

                    # For filtered performance keys, also get unfiltered weak performance
                    if y_data_key in FILTERED_PERFORMANCE_KEYS:
                        _, y_unfiltered = extract_x_and_y_values(
                            eval_data, x_data_key, "performance"
                        )
                        if len(y_unfiltered) > 0:
                            weak_first = float(y_unfiltered[0])
                            if do_normalize_y:
                                assert weak_for_norm is not None
                                assert perf_range_for_norm is not None
                                weak_first = (
                                    weak_first - weak_for_norm
                                ) / perf_range_for_norm
                            all_weak_performances.append(weak_first)
            except Exception as e:
                print(f"Warning: Failed to load {data_path}: {e}")
                continue

        if len(x_arrays) == 0:
            print(f"Warning: No valid data for {method}, skipping...")
            continue
        
        # Print AFHP values for wait policy
        if method == "wait" and x_data_key in ["step_afhp", "level_afhp"]:
            print("\n=== Wait Policy AFHP Values ===")
            for exp_idx, x_values in enumerate(x_arrays):
                print(f"Experiment {exp_ids[exp_idx]}: {sorted(set(x_values))}")
            all_afhp_values = set()
            for x in x_arrays:
                all_afhp_values.update(x)
            print(f"All unique AFHP values across experiments: {sorted(all_afhp_values)}")
            print(f"Total unique values: {len(all_afhp_values)}")

        if len(x_arrays) == 1 or no_aggregate:
            # Single experiment or no aggregation mode
            if no_aggregate and len(x_arrays) > 1:
                # Plot each experiment separately with slightly different shades
                base_color = method_color
                for i, (x, y, exp_id) in enumerate(zip(x_arrays, y_arrays, exp_ids)):
                    sort_idx = np.argsort(x)
                    # Vary alpha or lightness for different experiments
                    alpha = 0.7 + (i / len(x_arrays)) * 0.3
                    base_label = format_plot_label(
                        method,
                        paper_mode,
                        n_experiments=None,
                        hide_robust_suffix=hide_robust_label,
                    )
                    exp_label = f"{base_label} exp{exp_id}"
                    
                    # Add markers for wait policy
                    if method == "wait":
                        plt.plot(
                            x[sort_idx],
                            y[sort_idx],
                            label=exp_label,
                            color=base_color,
                            alpha=alpha,
                            marker="o",
                            markersize=3,
                            linewidth=1.5,
                            linestyle=line_styles[method_idx],
                        )
                    else:
                        plt.plot(
                            x[sort_idx],
                            y[sort_idx],
                            label=exp_label,
                            color=base_color,
                            alpha=alpha,
                            linewidth=1.5,
                            linestyle=line_styles[method_idx],
                        )
            else:
                # Single experiment
                x, y = x_arrays[0], y_arrays[0]
                sort_idx = np.argsort(x)
                x_sorted = x[sort_idx]
                y_sorted = y[sort_idx]
                
                # Format label using shared function
                plot_label = format_plot_label(
                    method,
                    paper_mode,
                    n_experiments=1 if not paper_mode else None,
                    hide_robust_suffix=hide_robust_label,
                )
                plt.plot(
                    x_sorted,
                    y_sorted,
                    label=plot_label,
                    color=method_color,
                    marker="o" if method == "wait" else None,
                    markersize=4,
                    linestyle=line_styles[method_idx],
                )
                
                # Calculate AUC for single experiment
                if calculate_auc:
                    # For single experiment, use the same value for all bands
                    auc = integrate.trapezoid(y_sorted, x_sorted)
                    method_auc_data[method] = (auc, auc, auc)
                    
                    # Update overall x-range
                    if len(x_sorted) > 0:
                        overall_x_min = min(overall_x_min, x_sorted.min())
                        overall_x_max = max(overall_x_max, x_sorted.max())
        else:
            # Multiple experiments, aggregate using quantile bands
            # Use 25th-75th percentile bands (interquartile range)
            common_x, y_median, y_lower_q, y_upper_q = calculate_minmax_bands(
                x_arrays, y_arrays, quantiles=(0.25, 0.75)
            )
            
            # Filter out NaN values
            valid_mask = ~np.isnan(y_median)
            common_x = common_x[valid_mask]
            y_median = y_median[valid_mask]
            y_lower_q = y_lower_q[valid_mask]
            y_upper_q = y_upper_q[valid_mask]
            
            n_exps = len(x_arrays)
            
            # Format label using shared function
            plot_label = format_plot_label(
                method,
                paper_mode,
                n_experiments=n_exps,
                hide_robust_suffix=hide_robust_label,
            )
            
            # Plot median line
            plt.plot(
                common_x,
                y_median,
                label=plot_label,
                color=method_color,
                linewidth=2,
                linestyle=line_styles[method_idx],
            )

            # Plot quantile band (25th-75th percentile)
            plt.fill_between(
                common_x,
                y_lower_q,
                y_upper_q,
                color=method_color,
                alpha=0.2,
            )
            
            # Calculate AUC if requested
            if calculate_auc:
                auc_median, auc_lower, auc_upper = calculate_auc_with_bands(
                    common_x, y_median, y_lower_q, y_upper_q
                )
                method_auc_data[method] = (auc_median, auc_lower, auc_upper)
                
                # Update overall x-range
                if len(common_x) > 0:
                    overall_x_min = min(overall_x_min, common_x.min())
                    overall_x_max = max(overall_x_max, common_x.max())

    # Add reference lines
    if not disable_horizontal_lines and all_first_performances:
        mean_first = np.mean(all_first_performances)
        mean_last = np.mean(all_last_performances)

        # For filtered performance keys, use unfiltered weak performance
        if y_data_key in FILTERED_PERFORMANCE_KEYS and all_weak_performances:
            weak_perf = np.mean(all_weak_performances)
        else:
            weak_perf = mean_first

        plt.axhline(
            y=weak_perf,
            color="red",
            linestyle="--",
            alpha=0.7,
            label=r"\textsc{Novice}",
        )
        plt.axhline(
            y=mean_last,
            color="blue",
            linestyle="--",
            alpha=0.7,
            label=r"\textsc{Expert}",
        )

    # Add random baseline diagonal line (from weak to oracle)
    if not disable_random_line and all_first_performances and all_x_min:
        mean_last = np.mean(all_last_performances)
        x_min = min(all_x_min)
        x_max = max(all_x_max)

        # For filtered performance keys, use unfiltered weak performance
        if y_data_key in FILTERED_PERFORMANCE_KEYS and all_weak_performances:
            random_start = np.mean(all_weak_performances)
        else:
            random_start = np.mean(all_first_performances)

        # Use a unique line style for random that's different from all methods
        # Since methods get assigned styles starting from index 0, we'll use a later style
        random_linestyle = (0, (2, 2, 10, 2))  # Custom pattern: short dash, short dash, long gap, short dash
        if paper_mode:
            # In paper mode, ensure random gets a distinct style
            random_linestyle = (0, (2, 2, 10, 2))
        else:
            random_linestyle = ":"  # Keep dotted in non-paper mode
            
        plt.plot(
            [x_min, x_max],
            [random_start, mean_last],
            color="black",
            linestyle=random_linestyle,
            alpha=0.9,
            linewidth=2,
            label=r"\textsc{Random}",
            zorder=10,  # Draw on top
        )

    # `--paper-app`: zoom y-axis to the [6, 10] Procgen reward region
    # (paired with the taller aspect ratio set earlier). Applied only when
    # the curves are in raw units; with --normalize_y the plot is already
    # in a [0, 1]-ish frame and a 6-10 window would be empty.
    if paper_app:
        if do_normalize_y:
            print(
                "Warning: --paper-app y-axis zoom ignored because "
                "--normalize_y is set."
            )
        else:
            plt.ylim(6, 10)

    # Labels and title
    env_str = env_filter if env_filter else "all"
    prefix_str = ",".join(prefix_filter) if prefix_filter else "all"
    # Update error type to reflect quantile bands
    error_type = "IQR"  # Interquartile range

    # Get display names for axis labels
    x_label = DATA_KEY_NAMES.get(x_data_key, x_data_key)
    y_label = DATA_KEY_NAMES.get(y_data_key, y_data_key)
    if do_normalize_y:
        y_label = f"{y_label} (normalized)"

    # When the robust qualifier has been moved out of per-method labels (or the
    # user explicitly filtered to a robust set), surface it as a title prefix.
    if robust_filter == "robust" and robust_variants_in_plot:
        if len(robust_variants_in_plot) == 1:
            robust_title_prefix = ROBUST_LABELS[next(iter(robust_variants_in_plot))]
        else:
            robust_title_prefix = "Robust"
    elif robust_filter == "non-robust":
        robust_title_prefix = "Non-Robust"
    else:
        robust_title_prefix = None

    plt.xlabel(x_label)
    plt.ylabel(y_label)

    # Use custom title if provided, otherwise generate one (skip if paper_mode)
    if not paper_mode:
        if title:
            final_title = title
        else:
            final_title = (
                f"{y_label} vs {x_label} "
                f"({env_str}, prefix={prefix_str}, shaded={error_type})"
            )
            if robust_title_prefix is not None:
                final_title = f"[{robust_title_prefix}] {final_title}"
        plt.title(textwrap.fill(final_title, width=TITLE_WRAP_WIDTH))

    # Build the legend explicitly so that PartialOracle entries sit just before
    # the Novice reference line, preceded by a thin visual separator.
    legend_handles, legend_labels, separator_indices = _build_legend_entries(
        plt.gca()
    )

    # Apply publication styling
    style_plot_for_publication(
        legend_outside=True,
        legend_location='center left',
        legend_bbox_to_anchor=(1.05, 0.5),
        handles=legend_handles,
        labels=legend_labels,
    )

    plt.tight_layout()

    # After layout settles, expand any separator rows into full-legend-width
    # horizontal rules.
    legend_obj = plt.gca().get_legend()
    if legend_obj is not None and separator_indices:
        _draw_legend_separator(legend_obj, separator_indices)

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches="tight")
        print(f"Saved figure to {save_path}")
    else:
        plt.show()

    # Print novice / expert baseline returns broken out by ID vs OOD using
    # the per-episode meta saved in the boundary points of each filtered
    # .npz. These are always raw mean returns — unaffected by
    # `--normalize_y`, which only rescales the plot's y-axis.
    _print_endpoint_baselines(
        results=results,
        valid_methods=valid_methods,
    )

    # Print AUC table if requested
    if calculate_auc and method_auc_data:
        x_range = (overall_x_min, overall_x_max) if overall_x_min != float('inf') else None

        # When the curves were normalized at plot time, AUC is already in
        # the (novice=0, expert=1) frame; skip the table's redundant
        # rescaling by passing None for weak/strong.
        if do_normalize_y:
            weak_perf = None
            strong_perf = None
        else:
            weak_perf = None
            strong_perf = None
            if all_first_performances:
                # For filtered performance keys, use unfiltered weak performance
                if y_data_key in FILTERED_PERFORMANCE_KEYS and all_weak_performances:
                    weak_perf = np.mean(all_weak_performances)
                else:
                    weak_perf = np.mean(all_first_performances)
            if all_last_performances:
                strong_perf = np.mean(all_last_performances)

        # OOD accuracy at AFHP=50% per method (median + IQR across exps).
        # Skipped entirely when the column is disabled to avoid the extra
        # `.npz` re-load pass.
        accuracy_data = (
            _compute_accuracy_at_afhp(
                results=results,
                valid_methods=valid_methods,
                x_data_key=x_data_key,
                target_afhp=0.5,
            )
            if show_accuracy
            else None
        )

        print_auc_latex_table(method_auc_data, x_label, y_label,
                            normalize_by_range=True, x_range=x_range,
                            weak_performance=weak_perf,
                            strong_performance=strong_perf,
                            hide_robust_suffix=hide_robust_label,
                            accuracy_data=accuracy_data,
                            show_iqr=show_iqr)


def list_available_methods(
    eval_dir: Path,
    prefix_filter: Optional[List[str]],
    env_filter: Optional[str],
    robust_filter: str = "all",
):
    """List available methods and their experiment coverage."""
    results = extract_dummy_results(eval_dir, prefix_filter, env_filter)

    if not results:
        print("No results found matching the filters.")
        return

    results = filter_results_by_robust(results, robust_filter)
    if not results:
        print(f"No results left after robust_filter={robust_filter!r}.")
        return

    print("\nAvailable methods:")
    print("-" * 50)
    for method in sorted(results.keys()):
        exp_ids = sorted(results[method].keys())
        print(f"  {method}: exp{exp_ids}")
    print()


def main():
    parser = argparse.ArgumentParser(
        description="Plot dummy evaluation results with aggregation across experiments"
    )
    parser.add_argument(
        "--eval_dir",
        type=str,
        required=True,
        help="Directory containing the evaluation files.",
    )
    parser.add_argument(
        "--prefix",
        type=str,
        nargs="+",
        default=None,
        help="Prefix filter(s) for experiment directories (e.g., 'dummy04' or 'dummy04 dummy05')",
    )
    parser.add_argument(
        "--env",
        type=str,
        default=None,
        choices=SUPPORTED_ENVS,
        help="Environment filter",
    )
    parser.add_argument(
        "--x_data_key",
        type=str,
        default="level_afhp",
        help="Key for the x data (default: level_afhp)",
    )
    parser.add_argument(
        "--y_data_key",
        type=str,
        default="performance",
        help="Key for the y data (default: performance)",
    )
    parser.add_argument(
        "--method_order",
        "-m",
        type=str,
        default=None,
        help="Comma-separated list of methods to plot in order",
    )
    parser.add_argument(
        "--method_filter",
        "-f",
        type=str,
        nargs="+",
        default=None,
        help="Methods to exclude from plot",
    )
    parser.add_argument(
        "--method_include_filter",
        type=str,
        nargs="+",
        default=None,
        help="Only include these methods in the plot",
    )
    parser.add_argument(
        "--use_std",
        action="store_true",
        help="Use standard deviation instead of standard error for shaded region",
    )
    parser.add_argument(
        "--disable_horizontal_lines",
        action="store_true",
        help="Disable weak/oracle reference lines",
    )
    parser.add_argument(
        "--disable_random_line",
        action="store_true",
        help="Disable random baseline diagonal line",
    )
    parser.add_argument(
        "--save",
        type=str,
        default=None,
        help="Path to save the figure (if not specified, displays interactively)",
    )
    parser.add_argument(
        "--title",
        type=str,
        default=None,
        help="Custom title for the plot (overrides auto-generated title)",
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="List available methods and exit",
    )
    parser.add_argument(
        "--no_aggregate",
        action="store_true",
        help="Plot experiments separately instead of aggregating",
    )
    parser.add_argument(
        "--paper",
        action="store_true",
        help="Paper mode: remove title and n=X from labels for cleaner figures",
    )
    parser.add_argument(
        "--auc",
        action="store_true",
        help="Calculate and display area under the curve (AUC) for each method with IQR bands",
    )
    parser.add_argument(
        "--normalize_y",
        "--normalize-y",
        action="store_true",
        help=(
            "Rescale plotted curves so that mean novice (curve y[0] across "
            "all method × exp) maps to 0 and mean expert (curve y[-1]) maps "
            "to 1. The y-axis label is suffixed with '(normalized)'. When "
            "combined with --auc, the AUC table reports values already in "
            "this normalized frame."
        ),
    )
    parser.add_argument(
        "--paper-app",
        "--paper_app",
        dest="paper_app",
        action="store_true",
        help=(
            "Appendix-figure preset: implies --paper, zooms the y-axis to "
            "[6, 10] (Procgen reward range), and uses a taller (8x6) "
            "figure. The y-axis zoom is skipped when --normalize_y is set."
        ),
    )
    parser.add_argument(
        "--hide-iqr",
        "--hide_iqr",
        dest="hide_iqr",
        action="store_true",
        help=(
            "In the --auc table, omit the [IQR] brackets and the IQR sub-columns. "
            "Headers collapse to 'Median AUC' / 'Median Accuracy' and only the "
            "median value is printed per cell."
        ),
    )
    parser.add_argument(
        "--show-accuracy",
        "--show_accuracy",
        dest="show_accuracy",
        action="store_true",
        help=(
            "Add the Accuracy column to the --auc table. Off by default; "
            "enabling it triggers an extra per-exp interpolation pass over "
            "the .npz files to compute OOD accuracy at AFHP=50%%."
        ),
    )
    parser.add_argument(
        "--robust-filter",
        choices=["all", "robust", "non-robust"],
        default="all",
        help=(
            "Filter runs by robust-strong-policy variant. "
            "'all' (default) keeps everything; "
            "'robust' keeps only runs whose method name carries a robust suffix; "
            "'non-robust' drops all such runs. "
            "When the resulting set contains a single robust variant, the "
            "'(Robust ...)' qualifier is stripped from per-method labels and "
            "moved into the plot title."
        ),
    )

    args = parser.parse_args()

    # `--paper-app` is cumulative with `--paper`: opting into the appendix
    # preset always also implies paper mode.
    if args.paper_app:
        args.paper = True

    eval_dir = Path(args.eval_dir)

    if args.list:
        list_available_methods(
            eval_dir, args.prefix, args.env, robust_filter=args.robust_filter
        )
        return

    # Parse method order
    method_order = None
    if args.method_order:
        method_order = [m.strip() for m in args.method_order.split(",")]

    plot_dummy_results(
        eval_dir=eval_dir,
        prefix_filter=args.prefix,
        env_filter=args.env,
        x_data_key=args.x_data_key,
        y_data_key=args.y_data_key,
        method_order=method_order,
        method_include_filter=args.method_include_filter,
        method_filter=args.method_filter,
        use_stderr=not args.use_std,
        disable_horizontal_lines=args.disable_horizontal_lines,
        disable_random_line=args.disable_random_line,
        save_path=args.save,
        title=args.title,
        no_aggregate=args.no_aggregate,
        paper_mode=args.paper,
        calculate_auc=args.auc,
        robust_filter=args.robust_filter,
        normalize_y=args.normalize_y,
        paper_app=args.paper_app,
        show_iqr=not args.hide_iqr,
        show_accuracy=args.show_accuracy,
    )


if __name__ == "__main__":
    main()
