"""SmTIAS 定量モード DNG 読み込みモジュール."""

from pathlib import Path

import numpy
import rawpy


# Rec.709 グレースケール変換係数（uniformity.py の REC709_COEFF_* と一致させる）
# 参照: ITU-R BT.709 (https://www.itu.int/rec/R-REC-BT.709/en)
_REC709_R = 0.2126
_REC709_G = 0.7152
_REC709_B = 0.0722

# demosaic 引数 → rawpy.DemosaicAlgorithm のマッピング
# 参照: https://letmaik.github.io/rawpy/api/rawpy.DemosaicAlgorithm.html
_DEMOSAIC_MAP: dict[str, rawpy.DemosaicAlgorithm] = {
    "nearest": rawpy.DemosaicAlgorithm.LINEAR,  # 最近傍は rawpy に専用 enum がないため LINEAR で代替
    "bilinear": rawpy.DemosaicAlgorithm.LINEAR,
    "ahd": rawpy.DemosaicAlgorithm.AHD,
}

# 有効な output_format の一覧
_VALID_OUTPUT_FORMATS = {"rgb_float32", "luma_float32", "rgb_uint8"}

# rawpy が output_bps=16 で返す uint16 の最大値（0〜65535 を [0, 1] に正規化するため）
_UINT16_MAX = 65535.0


