"""再現性評価モジュール."""

from itertools import combinations

import numpy as np
from skimage.metrics import structural_similarity as ssim

from src.analysis.registration import apply_shift, compute_valid_mask, phase_correlate

# Rec.709 輝度（uint8 RGB → float64）の取りうる範囲は 0.0〜255.0
_DATA_RANGE = 255.0


def _validate_luminance_maps(luminance_maps: list[np.ndarray]) -> None:
    """輝度マップが SSIM 算出に必要な枚数（2 枚以上）あるか検証する．

    Args:
        luminance_maps: 輝度マップのリスト．

    Raises:
        ValueError: 輝度マップが 2 枚未満の場合．
    """
    if len(luminance_maps) < 2:
        raise ValueError(
            f"輝度マップは 2 枚以上必要です（現在: {len(luminance_maps)} 枚）．"
        )


def _aggregate_scores(scores: np.ndarray, prefix: str) -> dict:
    """スコア配列から mean/min/max/std を float で集約した辞書を返す．

    Args:
        scores: 集約対象のスコア配列．
        prefix: 辞書キーの接尾辞（例 "ssim" → "mean_ssim" 等）．

    Returns:
        {f"mean_{prefix}": ..., f"min_{prefix}": ...,
         f"max_{prefix}": ..., f"std_{prefix}": ...} 形式の辞書．
    """
    return {
        f"mean_{prefix}": float(np.mean(scores)),
        f"min_{prefix}": float(np.min(scores)),
        f"max_{prefix}": float(np.max(scores)),
        f"std_{prefix}": float(np.std(scores, ddof=0)),
    }


def calc_pairwise_ssim(luminance_maps: list[np.ndarray]) -> dict:
    """複数の輝度マップ間のペアワイズ SSIM を算出する．

    Args:
        luminance_maps: 輝度マップのリスト（各要素は同一サイズの 2D numpy 配列）．

    Returns:
        {
            "pairwise_scores": list[dict] — 各ペアの SSIM スコア
                [{"image_a": 0, "image_b": 1, "ssim": 0.998}, ...],
            "mean_ssim": float — 全ペアの平均 SSIM,
            "min_ssim": float — 最小 SSIM,
            "max_ssim": float — 最大 SSIM,
            "std_ssim": float — SSIM の標準偏差,
            "n_pairs": int — ペア数,
        }

    Raises:
        ValueError: 輝度マップが 2 枚未満の場合．
    """
    _validate_luminance_maps(luminance_maps)

    pairwise_scores: list[dict] = []
    for idx_a, idx_b in combinations(range(len(luminance_maps)), 2):
        score = float(
            ssim(
                luminance_maps[idx_a],
                luminance_maps[idx_b],
                data_range=_DATA_RANGE,
            )
        )
        pairwise_scores.append({"image_a": idx_a, "image_b": idx_b, "ssim": score})

    scores = np.array([p["ssim"] for p in pairwise_scores])

    return {
        "pairwise_scores": pairwise_scores,
        **_aggregate_scores(scores, "ssim"),
        "n_pairs": len(pairwise_scores),
    }


def calc_pairwise_ssim_registered(luminance_maps: list[np.ndarray]) -> dict:
    """位相相関による位置合わせ付きペアワイズ SSIM を算出する．

    各ペアに対して，位置合わせなしの SSIM と位置合わせ後の SSIM を両方算出する．
    位置合わせは phase_correlate でシフトを推定し，apply_shift で適用する．
    compute_valid_mask で有効領域を限定することで，
    境界のゼロパディング領域が SSIM に影響しないようにする．

    Args:
        luminance_maps: 輝度マップのリスト（各要素は同一サイズの 2D numpy 配列）．

    Returns:
        {
            "pairwise_scores": list[dict] — 各ペアのスコアと位置合わせ情報
                [{"image_a": 0, "image_b": 1,
                  "ssim": 0.911, "ssim_registered": 0.985,
                  "dx": 1.2, "dy": -0.5, "shift_magnitude": 1.3}, ...],
            "mean_ssim": float — 位置合わせなしの平均 SSIM,
            "min_ssim": float — 位置合わせなしの最小 SSIM,
            "max_ssim": float — 位置合わせなしの最大 SSIM,
            "std_ssim": float — 位置合わせなしの SSIM 標準偏差,
            "mean_ssim_registered": float — 位置合わせ後の平均 SSIM,
            "min_ssim_registered": float — 位置合わせ後の最小 SSIM,
            "max_ssim_registered": float — 位置合わせ後の最大 SSIM,
            "std_ssim_registered": float — 位置合わせ後の SSIM 標準偏差,
            "mean_shift_magnitude": float — 平均シフト量（ピクセル），
            "max_shift_magnitude": float — 最大シフト量（ピクセル），
            "n_pairs": int — ペア数,
        }

    Raises:
        ValueError: 輝度マップが 2 枚未満の場合．
    """
    _validate_luminance_maps(luminance_maps)

    pairwise_scores: list[dict] = []
    for idx_a, idx_b in combinations(range(len(luminance_maps)), 2):
        img_a = luminance_maps[idx_a]
        img_b = luminance_maps[idx_b]

        # 1. 位置合わせなし SSIM
        score_raw = float(ssim(img_a, img_b, data_range=_DATA_RANGE))

        # 2. 位相相関でシフト推定（img_b を img_a に合わせる）
        dy, dx = phase_correlate(img_a, img_b)

        # 3. シフト適用
        img_b_shifted = apply_shift(img_b, dy, dx)

        # 4. 有効領域マスクを生成し，有効ピクセル数を確認してクロップ
        mask = compute_valid_mask(img_a.shape, dy, dx)
        valid_count = int(np.sum(mask))

        # 有効領域が小さすぎる場合は位置合わせなしの値で代替
        # SSIM の win_size=7 を考慮し，有効ピクセル数 > 49 を最低条件とする
        if valid_count > 49:
            # 有効領域を 2D に再構成して SSIM を算出
            rows = np.where(mask.any(axis=1))[0]
            cols = np.where(mask.any(axis=0))[0]
            img_a_cropped = img_a[rows[0] : rows[-1] + 1, cols[0] : cols[-1] + 1]
            img_b_cropped = img_b_shifted[
                rows[0] : rows[-1] + 1, cols[0] : cols[-1] + 1
            ]
            score_registered = float(
                ssim(img_a_cropped, img_b_cropped, data_range=_DATA_RANGE)
            )
        else:
            score_registered = score_raw

        shift_magnitude = float(np.sqrt(dy**2 + dx**2))

        pairwise_scores.append({
            "image_a": idx_a,
            "image_b": idx_b,
            "ssim": score_raw,
            "ssim_registered": score_registered,
            "dx": float(dx),
            "dy": float(dy),
            "shift_magnitude": shift_magnitude,
        })

    scores = np.array([p["ssim"] for p in pairwise_scores])
    scores_reg = np.array([p["ssim_registered"] for p in pairwise_scores])
    shifts = np.array([p["shift_magnitude"] for p in pairwise_scores])

    return {
        "pairwise_scores": pairwise_scores,
        **_aggregate_scores(scores, "ssim"),
        **_aggregate_scores(scores_reg, "ssim_registered"),
        "mean_shift_magnitude": float(np.mean(shifts)),
        "max_shift_magnitude": float(np.max(shifts)),
        "n_pairs": len(pairwise_scores),
    }