"""画像全体で評価する，ノイズ補正付き最小/最大比のモジュール．

輪帯平均の min/max（``radial.calc_radial_min_max_ratio``）は方位方向に平均するため，
同心円状でない輝度分布（TEST_18 のオフ軸 LED 配置）では不均一を平均で打ち消してしまう．
一方，画素単位の min/max や平滑化後のパーセンタイルはモンテカルロ（MC）ノイズの裾を拾い，
生の輝度が低い大きなキャビティで値が頭打ちになる．

本モジュールはその中間として「ガウス平滑化した像の min/max を，同じ平滑化を掛けた
ノイズ場の極値で補正する」定義を提供する．平滑化で方位方向に平均せずノイズだけを落とし，
残ったノイズの寄与（平滑化後でも残る極値のバイアス）を統計的に差し引く．

補正に使うノイズ場は，残差（生 − メディアン/ガウスで均した像）を空間シャッフルして
同じ σ で平滑化したもの．シャッフルで構造を壊し，ノイズの振幅分布だけを残す
（``evaluation`` 外の ``scripts/build_radius_sweep_visual.py:decompose_cov`` と同じ発想）．

仮定:
    ノイズが空間的に独立（画素間で無相関）であること．シャッフルは画素の並べ替えなので，
    もともと空間相関を持つノイズ（デノイザ適用後の像等）では補正量を過小評価する．
    MC レンダの生出力（デノイズ無し）を前提とする．

検証:
    軸上スイープ（TEST_15 球 / TEST_16 立方体の計 20 条件，分布が同心円対称なので
    輪帯平均が基準として使える）で，本定義は輪帯平均と ±0.01 以内で一致し，
    一様像にノイズだけを載せたときの見かけの不均一（ノイズ床）は ±0.004 だった．
    検証スクリプト: ``scripts/build_fullimage_minmax.py``（定義名 ``smooth_corr``）．
"""

from __future__ import annotations

import numpy
from scipy.ndimage import gaussian_filter, median_filter


def estimate_noise_extremes(
    image: numpy.ndarray,
    sigma: float,
    residual_sigma: float = 4.0,
    median_size: int = 5,
    n_shuffle: int = 4,
    seed: int = 0,
) -> tuple[float, float]:
    """ノイズだけの場に σ=``sigma`` の平滑化を掛けたときの (min, max) の期待値を返す．

    手順:
        1. 残差 $r = I - G_{\\sigma_r}(M_k(I))$ を取る（$M_k$ = サイズ $k$ のメディアン
           フィルタ，$G_{\\sigma_r}$ = σ=``residual_sigma`` のガウス平滑化）．
           メディアンで firefly を落としてからガウスで均すことで，残差に構造を残さない．
        2. 残差を空間シャッフルして構造を壊し，σ=``sigma`` のガウス平滑化を掛ける．
        3. その場の min / max を ``n_shuffle`` 回平均する．

    乱数は ``numpy.random.default_rng(seed)`` をこの関数内で生成するため，
    同じ入力・同じ ``seed`` なら常に同じ値を返す（呼び出し順に依存しない）．

    Args:
        image: 2D 輝度配列（H x W）．
        sigma: 本体の平滑化と同じガウス σ（画素単位）．
        residual_sigma: 残差を取るときのガウス σ（画素単位）．
        median_size: 残差を取るときのメディアンフィルタのサイズ（画素）．
        n_shuffle: シャッフルの反復回数（多いほど期待値の推定が安定する）．
        seed: 乱数シード．

    Returns:
        (ノイズ場の最小値の期待値, ノイズ場の最大値の期待値)．
        それぞれ平滑化後の像の min / max から差し引くための量（min 側は負，max 側は正）．
    """
    img = numpy.asarray(image, dtype=numpy.float64)
    rng = numpy.random.default_rng(seed)

    residual = img - gaussian_filter(median_filter(img, size=median_size), sigma=residual_sigma)

    mins: list[float] = []
    maxs: list[float] = []
    for _ in range(n_shuffle):
        flat = residual.ravel().copy()
        rng.shuffle(flat)
        field = gaussian_filter(flat.reshape(img.shape), sigma=sigma)
        mins.append(float(field.min()))
        maxs.append(float(field.max()))

    return float(numpy.mean(mins)), float(numpy.mean(maxs))


