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

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


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

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

    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,
        "mean_ssim": float(np.mean(scores)),
        "min_ssim": float(np.min(scores)),
        "max_ssim": float(np.max(scores)),
        "std_ssim": float(np.std(scores, ddof=0)),
        "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 枚未満の場合．
    """
    if len(luminance_maps) < 2:
        raise ValueError(
            f"輝度マップは 2 枚以上必要です（現在: {len(luminance_maps)} 枚）．"
        )

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

    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,
        "mean_ssim": float(np.mean(scores)),
        "min_ssim": float(np.min(scores)),
        "max_ssim": float(np.max(scores)),
        "std_ssim": float(np.std(scores, ddof=0)),
        "mean_ssim_registered": float(np.mean(scores_reg)),
        "min_ssim_registered": float(np.min(scores_reg)),
        "max_ssim_registered": float(np.max(scores_reg)),
        "std_ssim_registered": float(np.std(scores_reg, ddof=0)),
        "mean_shift_magnitude": float(np.mean(shifts)),
        "max_shift_magnitude": float(np.max(shifts)),
        "n_pairs": len(pairwise_scores),
    }