def load_image_dng(
    path: str,
    demosaic: str = "bilinear",
    black_level: int | None = None,
    output_format: str = "rgb_float32",
) -> numpy.ndarray:
    """SmTIAS 定量モード DNG を読み込み配列として返す．

    rawpy の postprocess() を linear パラメータで呼び出すことで，
    ISP（ガンマ・自動輝度・カメラ WB）の介入を最小化した linear raw 値を取り出す．
    具体的には gamma=(1,1), no_auto_bright=True, output_bps=16, use_camera_wb=False を指定する．

    Args:
        path: DNG ファイルのパス．
        demosaic: デモザイクアルゴリズム．既定 "bilinear"．
            "nearest"（最近傍）/ "bilinear"（双線形）/ "ahd"（適応的）から選択．
            "nearest" は rawpy に専用実装がないため "bilinear" と同一の LINEAR を使用する．
            bilinear（既定）は線形演算で比例性を保つ．
            **AHD は適応的＝非線形でチャンネル間の比例が崩れうるため量的評価では非推奨**．
        black_level: 黒レベル（整数）．None なら DNG Exif の
            black_level_per_channel[0] から自動取得する．通常は 64（SH-02M）．
        output_format: 出力形式．既定 "rgb_float32"．
            "rgb_float32": (H, W, 3) float32, linear, 値域 [0.0, 1.0]．
                output_color=raw（カラーマトリクス非適用）でカメラ固有チャンネルの
                linear raw 信号を出力する．各チャンネルがそのセンサチャンネルの
                入射光量に比例する（チャンネル混合なし）．
                飽和画素は 1.0 を超えることがある（クリップしない）．
            "luma_float32": (H, W) float32, linear, 値域 [0.0, 1.0]．
                output_color=raw（カラーマトリクス非適用）でカメラ固有チャンネルの
                linear raw 信号を取得し，Rec.709 重みで輝度に変換する
                （uniformity.py の to_grayscale と同じ係数）．
                **output_color=raw のとき Rec.709 luma は厳密な CIE 輝度ではなく，
                既存 PNG 経路（uniformity.py の to_grayscale）と互換の固定加重輝度プロキシ
                である**．Rec.709 係数を維持する理由は，既存パイプラインとの方法論的一貫性
                と DNG/PNG 評価指標の直接比較のためである．
            "rgb_uint8": (H, W, 3) uint8, カラーマトリクス＋sRGB OETF 適用後．
                output_color=sRGB（カラーマトリクス適用）＋ sRGB OETF を維持する．
                表示・PNG 互換専用であり，量的評価には linear 形式（rgb_float32 /
                luma_float32）を使うこと．

    Returns:
        numpy.ndarray. 形状・dtype は output_format に依存．

    Raises:
        FileNotFoundError: ファイルが存在しない場合．
        ValueError: demosaic または output_format の値が不正な場合，
            または rawpy が DNG を読み込めなかった場合（形式が不正等）．
    """
    if not Path(path).exists():
        raise FileNotFoundError(f"DNG ファイルが見つかりません: {path}")

    if demosaic not in _DEMOSAIC_MAP:
        raise ValueError(
            f"demosaic の値が不正です: '{demosaic}'．"
            f"有効な値: {sorted(_DEMOSAIC_MAP.keys())}"
        )

    if output_format not in _VALID_OUTPUT_FORMATS:
        raise ValueError(
            f"output_format の値が不正です: '{output_format}'．"
            f"有効な値: {sorted(_VALID_OUTPUT_FORMATS)}"
        )

    try:
        raw = rawpy.imread(str(path))
    except Exception as exc:
        raise ValueError(
            f"DNG の読み込みに失敗しました（形式が不正の可能性があります）: {path}"
        ) from exc

    # rawpy.postprocess() は DNG Exif の black_level_per_channel / white_level を
    # 内部で自動取得して黒レベル除去・正規化を行う．
    # postprocess に black_level を直接渡す API はないため，引数値の検証は行わない．
    # （将来の自前デモザイク実装時に black_level 引数を活用できるよう API として保持している）
    # SH-02M: black_level_per_channel=[64,64,64,64], white_level=1023 (scale=959)

    # 共通パラメータ: linear 出力のための ISP 無効化設定
    # - gamma=(1,1): ガンマ補正なし（linear のまま出力）
    # - no_auto_bright=True: 自動輝度調整を無効化
    # - output_bps=16: 16-bit 出力（精度確保のため）
    # - use_camera_wb=False + user_wb=[1,1,1,1]: WB ゲインをすべて 1 に固定（無補正）
    # - demosaic_algorithm: ユーザー指定アルゴリズム
    # rawpy.postprocess は DNG Exif から黒レベルと white_level を読み取り，
    # (raw - black_level) / (white_level - black_level) を 0〜65535 にスケールして出力する．
    # gamma=(1,1) により出力は linear のまま保たれる．
    # SH-02M の定量モード DNG は TONEMAP_MODE=CONTRAST_CURVE_LINEAR であり，
    # sRGB OETF は適用されていないため，このパラメータで正しく linear raw が得られる．
    _postprocess_common = {
        "demosaic_algorithm": _DEMOSAIC_MAP[demosaic],
        "gamma": (1, 1),
        "no_auto_bright": True,
        "output_bps": 16,
        "use_camera_wb": False,
        "user_wb": [1.0, 1.0, 1.0, 1.0],
    }

    if output_format in ("rgb_float32", "luma_float32"):
        # output_color=raw: カラーマトリクスを通さずカメラ固有チャンネルの linear raw 信号を出力．
        # これにより各チャンネルがそのセンサチャンネルの入射光量に比例する（チャンネル混合なし）．
        # SmTIAS-LightSim がチャンネル単位でシミュ vs 実機を突合するため per-channel の
        # センサ忠実度を優先する．output_color=sRGB（既定）はカラーマトリクスを適用するため
        # 各出力チャンネルが他チャンネルの線形結合となり，量的評価に不適切である．
        rgb16 = raw.postprocess(
            **_postprocess_common,
            output_color=rawpy.ColorSpace.raw,
        )
        # rgb16: (H, W, 3) uint16，値域 0〜65535（black_level 補正済み linear，white_level で正規化）
        rgb_float = rgb16.astype(numpy.float32) / _UINT16_MAX
        # rgb_float: (H, W, 3) float32, linear, 値域 [0, 1]

        raw.close()

        if output_format == "rgb_float32":
            return rgb_float

        # output_format == "luma_float32"
        # Rec.709 重みで輝度に変換（uniformity.py の to_grayscale と同じ係数）．
        # output_color=raw のとき Rec.709 luma は厳密な CIE 輝度ではなく固定加重輝度プロキシ．
        # 既存 PNG 経路（uniformity.py の to_grayscale）との方法論的一貫性と
        # DNG/PNG 評価指標の直接比較のために Rec.709 係数を維持する．
        luma = (
            _REC709_R * rgb_float[:, :, 0]
            + _REC709_G * rgb_float[:, :, 1]
            + _REC709_B * rgb_float[:, :, 2]
        )
        return luma  # (H, W) float32

    # output_format == "rgb_uint8"
    # output_color=sRGB: カラーマトリクス＋sRGB OETF 適用（既存 PNG 経路と互換）．
    # 表示・PNG 互換専用であり，量的評価には rgb_float32 / luma_float32 を使うこと．
    # sRGB OETF: V <= 0.0031308 → 12.92 * V, else 1.055 * V^(1/2.4) - 0.055
    # 参照: IEC 61966-2-1 (sRGB standard)
    rgb16 = raw.postprocess(
        **_postprocess_common,
        output_color=rawpy.ColorSpace.sRGB,
    )
    rgb_float = rgb16.astype(numpy.float32) / 65535.0
    raw.close()

    rgb_srgb = _apply_srgb_oetf(rgb_float)
    rgb_uint8 = numpy.clip(rgb_srgb * 255.0, 0, 255).astype(numpy.uint8)
    return rgb_uint8  # (H, W, 3) uint8


def _apply_srgb_oetf(linear: numpy.ndarray) -> numpy.ndarray:
    """Linear float 画像に sRGB OETF を適用する．

    IEC 61966-2-1 に基づく区分関数:
        V <= 0.0031308: out = 12.92 * V
        V >  0.0031308: out = 1.055 * V^(1/2.4) - 0.055

    Args:
        linear: (H, W, 3) float32, 値域 [0.0, 1.0]．

    Returns:
        sRGB OETF 適用後の (H, W, 3) float32, 値域 [0.0, 1.0]．
    """
    # ゼロ以下の値はそのまま（負値クリップはしない）
    out = numpy.where(
        linear <= 0.0031308,
        12.92 * linear,
        1.055 * numpy.power(numpy.maximum(linear, 0.0), 1.0 / 2.4) - 0.055,
    )
    return out.astype(numpy.float32)