trigger-off/src/evaluation/visualize.py

"""
Visualisation utilities for Trigger-Off experiment results.

Requires: matplotlib
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _ensure_matplotlib() -> None:
    try:
        import matplotlib  # noqa: F401
    except ImportError as exc:
        raise ImportError(
            "matplotlib is required for visualisation. "
            "Install with: pip install matplotlib"
        ) from exc


# ---------------------------------------------------------------------------
# Main results bar chart
# ---------------------------------------------------------------------------

def plot_main_results(
    results_dict: Dict[str, Dict[str, float]],
    metric: str = "step_acc",
    title: str = "Trigger-Off: Main Results",
    save_path: Optional[str] = None,
) -> None:
    """
    Plot a bar chart comparing Step-Acc (or another metric) across baselines.

    Args:
        results_dict: {baseline_name: {metric_name: value, ...}}
            e.g. {"SeeClick-Clean": {"step_acc": 0.72, ...},
                  "SeeClick+TriggerOff": {"step_acc": 0.68, ...}}
        metric: Which metric to plot (default: "step_acc").
        title: Chart title.
        save_path: If provided, save the figure to this path.
    """
    _ensure_matplotlib()
    import matplotlib.pyplot as plt
    import numpy as np

    labels = list(results_dict.keys())
    values = [results_dict[k].get(metric, 0.0) for k in labels]

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

    fig, ax = plt.subplots(figsize=(max(6, len(labels) * 1.4), 5))
    bars = ax.bar(x, values, width, color="steelblue", edgecolor="white", alpha=0.85)

    ax.set_ylabel(metric.replace("_", " ").title())
    ax.set_title(title)
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=9)
    ax.set_ylim(0, min(1.05, max(values) * 1.2 + 0.05))

    for bar, val in zip(bars, values):
        ax.text(
            bar.get_x() + bar.get_width() / 2.0,
            bar.get_height() + 0.01,
            f"{val:.3f}",
            ha="center",
            va="bottom",
            fontsize=8,
        )

    fig.tight_layout()

    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    else:
        plt.show()

    plt.close(fig)


# ---------------------------------------------------------------------------
# Ablation line chart
# ---------------------------------------------------------------------------

def plot_ablation(
    ablation_results: Dict[str, List[float]],
    variable_name: str,
    x_values: List[Any],
    metric: str = "step_acc",
    title: Optional[str] = None,
    save_path: Optional[str] = None,
) -> None:
    """
    Plot a line chart for ablation experiments.

    Args:
        ablation_results: {series_name: [value_at_x1, value_at_x2, ...]}
            e.g. {"clean_acc": [0.72, 0.71, 0.70],
                  "attack_acc": [0.40, 0.32, 0.28],
                  "immunity_acc": [0.70, 0.69, 0.68]}
        variable_name: X-axis label, e.g. "poison_ratio".
        x_values: Values on the X-axis, e.g. [0.10, 0.15, 0.20, 0.25].
        metric: Metric label for Y axis (default: "step_acc").
        title: Chart title.
        save_path: If provided, save the figure to this path.
    """
    _ensure_matplotlib()
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(figsize=(7, 4))

    markers = ["o", "s", "^", "D", "v", "P"]
    for i, (series_name, y_values) in enumerate(ablation_results.items()):
        marker = markers[i % len(markers)]
        ax.plot(x_values, y_values, marker=marker, label=series_name, linewidth=2)

    ax.set_xlabel(variable_name.replace("_", " ").title())
    ax.set_ylabel(metric.replace("_", " ").title())
    ax.set_title(title or f"Ablation: {variable_name}")
    ax.legend(loc="best", fontsize=9)
    ax.set_ylim(0, 1.05)
    ax.grid(True, linestyle="--", alpha=0.5)

    fig.tight_layout()

    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
    else:
        plt.show()

    plt.close(fig)


# ---------------------------------------------------------------------------
# Results table
# ---------------------------------------------------------------------------

def print_results_table(results: Dict[str, Any]) -> None:
    """
    Print a nicely formatted results table to stdout.

    Args:
        results: Dict of {scenario_name: EvalResult_or_dict}.
    """
    # Normalise to dicts
    rows: Dict[str, Dict[str, Any]] = {}
    for name, r in results.items():
        if hasattr(r, "to_dict"):
            rows[name] = r.to_dict()
        elif isinstance(r, dict):
            rows[name] = r
        else:
            rows[name] = {"value": r}

    if not rows:
        print("(No results to display.)")
        return

    # Determine columns
    columns = ["step_acc", "task_sr", "asr", "immunity_rate"]
    col_labels = {
        "step_acc": "Step-Acc",
        "task_sr": "Task-SR",
        "asr": "ASR",
        "immunity_rate": "Imm-Rate",
    }

    # Column widths
    name_width = max(len(n) for n in rows) + 2
    col_width = 10

    header = f"{'Scenario':<{name_width}}" + "".join(
        f"{col_labels.get(c, c):>{col_width}}" for c in columns
    )
    separator = "-" * len(header)

    print("\n" + separator)
    print(header)
    print(separator)

    for name, row in rows.items():
        line = f"{name:<{name_width}}"
        for col in columns:
            val = row.get(col, None)
            if val is None:
                cell = "N/A"
            else:
                cell = f"{val:.4f}"
            line += f"{cell:>{col_width}}"
        print(line)

    print(separator + "\n")