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

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

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

カラーマップ・軸ラベル・全体のスタイルは MiniTias-Evaluation の出力と揃えている
（実測 PNG とシミュ PNG を並べて比較できるように）．
"""

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.pyplot as plt  # noqa: E402
import numpy  # noqa: E402

# MiniTias-Evaluation の出力に合わせて hot カラーマップを統一使用する．
_LUMINANCE_CMAP = "hot"

# ゾーン境界の表示色（中心-中間，中間-周辺の 2 本）
_ZONE_BOUNDARY_COLORS = ["cyan", "lime"]

# フォールバック描画（luminance=None 時）のゾーン塗り分け色（center/middle/periphery）
_ZONE_FALLBACK_COLORS_HEX = ["#e74c3c", "#f39c12", "#2ecc71"]
# hex カラーを [0, 1] RGB タプルに変換済み（save_zone_map 呼び出しごとの変換を避けるため）
_ZONE_FALLBACK_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_FALLBACK_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=_LUMINANCE_CMAP, 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 として保存する．

    横軸は uint8 全範囲（0〜255）に固定して MiniTias-Evaluation と比較しやすくする．

    Args:
        luminance: 輝度画像（H x W, float）．
        output_path: 保存先 PNG パス．
        bins: ヒストグラムのビン数．デフォルトは 50．
    """
    fig, ax = plt.subplots(figsize=(8, 5))
    ax.hist(luminance.ravel(), bins=bins, color="#1f77b4", edgecolor="none")
    ax.set_title("Luminance Histogram")
    ax.set_xlabel("Luminance")
    ax.set_ylabel("Pixel Count")
    ax.set_xlim(0, 255)
    _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=(8, 5))
    ax.plot(distances, mean_luminance, marker="o", markersize=6, linewidth=2, color="#1f77b4")
    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,
    luminance: numpy.ndarray | None = None,
) -> None:
    """ゾーンマップを PNG として保存する．

    MiniTias-Evaluation 互換：背景に輝度マップ（hot）を表示し、
    その上にゾーン境界（center-middle，middle-periphery）を線で重ねる．

    `luminance` が None の場合は zone_map のみを 3 色で塗り分けるフォールバック動作．

    Args:
        zone_map: ゾーンラベルマップ（shape: (H, W), dtype: int8）．
            値は 0（center）, 1（middle）, 2（periphery）．
        output_path: 保存先 PNG パス．
        luminance: 背景に重ねる輝度マップ（H x W, float）．None ならゾーンの色塗りのみ．
    """
    fig, ax = plt.subplots(figsize=(6, 5))

    if luminance is not None:
        im = ax.imshow(luminance, cmap=_LUMINANCE_CMAP, origin="upper")
        fig.colorbar(im, ax=ax, label="Luminance")
        # ゾーン境界線（zone_map のラベル変化を contour で抽出）
        # zone_map: 0/1/2 → 境界は level=0.5（0-1）と level=1.5（1-2）
        ax.contour(
            zone_map.astype(float),
            levels=[0.5, 1.5],
            colors=_ZONE_BOUNDARY_COLORS,
            linewidths=2,
        )
    else:
        # フォールバック: 3 色塗り分け
        height, width = zone_map.shape
        rgb = numpy.zeros((height, width, 3), dtype=numpy.float32)
        for label, color_rgb in enumerate(_ZONE_FALLBACK_COLORS_RGB):
            mask = zone_map == label
            rgb[mask] = color_rgb
        ax.imshow(rgb, origin="upper")

    ax.set_title("Zone Map (center / middle / periphery)")
    ax.set_xlabel("X (px)")
    ax.set_ylabel("Y (px)")

    _save_figure(fig, output_path)