src/personalization_module/decision_prediction/common/plotting.py

from __future__ import annotations

from pathlib import Path
from typing import Any


def save_training_curves(
    *,
    history: list[dict[str, float]],
    output_path: Path,
    title: str = "Training Curves",
) -> Path | None:
    if not history:
        return None

    plt, sns = _load_plot_libs()
    sns.set_theme(style="whitegrid")

    epochs: list[float] = []
    train_loss: list[float] = []
    val_loss: list[float] = []

    for row in history:
        if not isinstance(row, dict):
            continue
        epoch = row.get("epoch")
        tr = row.get("train_loss")
        if epoch is None or tr is None:
            continue
        epochs.append(float(epoch))
        train_loss.append(float(tr))
        val = row.get("val_loss")
        val_loss.append(float(val) if val is not None else float("nan"))

    if not epochs:
        return None

    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig, ax = plt.subplots(figsize=(10, 5))
    sns.lineplot(x=epochs, y=train_loss, marker="o", label="train_loss", ax=ax)
    if any(_is_finite(v) for v in val_loss):
        sns.lineplot(x=epochs, y=val_loss, marker="o", label="val_loss", ax=ax)
    ax.set_title(title)
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Loss")
    ax.legend()
    fig.tight_layout()
    fig.savefig(output_path, dpi=140)
    plt.close(fig)
    return output_path


def save_confusion_matrix_plot(
    *,
    y_true: list[str],
    y_pred: list[str],
    labels: list[str],
    output_path: Path,
    title: str = "Confusion Matrix",
) -> Path | None:
    if not y_true or not y_pred or len(y_true) != len(y_pred):
        return None

    plt, sns = _load_plot_libs()
    sns.set_theme(style="white")

    matrix = _build_confusion_matrix(y_true=y_true, y_pred=y_pred, labels=labels)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    fig, ax = plt.subplots(figsize=(8, 6))
    sns.heatmap(
        matrix,
        annot=True,
        fmt="d",
        cmap="Blues",
        xticklabels=labels,
        yticklabels=labels,
        cbar=True,
        ax=ax,
    )
    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")
    ax.set_title(title)
    fig.tight_layout()
    fig.savefig(output_path, dpi=140)
    plt.close(fig)
    return output_path


def save_retrieval_recall_plot(
    *,
    recall_at_k: dict[int, float],
    hit_at_k: dict[int, float],
    output_path: Path,
    title: str = "Retrieval Recall@K / Hit@K",
) -> Path | None:
    if not recall_at_k and not hit_at_k:
        return None

    plt, sns = _load_plot_libs()
    sns.set_theme(style="whitegrid")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    ks = sorted(set(recall_at_k.keys()) | set(hit_at_k.keys()))
    if not ks:
        return None
    recall_values = [float(recall_at_k.get(k, float("nan"))) for k in ks]
    hit_values = [float(hit_at_k.get(k, float("nan"))) for k in ks]

    fig, ax = plt.subplots(figsize=(9, 5))
    sns.lineplot(x=ks, y=recall_values, marker="o", label="Recall@K", ax=ax)
    sns.lineplot(x=ks, y=hit_values, marker="o", label="Hit@K", ax=ax)
    ax.set_ylim(0.0, 1.0)
    ax.set_xlabel("K")
    ax.set_ylabel("Score")
    ax.set_title(title)
    ax.legend()
    fig.tight_layout()
    fig.savefig(output_path, dpi=140)
    plt.close(fig)
    return output_path


def save_retrieval_mrr_by_label_plot(
    *,
    mrr_by_label: dict[str, float],
    output_path: Path,
    title: str = "Retrieval MRR by Decision Label",
) -> Path | None:
    if not mrr_by_label:
        return None
    plt, sns = _load_plot_libs()
    sns.set_theme(style="whitegrid")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    labels = list(mrr_by_label.keys())
    values = [float(mrr_by_label[label]) for label in labels]

    fig, ax = plt.subplots(figsize=(9, 5))
    sns.barplot(x=labels, y=values, ax=ax, palette="Blues_d")
    ax.set_ylim(0.0, 1.0)
    ax.set_xlabel("Decision label")
    ax.set_ylabel("MRR")
    ax.set_title(title)
    fig.tight_layout()
    fig.savefig(output_path, dpi=140)
    plt.close(fig)
    return output_path


