"""Module containing the implementation of the distance estimation algorithm to
estimate the shift between thermal and optical images."""

from typing import Tuple, Union

import numpy as np
from PIL import Image, ImageOps

from flyr.dtypes import (
    GrayArray,
    MaskArray,
    RGBArray,
)
from flyr.msx.algorithms import calculate_translation
from flyr.picture_in_picture import PictureInPicture

__EPS = 1e-12  # Division by 0 prevention


def estimate_shift(
    thermal_image: GrayArray,
    optical_image: RGBArray,
    pip_info: PictureInPicture,
) -> Tuple[int, int]:
    """Estimate the shift between thermal and optical images using a distance estimation algorithm.

    Parameters
    ----------
    thermal_image: GrayArray
        The thermal image.
    optical_image: RGBArray
        The optical image.
    pip_info: PictureInPicture
        The picture in picture information.
    masks: Tuple[Optional[MaskArray], Optional[MaskArray]]
        The masks to be used in the phase correlation algorithm. If None, the masks will be
        built using the gradient edge maps of the images.

    Raises
    ------
    ValueError
        If the masks are provided but one of them is empty (contains no objects).
    ValueError
        If only one the optical mask is provided and the thermal mask is not.

    Returns
    -------
    Tuple[int, int]
        The estimated translation in the x-axis and the y-axis.
    """
    optical_crop, thermal_crop = _find_optical_thermal_matched_region(
        thermal_image, optical_image, pip_info
    )

    size = thermal_crop.shape[1], thermal_crop.shape[0]
    cover = ImageOps.cover(Image.fromarray(optical_crop), size)
    optical_crop = np.array(cover)
    thermal_edges, optical_edges = _calculate_edges(thermal_crop, optical_crop)

    if not np.any(thermal_edges):
        msg = "The mask for the thermal image must contain at least one object. But, it's empty."
        raise ValueError(msg)

    if not np.any(optical_edges):
        msg = "The mask for the optical image must contain at least one object. But, it's empty."
        raise ValueError(msg)

    return calculate_translation(optical_edges, thermal_edges)


def _calculate_edges(
    thermal_image: GrayArray,
    optical_image: RGBArray,
) -> Tuple[MaskArray, MaskArray]:
    """Generate the masks to be used in the distance estimation algorithm using
    gradient edge maps.

    Parameters
    ----------
    thermal_image: GrayArray
        The thermal image.
    optical_image: RGBArray
        The optical image.
    masks: Tuple[Optional[MaskArray], Optional[MaskArray]]
        The masks to be used in the phase correlation algorithm.

    Raises
    ------
    ValueError
        If only the optical mask is provided and the thermal mask is not.

    Returns
    -------
    Tuple[MaskArray, MaskArray]
        The masks to be used in the phase correlation algorithm.
    """
    thermal_mask = _build_gradient_edge_maps(thermal_image)
    optical_mask = _build_gradient_edge_maps(optical_image)
    return thermal_mask, optical_mask


def _find_optical_thermal_matched_region(
    thermal: GrayArray, optical: RGBArray, pip_info: PictureInPicture
) -> Tuple[RGBArray, GrayArray]:
    """
    Finds the optical and thermal images scene regions aligned to the thermal PiP
    crop box.

    Parameters
    ----------
    thermal: GrayArray
        The thermal image.
    optical: RGBArray
        The optical image.
    pip_info: PictureInPicture
        The picture in picture information.

    Returns
    -------
    Tuple[RGBArray, GrayArray]
        The cropped optical and thermal images aligned to the thermal PiP crop box.
    """
    # Compute optical crop bounds using picture_in_picture_pil math.
    x1, y1, x2, y2 = pip_info.crop_box
    ratio = pip_info.real_to_ir
    render_w = thermal.shape[1]
    ratio = optical.shape[1] / render_w / ratio
    dst_w = round((x2 - x1) * ratio)
    dst_h = round((y2 - y1) * ratio)

    width = optical.shape[1]
    height = optical.shape[0]
    origin_x = width // 2 - dst_w // 2 + pip_info.offset_x
    origin_y = height // 2 - dst_h // 2 + pip_info.offset_y

    optical_pil = Image.fromarray(optical)
    optical_crop = optical_pil.crop(
        (origin_x, origin_y, origin_x + dst_w, origin_y + dst_h)
    )

    return np.array(optical_crop), thermal


def _build_gradient_edge_maps(image: Union[GrayArray, RGBArray]) -> MaskArray:
    """Build the gradient edge map of the image using a thresholding method.

    An L1 normalization is applied to the gradients to amplify edges.

    Parameters
    ----------
    image : Union[GrayArray, RGBArray]
        The image to build the gradient edge map from.

    Returns
    -------
    MaskArray
        The gradient edge map of the image.
    """
    thresholded_edges = _otsu_thresold(image)
    gradients = np.gradient(thresholded_edges.astype(np.float32))
    gradient_magnitude = np.sum(np.abs(gradients), axis=0)
    return gradient_magnitude.astype(np.uint8) > 0


def _otsu_thresold(image: Union[GrayArray, RGBArray]) -> GrayArray:
    """Applies Otsu's thresholding method to the input image.

    Parameters
    ----------
    image: Union[GrayArray, RGBArray]
        The input image to be thresholded, either grayscale or RGB which will
        then be converted to grayscale.

    Returns
    -------
    GrayArray
        The thresholded image.
    """
    image = np.array(Image.fromarray(image).convert("L"))
    if image.size == 0:
        return image.astype(np.uint8)

    min_val, max_val = float(image.min()), float(image.max())
    hist, bin_edges = np.histogram(image.ravel(), bins=256, range=(min_val, max_val))
    hist = hist.astype(np.float64)
    bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2.0

    prob = hist / hist.sum()
    cum_prob = np.cumsum(prob)
    cum_mean = np.cumsum(prob * bin_centers)
    global_mean = cum_mean[-1]

    between_var = (global_mean * cum_prob - cum_mean) ** 2 / (
        cum_prob * (1.0 - cum_prob) + __EPS
    )
    threshold = bin_centers[int(np.argmax(between_var))]

    return ~((image > threshold).astype(np.uint8) * 255)
