"""交差検証の図（学習曲線・混同行列）.

`figures/` は回収対象であるため，集計値の図だけを描き，画像そのものは描かない（SPEC_01
「実験の出力」）．`data/survey_figures.py` と同じく pyplot を使わず `Figure` ＋
`FigureCanvasAgg` で描く．ラベルは日本語フォントの有無に左右されないよう英字で書く．
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from pathlib import Path

import numpy as np
from matplotlib.backends.backend_agg import FigureCanvasAgg
from matplotlib.figure import Figure

from tias_judge_model.data.classes import CLASS_SLUGS

LEARNING_CURVES_FILENAME = "learning_curves.png"
_DPI = 120


def confusion_matrix_filename(subset: str) -> str:
    """部分集合の混同行列の図のファイル名を返す（例: `confusion_matrix_all.png`）."""
    return f"confusion_matrix_{subset}.png"


def plot_learning_curves(history: Sequence[Mapping[str, float]], path: Path) -> None:
    """fold ごとの学習曲線（損失・検証 macro F1）を描いて PNG に保存する.

    Args:
        history: エポックごとの記録（`fold`・`epoch`・`train_loss`・`val_loss`・`val_macro_f1`）．
        path: 出力先の PNG パス．親ディレクトリは作成される．
    """
    folds = sorted({int(h["fold"]) for h in history})
    fig = Figure(figsize=(12, 5), dpi=_DPI)
    FigureCanvasAgg(fig)
    ax_loss, ax_f1 = fig.subplots(1, 2)
    for fold in folds:
        rows = [h for h in history if int(h["fold"]) == fold]
        epochs = [int(h["epoch"]) for h in rows]
        line = ax_loss.plot(epochs, [h["train_loss"] for h in rows], label=f"fold{fold} train")[0]
        ax_loss.plot(
            epochs,
            [h["val_loss"] for h in rows],
            linestyle="--",
            color=line.get_color(),
            label=f"fold{fold} val",
        )
        ax_f1.plot(epochs, [h["val_macro_f1"] for h in rows], label=f"fold{fold}")
    ax_loss.set_xlabel("epoch")
    ax_loss.set_ylabel("loss")
    ax_loss.set_title("loss (solid: train, dashed: val)")
    ax_loss.legend(fontsize=7)
    ax_f1.set_xlabel("epoch")
    ax_f1.set_ylabel("val macro F1")
    ax_f1.set_ylim(0.0, 1.0)
    ax_f1.set_title("validation macro F1")
    ax_f1.legend(fontsize=7)
    fig.tight_layout()
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path)


def plot_confusion_matrix(matrix: Sequence[Sequence[int]], title: str, path: Path) -> None:
    """混同行列（件数と行方向の割合）を描いて PNG に保存する.

    Args:
        matrix: 5×5 の混同行列（行が正解・列が予測．クラス ID 順）．
        title: 図のタイトル．
        path: 出力先の PNG パス．親ディレクトリは作成される．
    """
    cm = np.asarray(matrix, dtype=np.float64)
    row_sums = cm.sum(axis=1, keepdims=True)
    ratio = np.divide(cm, row_sums, out=np.zeros_like(cm), where=row_sums > 0)
    fig = Figure(figsize=(7, 6), dpi=_DPI)
    FigureCanvasAgg(fig)
    ax = fig.subplots()
    image = ax.imshow(ratio, vmin=0.0, vmax=1.0, cmap="Blues")
    fig.colorbar(image, ax=ax, label="row ratio")
    n = len(CLASS_SLUGS)
    ax.set_xticks(range(n), CLASS_SLUGS, rotation=30)
    ax.set_yticks(range(n), CLASS_SLUGS)
    ax.set_xlabel("predicted")
    ax.set_ylabel("true")
    ax.set_title(title)
    for i in range(n):
        for j in range(n):
            color = "white" if ratio[i, j] > 0.5 else "black"
            ax.text(
                j,
                i,
                f"{int(cm[i, j])}\n{ratio[i, j]:.2f}",
                ha="center",
                va="center",
                color=color,
                fontsize=8,
            )
    fig.tight_layout()
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path)