def save_retrieval_user_signal_heatmap(
    *,
    user_signal_scores: dict[str, dict[str, float]],
    output_path: Path,
    title: str = "Retrieval Top-K Mean Match Score (User x Signal)",
) -> Path | None:
    if not user_signal_scores:
        return None

    plt, sns = _load_plot_libs()
    sns.set_theme(style="white")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    users = sorted(user_signal_scores.keys())
    signals = sorted({signal for row in user_signal_scores.values() for signal in row.keys()})
    if not users or not signals:
        return None

    matrix: list[list[float]] = []
    for user in users:
        row = user_signal_scores.get(user, {})
        matrix.append([float(row.get(signal, 0.0)) for signal in signals])

    fig_w = max(10, int(0.45 * len(signals)) + 4)
    fig_h = max(4, int(0.6 * len(users)) + 3)
    fig, ax = plt.subplots(figsize=(fig_w, fig_h))
    sns.heatmap(
        matrix,
        cmap="YlGnBu",
        annot=False,
        xticklabels=signals,
        yticklabels=users,
        cbar=True,
        ax=ax,
    )
    ax.set_xlabel("Signal")
    ax.set_ylabel("User")
    ax.set_title(title)
    fig.tight_layout()
    fig.savefig(output_path, dpi=140)
    plt.close(fig)
    return output_path


def save_ablation_comparison_plot(
    *,
    with_retrieval: dict[str, float],
    without_retrieval: dict[str, float],
    output_path: Path,
    title: str = "Classification Ablation: With vs Without Retrieval",
) -> Path | None:
    keys = ["accuracy", "macro_f1"]
    if any(key not in with_retrieval for key in keys) or any(
        key not in without_retrieval for key in keys
    ):
        return None

    plt, sns = _load_plot_libs()
    sns.set_theme(style="whitegrid")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    labels = ["Accuracy", "Macro-F1"]
    x = [0, 1]
    with_vals = [float(with_retrieval["accuracy"]), float(with_retrieval["macro_f1"])]
    without_vals = [
        float(without_retrieval["accuracy"]),
        float(without_retrieval["macro_f1"]),
    ]

    fig, ax = plt.subplots(figsize=(8, 5))
    width = 0.34
    ax.bar([v - width / 2 for v in x], with_vals, width=width, label="with retrieval")
    ax.bar([v + width / 2 for v in x], without_vals, width=width, label="without retrieval")
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.set_ylim(0.0, 1.0)
    ax.set_ylabel("Score")
    ax.set_title(title)
    ax.legend()
    fig.tight_layout()
    fig.savefig(output_path, dpi=140)
    plt.close(fig)
    return output_path


def _build_confusion_matrix(
    *,
    y_true: list[str],
    y_pred: list[str],
    labels: list[str],
) -> list[list[int]]:
    index = {label: i for i, label in enumerate(labels)}
    matrix = [[0 for _ in labels] for _ in labels]

    for gold, pred in zip(y_true, y_pred):
        if gold not in index or pred not in index:
            continue
        matrix[index[gold]][index[pred]] += 1
    return matrix


def _is_finite(value: float) -> bool:
    return value == value and value not in (float("inf"), float("-inf"))


def _load_plot_libs() -> tuple[Any, Any]:
    try:
        import matplotlib.pyplot as plt
        import seaborn as sns
    except Exception as exc:  # pragma: no cover - runtime dependency guard
        raise RuntimeError(
            "Plotting requires matplotlib and seaborn. "
            "Install them with: pip install matplotlib seaborn"
        ) from exc
    return plt, sns