"""Module containing the implementation of image alignment algorithms."""

import numpy as np

from flyr.dtypes import Generic2DArray, MaskArray

__EPS: float = 1e-9  # Division by 0 prevention


def calculate_translation(
    image_1: MaskArray, image_2: MaskArray, threshold: int = 8
) -> tuple[int, int]:
    """Calculate the translation between two images using phase correlation. Return
    (dy, dx) translation that aligns `image_2` to `image_1`.

    Parameters
    ----------
    image_1 : MaskArray
        The first image.
    image_2 : MaskArray
        The second image.
    threshold : int, optional
        The minimum Peak to Second Peak Ratio to consider the translation valid, by default 8.

    Returns
    -------
    tuple[int, int]
        The translation between the two images.
    """
    if image_1.shape != image_2.shape:
        msg = "To estimate phase correlation, the inputs must have the same shape"
        raise ValueError(msg)

    matrix = _calculate_phase_correlation_matrix(image_1, image_2)
    matrix = _filter_non_maximum_peaks(matrix)
    max_y, max_x = _find_maximum_peaks_coords(matrix)

    metric = _calculate_metric(matrix, max_x, max_y)
    if metric > threshold:
        h, w = image_1.shape
        dy = max_y if max_y < h // 2 else max_y - h
        dx = max_x if max_x < w // 2 else max_x - w
        return int(dx), int(dy)
    return 0, 0


def _calculate_metric(
    matrix: Generic2DArray, max_x: int, max_y: int, window: int = 5
) -> float:
    """Calculate the Peak to Second Peak Ratio to understand if the maximum peak is stable.

    Parameters
    ----------
    max_x : int
        The x coordinate of the maximum peak.
    max_y : int
        The y coordinate of the maximum peak.
    window : int, optional
        The size of the window to exclude around the maximum peak, by default 5.

    Returns
    -------
    float
        The Peak to Second Peak Ratio.
    """

    h, w = matrix.shape
    peak = matrix[max_y, max_x]

    mask = np.ones_like(matrix, dtype=bool)
    y0 = max(0, max_y - window)
    y1 = min(h, max_y + window + 1)
    x0 = max(0, max_x - window)
    x1 = min(w, max_x + window + 1)

    mask[y0:y1, x0:x1] = False

    sidelobes = matrix[mask]
    mean = sidelobes.mean()
    std = sidelobes.std()

    return (peak - mean) / (std + __EPS)


def _calculate_phase_correlation_matrix(
    image_1: MaskArray, image_2: MaskArray
) -> Generic2DArray:
    """Calculate the phase correlation matrix between two images.

    Parameters
    ----------
    image_1 : MaskArray
        The first image.
    image_2 : MaskArray
        The second image.

    Returns
    -------
    Generic2DArray
        The phase correlation matrix.
    """
    f_image_1 = np.fft.fft2(image_1)
    f_image_2 = np.fft.fft2(image_2)

    r = f_image_1 * np.conj(f_image_2)
    r /= np.maximum(np.abs(r), __EPS)
    return np.abs(np.fft.ifft2(r))


def _filter_non_maximum_peaks(matrix: Generic2DArray) -> Generic2DArray:
    """Subtract the matrix by estimating the background as the median to filter the non-maximum peaks.

    Parameters
    ----------
    matrix : Generic2DArray
        The matrix to filter.

    Returns
    -------
    Generic2DArray
        The filtered matrix.
    """
    matrix -= np.median(matrix)
    return np.clip(matrix, 0, None)


def _find_maximum_peaks_coords(matrix: Generic2DArray) -> tuple[int, int]:
    """Find the coordinates of the maximum peak in the matrix.

    Returns
    -------
    Generic2DArray
        The coordinates of the maximum peak.
    """
    indices = np.unravel_index(np.argmax(matrix), matrix.shape)
    return (int(indices[0]), int(indices[1]))
