"""可視化モジュール．

matplotlib を使って輝度マップ・ヒストグラム・放射状プロファイル・ゾーンマップを
PNG として保存する．画面表示は行わない（Agg バックエンド前提）．

各関数は figure を生成 → savefig → plt.close(fig) の形式で実装し，
メモリリークを防ぐ．
"""

from __future__ import annotations

from pathlib import Path

import matplotlib

# headless 環境（サーバ・CI 等）でも動くよう Agg バックエンドを強制する．
# pyplot を import する前に設定しないと既存バックエンドが優先されてしまうため，
# このモジュール内では matplotlib.use() を pyplot import より前に置く．
matplotlib.use("Agg")

import matplotlib.figure  # noqa: E402
import matplotlib.patches  # noqa: E402
import matplotlib.pyplot as plt  # noqa: E402
import numpy  # noqa: E402

# ゾーンラベルと RGB 値（classify_zones の 0=center/1=middle/2=periphery に対応）
_ZONE_LABELS = ["center", "middle", "periphery"]
_ZONE_COLORS_HEX = ["#e74c3c", "#f39c12", "#2ecc71"]
# hex カラーを [0, 1] RGB タプルに変換しておく（save_zone_map 呼び出しごとの変換を避けるため）
_ZONE_COLORS_RGB: list[tuple[float, float, float]] = [
    (int(h[1:3], 16) / 255.0, int(h[3:5], 16) / 255.0, int(h[5:7], 16) / 255.0)
    for h in _ZONE_COLORS_HEX
]


def _save_figure(fig: matplotlib.figure.Figure, output_path: Path, dpi: int = 150) -> None:
    """figure を PNG として保存してメモリを解放する．

    tight_layout → savefig → plt.close の順序を統一するためのヘルパ．

    Args:
        fig: 保存対象の Figure．
        output_path: 保存先 PNG パス．
        dpi: 解像度（デフォルト 150）．
    """
    fig.tight_layout()
    fig.savefig(output_path, dpi=dpi)
    plt.close(fig)


def save_luminance_map(luminance: numpy.ndarray, output_path: Path) -> None:
    """輝度マップを PNG として保存する．

    Args:
        luminance: 輝度画像（H x W, float）．
        output_path: 保存先 PNG パス．
    """
    fig, ax = plt.subplots(figsize=(6, 5))
    im = ax.imshow(luminance, cmap="viridis", origin="upper")
    fig.colorbar(im, ax=ax, label="Luminance")
    ax.set_title("Luminance Map")
    ax.set_xlabel("X [px]")
    ax.set_ylabel("Y [px]")
    _save_figure(fig, output_path)


def save_histogram(luminance: numpy.ndarray, output_path: Path, bins: int = 50) -> None:
    """輝度ヒストグラムを PNG として保存する．

    Args:
        luminance: 輝度画像（H x W, float）．
        output_path: 保存先 PNG パス．
        bins: ヒストグラムのビン数．デフォルトは 50．
    """
    fig, ax = plt.subplots(figsize=(6, 4))
    ax.hist(luminance.ravel(), bins=bins, color="#3498db", edgecolor="none", alpha=0.8)
    ax.set_title("Luminance Histogram")
    ax.set_xlabel("Luminance")
    ax.set_ylabel("Pixel Count")
    _save_figure(fig, output_path)


def save_radial_profile(profile: numpy.ndarray, output_path: Path) -> None:
    """放射状輝度プロファイルを線グラフとして PNG に保存する．

    `evaluation/radial.calc_radial_profile` が返す (n_bins, 2) 配列を受け取る．

    Args:
        profile: 放射状プロファイル配列（shape: (n_bins, 2)）．
            列 0: 正規化距離（ビン中心），列 1: 平均輝度．空ビンは nan．
        output_path: 保存先 PNG パス．
    """
    distances = profile[:, 0]
    mean_luminance = profile[:, 1]

    fig, ax = plt.subplots(figsize=(6, 4))
    ax.plot(distances, mean_luminance, marker="o", markersize=3, linewidth=1.5, color="#2c3e50")
    ax.set_title("Radial Luminance Profile")
    ax.set_xlabel("Normalized Distance from Center")
    ax.set_ylabel("Mean Luminance")
    ax.set_xlim(0.0, 1.0)
    ax.grid(True, linestyle="--", alpha=0.5)
    _save_figure(fig, output_path)


def save_zone_map(zone_map: numpy.ndarray, output_path: Path) -> None:
    """ゾーンマップを 3 色で PNG として保存する．

    `evaluation/zones.classify_zones` が返す int8 配列（0=center, 1=middle, 2=periphery）
    を 3 色に着色して凡例付きで保存する．

    Args:
        zone_map: ゾーンラベルマップ（shape: (H, W), dtype: int8）．
            値は 0（center）, 1（middle）, 2（periphery）．
        output_path: 保存先 PNG パス．
    """
    height, width = zone_map.shape
    rgb = numpy.zeros((height, width, 3), dtype=numpy.float32)

    for label, color_rgb in enumerate(_ZONE_COLORS_RGB):
        mask = zone_map == label
        rgb[mask] = color_rgb

    fig, ax = plt.subplots(figsize=(6, 5))
    ax.imshow(rgb, origin="upper")
    ax.set_title("Zone Map")
    ax.set_xlabel("X [px]")
    ax.set_ylabel("Y [px]")

    # imshow は自動凡例を持たないため，凡例パッチを手動で生成する
    legend_patches = [
        matplotlib.patches.Patch(color=_ZONE_COLORS_HEX[i], label=_ZONE_LABELS[i])
        for i in range(len(_ZONE_LABELS))
    ]
    ax.legend(handles=legend_patches, loc="upper right")

    _save_figure(fig, output_path)