"""データセット集計の図（グループごとの Lab 分布・グループ平均の a*-b* 平面）.

`figures/` は回収ホワイトリストの対象であるため，**集計値の図だけ**を描き，画像そのもの・
切り出しは描かない（SPEC_01「実験の出力」）．個々の画像の値が点として現れないよう，
箱ひげ図は外れ値の点を描かず，ひげを 10/90 パーセンタイルにする（`dataset_summary.json` の
p10・p90 と一致する）．

pyplot を使わず `matplotlib.figure.Figure` を直接使う（GUI バックエンドを選ばず Agg で描くため．
サーバには表示環境が無い）．ラベルは日本語フォントの有無に左右されないよう英字 slug で書く．
"""

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 ALL_DIR_NAMES, parse_dir_name

LAB_BOXPLOT_FILENAME = "lab_by_group.png"
LAB_MEANS_AB_FILENAME = "lab_group_means_ab.png"

_CHANNEL_TITLES: tuple[str, ...] = ("L*", "a*", "b*")
_WHISKER_PERCENTILES = (10.0, 90.0)
_DPI = 120
# 程度ごとの目印（同じクラスの strong と mild を同じ色で並べ，形で見分ける）．
_SEVERITY_MARKERS: dict[str, str] = {"strong": "o", "mild": "^", "normal": "s"}
_CLASS_COLORS: dict[str, str] = {
    "pale": "#c9a0a0",
    "light_red": "#e07b7b",
    "red": "#d62728",
    "dark_red": "#7f1d1d",
    "purple": "#7b3fa0",
}


def plot_lab_by_group(lab_by_group: Mapping[str, np.ndarray], path: Path) -> None:
    """グループごとの L*・a*・b* の分布を箱ひげ図で描いて PNG に保存する.

    Args:
        lab_by_group: グループ名 → 画像ごとの舌領域の平均 Lab（`(N, 3)`, float64）．
        path: 出力先の PNG パス．親ディレクトリは作成される．
    """
    groups = [g for g in ALL_DIR_NAMES if g in lab_by_group]
    fig = Figure(figsize=(12, 10), dpi=_DPI)
    FigureCanvasAgg(fig)
    axes = fig.subplots(3, 1, sharex=True)
    for channel, (ax, title) in enumerate(zip(axes, _CHANNEL_TITLES, strict=True)):
        data = [lab_by_group[g][:, channel] for g in groups]
        positions = [i for i, d in enumerate(data) if d.size > 0]
        if positions:
            ax.boxplot(
                [data[i] for i in positions],
                positions=positions,
                whis=_WHISKER_PERCENTILES,
                showfliers=False,
                showmeans=True,
                widths=0.6,
            )
        counts = [d.size for d in data]
        ax.set_xticks(range(len(groups)))
        ax.set_xticklabels(
            [f"{g}\n(n={n})" for g, n in zip(groups, counts, strict=True)], fontsize=8
        )
        ax.set_xlim(-0.6, len(groups) - 0.4)
        ax.set_ylabel(title)
        ax.grid(axis="y", alpha=0.3)
    axes[0].set_title(
        "Per-image mean Lab of tongue region by group "
        "(box: quartiles, whiskers: 10-90th pct, triangle: mean)"
    )
    _save(fig, path)


def plot_group_means_ab(
    group_means: Mapping[str, Sequence[float] | None],
    group_stds: Mapping[str, Sequence[float] | None],
    path: Path,
) -> None:
    """グループ平均の a*-b* 平面の図（標準偏差のエラーバー付き）を PNG に保存する.

    mild グループが淡紅（`light_red_normal`）と同じクラスの strong のどちらに近いかを
    目で確かめるための図．

    Args:
        group_means: グループ名 → 平均 (L*, a*, b*)（無ければ None）．
        group_stds: グループ名 → 標準偏差 (L*, a*, b*)（無ければ None．要素が None でもよい）．
        path: 出力先の PNG パス．親ディレクトリは作成される．
    """
    fig = Figure(figsize=(8, 7), dpi=_DPI)
    FigureCanvasAgg(fig)
    ax = fig.subplots()
    for group in ALL_DIR_NAMES:
        mean = group_means.get(group)
        if mean is None:
            continue
        label, severity = parse_dir_name(group)
        std = group_stds.get(group) or (None, None, None)
        ax.errorbar(
            mean[1],
            mean[2],
            xerr=std[1] if std[1] is not None else None,
            yerr=std[2] if std[2] is not None else None,
            fmt=_SEVERITY_MARKERS[severity],
            color=_CLASS_COLORS[label],
            markersize=9,
            capsize=3,
            label=group,
        )
    ax.set_xlabel("a* (group mean, error bar: std)")
    ax.set_ylabel("b* (group mean, error bar: std)")
    ax.set_title("Group means on the a*-b* plane (o: strong, ^: mild, s: normal)")
    ax.grid(alpha=0.3)
    ax.legend(fontsize=8, loc="best")
    _save(fig, path)


def _save(fig: Figure, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fig.tight_layout()
    fig.savefig(path, format="png")