def calc_fullimage_min_max_ratio(
    image: numpy.ndarray,
    sigma: float = 12.0,
    residual_sigma: float = 4.0,
    median_size: int = 5,
    n_shuffle: int = 4,
    seed: int = 0,
) -> dict[str, float]:
    """画像全体で評価する，ノイズ補正付き最小/最大比を算出する．

    $\\mathrm{ratio} = \\dfrac{\\min G_\\sigma(I) - n_{\\min}}{\\max G_\\sigma(I) - n_{\\max}}$

    ここで $G_\\sigma$ は σ=``sigma`` のガウス平滑化，$(n_{\\min}, n_{\\max})$ は
    ``estimate_noise_extremes`` が返すノイズ場の極値の期待値．
    平滑化で MC ノイズを落としつつ方位方向には平均しないので，
    オフ軸 LED のような方位依存の不均一も min/max に反映される．

    検証（``scripts/build_fullimage_minmax.py``）: 同心円対称な軸上スイープ 20 条件で
    輪帯平均の min/max と ±0.01 以内で一致し，ノイズ床（一様像 + ノイズでの 1 からのずれ）は
    ±0.004．ノイズが空間的に独立であることを仮定する（モジュール docstring 参照）．

    Args:
        image: 2D 輝度配列（H x W）．ROI 抽出後の輝度を渡すこと．
        sigma: ガウス平滑化の σ（画素単位）．既定 12.0 は 464x450 の ROI で検証した値．
        residual_sigma: ノイズ場推定の残差を取るガウス σ（画素単位）．
        median_size: ノイズ場推定の残差を取るメディアンフィルタのサイズ（画素）．
        n_shuffle: ノイズ場の極値を平均する回数．
        seed: ノイズ場推定の乱数シード（同じ入力・同じ seed なら決定論的）．

    Returns:
        以下のキーを持つ辞書:
        - ``"fullimage_min_max_ratio"`` (float): 補正後の最小/最大比（1 が理想）．
        - ``"fullimage_min"`` (float): 補正後の最小値（平滑化後 min − ノイズ場 min）．
        - ``"fullimage_max"`` (float): 補正後の最大値（平滑化後 max − ノイズ場 max）．
        - ``"fullimage_noise_min"`` (float): ノイズ場の最小値の期待値．
        - ``"fullimage_noise_max"`` (float): ノイズ場の最大値の期待値．

    Raises:
        ValueError: image が 2D の数値配列でない場合．
        ValueError: sigma が 0 以下の場合．
        ValueError: 画像が平滑化に対して小さすぎる場合（min(H, W) < 4 * sigma）．
        ValueError: 補正後の最大値が 0 以下の場合（画像がノイズだけとみなせる）．
    """
    img = numpy.asarray(image)
    if img.ndim != 2:
        msg = f"image は 2D 配列を期待します（ndim={img.ndim}）"
        raise ValueError(msg)
    if not numpy.issubdtype(img.dtype, numpy.number):
        msg = f"image は数値配列を期待します（dtype={img.dtype}）"
        raise ValueError(msg)
    img = img.astype(numpy.float64)

    if sigma <= 0.0:
        msg = f"sigma は正の値を期待します（sigma={sigma}）"
        raise ValueError(msg)

    # 平滑化カーネルが画像からはみ出すと端の値が境界条件に支配されるため下限を設ける
    min_side = min(img.shape)
    if min_side < 4.0 * sigma:
        msg = (
            f"画像が平滑化 σ={sigma:g} に対して小さすぎます"
            f"（min(H, W)={min_side} < 4σ={4.0 * sigma:g}）．sigma を小さくしてください．"
        )
        raise ValueError(msg)

    smooth = gaussian_filter(img, sigma=sigma)
    noise_min, noise_max = estimate_noise_extremes(
        img,
        sigma=sigma,
        residual_sigma=residual_sigma,
        median_size=median_size,
        n_shuffle=n_shuffle,
        seed=seed,
    )

    corrected_min = float(smooth.min()) - noise_min
    corrected_max = float(smooth.max()) - noise_max
    if corrected_max <= 0.0:
        msg = (
            "ノイズ補正後の最大値が 0 以下です"
            f"（平滑化後 max={smooth.max():.6g}, ノイズ場 max={noise_max:.6g}）．"
            "画像が実質ノイズのみの可能性があります．"
        )
        raise ValueError(msg)

    return {
        "fullimage_min_max_ratio": corrected_min / corrected_max,
        "fullimage_min": corrected_min,
        "fullimage_max": corrected_max,
        "fullimage_noise_min": noise_min,
        "fullimage_noise_max": noise_max,
    }