"""HDR 画像出力モジュール．

レンダリング結果の HDR float 配列を PNG または luminance.npy として保存する．
"""

from __future__ import annotations

from pathlib import Path

import numpy
from PIL import Image

from sm_tias_lightsim.uniformity import REC709_COEFF_B, REC709_COEFF_G, REC709_COEFF_R


def save_luminance(image: numpy.ndarray, path: Path) -> None:
    """HDR 画像を Rec.709 輝度に変換して .npy として保存する．

    uniformity.py の Rec.709 係数定数を再利用する（DRY）．
    HDR 値域（0〜∞）をそのまま保持するため clip は行わない．

    Args:
        image: HDR 画像 (H, W, 3) float．
        path: 出力 .npy ファイルパス．
    """
    luminance = (
        REC709_COEFF_R * image[..., 0]
        + REC709_COEFF_G * image[..., 1]
        + REC709_COEFF_B * image[..., 2]
    )
    numpy.save(path, luminance.astype(numpy.float64))


def save_png(image: numpy.ndarray, path: Path, gain: float = 1.0) -> None:
    """HDR 画像をトーンマップして 8-bit PNG として保存する．

    線形スケール（gain 倍）して [0, 1] にクリップ後，uint8 [0, 255] に変換する．
    物理的に妥当なトーンマッピングではなく，目視確認・デバッグ用途の簡易実装．

    Args:
        image: HDR 画像 (H, W, 3) float．
        path: 出力 PNG パス．
        gain: 露出補正係数．既定 1.0．
    """
    # gain を掛けて [0, 1] にクリップし uint8 化
    tonemapped = numpy.clip(image * gain, 0.0, 1.0)
    uint8_image = (tonemapped * 255.0).astype(numpy.uint8)
    pil_image = Image.fromarray(uint8_image, mode="RGB")
    pil_image.save(str(path))