"""位相相関法による画像位置合わせモジュール."""

import numpy as np
from scipy import ndimage
from scipy.fft import fft2, ifft2


def phase_correlate(reference: np.ndarray, target: np.ndarray) -> tuple[float, float]:
    """位相相関法で 2 枚の輝度マップ間の並進シフトを推定する．

    クロスパワースペクトルのピーク位置から並進シフトを算出し，
    ピーク周辺の放物線補間でサブピクセル精度を達成する．
    Foroosh et al., 2002 の手法に基づく [10]．

    Args:
        reference: 基準輝度マップ（H x W, float64）．
        target: ターゲット輝度マップ（H x W, float64）．

    Returns:
        (dy, dx) — target を reference に合わせるためのシフト量（ピクセル単位）．
        正値は正方向（下・右），負値は負方向（上・左）のシフトを示す．
    """
    # FFT でクロスパワースペクトルを算出
    f_ref = fft2(reference)
    f_tgt = fft2(target)
    cross_power = f_ref * np.conj(f_tgt)

    # ゼロ除算を防ぐため微小値を加算
    denom = np.abs(cross_power)
    denom = np.where(denom < 1e-10, 1e-10, denom)
    cross_power_norm = cross_power / denom

    # 逆 FFT で相関マップを算出（実部のみ使用）
    correlation = np.real(ifft2(cross_power_norm))

    # ピーク位置を取得
    peak_idx = np.unravel_index(np.argmax(correlation), correlation.shape)
    py, px = int(peak_idx[0]), int(peak_idx[1])
    h, w = correlation.shape

    # 放物線補間でサブピクセル精度を達成
    # y 方向の補間
    dy_sub = 0.0
    py_prev = (py - 1) % h
    py_next = (py + 1) % h
    c_prev = correlation[py_prev, px]
    c_curr = correlation[py, px]
    c_next = correlation[py_next, px]
    denom_y = c_prev - 2 * c_curr + c_next
    if abs(denom_y) > 1e-10:
        dy_sub = 0.5 * (c_prev - c_next) / denom_y

    # x 方向の補間
    dx_sub = 0.0
    px_prev = (px - 1) % w
    px_next = (px + 1) % w
    c_prev = correlation[py, px_prev]
    c_curr = correlation[py, px]
    c_next = correlation[py, px_next]
    denom_x = c_prev - 2 * c_curr + c_next
    if abs(denom_x) > 1e-10:
        dx_sub = 0.5 * (c_prev - c_next) / denom_x

    dy_float = float(py) + dy_sub
    dx_float = float(px) + dx_sub

    # FFT の周期境界により折り返し補正を適用
    # シフト量が画像サイズの半分を超える場合は反対方向として解釈する
    if dy_float > h / 2:
        dy_float -= h
    if dx_float > w / 2:
        dx_float -= w

    return dy_float, dx_float


def apply_shift(image: np.ndarray, dy: float, dx: float) -> np.ndarray:
    """サブピクセル精度でシフトを適用する．

    scipy.ndimage.shift（3 次スプライン補間）を使用する．
    境界領域は 0 でパディングする（mode='constant', cval=0.0）．

    Args:
        image: 入力輝度マップ（H x W, float64）．
        dy: y 方向のシフト量（ピクセル単位，正値は下方向）．
        dx: x 方向のシフト量（ピクセル単位，正値は右方向）．

    Returns:
        シフト後の輝度マップ（H x W, float64）．
    """
    return ndimage.shift(image, shift=(dy, dx), order=3, mode="constant", cval=0.0)


def compute_valid_mask(
    shape: tuple[int, int], dy: float, dx: float
) -> np.ndarray:
    """シフトにより無効になる境界領域を除いた有効領域のブールマスクを生成する．

    シフト量に基づいて境界から何ピクセル分を無効とするかを決定する．
    有効領域は両画像（シフト前・後）で有効なピクセルの積集合となる．

    Args:
        shape: 画像の形状 (height, width)．
        dy: y 方向のシフト量（ピクセル単位）．
        dx: x 方向のシフト量（ピクセル単位）．

    Returns:
        有効ピクセルを True とするブールマスク（H x W, bool）．
    """
    h, w = shape
    mask = np.ones((h, w), dtype=bool)

    # dy に応じた行方向の無効領域
    idy = int(np.ceil(abs(dy)))
    if dy > 0:
        # 下方向シフト → 上端が無効
        mask[:idy, :] = False
    elif dy < 0:
        # 上方向シフト → 下端が無効
        mask[h - idy :, :] = False

    # dx に応じた列方向の無効領域
    idx = int(np.ceil(abs(dx)))
    if dx > 0:
        # 右方向シフト → 左端が無効
        mask[:, :idx] = False
    elif dx < 0:
        # 左方向シフト → 右端が無効
        mask[:, w - idx :] = False

    return mask