from typing import Optional, Tuple

import numpy as np
from PIL import Image

import flyr.picture_in_picture as pip
from flyr.dtypes import GrayArray, MaskArray, RGBArray


def emphasize_edges(
    bg_render_pil: Image.Image,
    optical: RGBArray,
    pip_info: pip.PictureInPicture,
    edge_emphasis: float,
    color_mask: Optional[MaskArray],
    custom_shift_x: int,
    custom_shift_y: int,
) -> Image.Image:
    ratio = pip_info.real_to_ir
    offset_x = pip_info.offset_x
    offset_y = pip_info.offset_y

    ratio = optical.shape[1] / bg_render_pil.size[0] / ratio
    src_size = bg_render_pil.size
    dst_size = (
        round(bg_render_pil.size[0] * ratio),
        round(bg_render_pil.size[1] * ratio),
    )

    height, width = optical.shape[:2]
    origin_x = (width // 2) + custom_shift_x
    origin_y = (height // 2) + custom_shift_y
    origin_x = origin_x - dst_size[0] // 2 + offset_x
    origin_y = origin_y - dst_size[1] // 2 + offset_y

    thermal_surface = (
        origin_x,
        origin_y,
        origin_x + dst_size[0],
        origin_y + dst_size[1],
    )
    embossment, opacity = _embossment_masks_pil(optical, edge_emphasis)
    embossment = embossment.crop(thermal_surface).resize(src_size)
    opacity = opacity.crop(thermal_surface).resize(src_size)

    if color_mask is not None:
        embossment = _apply_mask_pil(embossment, color_mask)
        opacity = _apply_mask_pil(opacity, color_mask)

    bg_render_pil = bg_render_pil.convert("RGBA")
    bg_render_pil.paste(embossment, (0, 0), opacity)
    return bg_render_pil


def _embossment_masks(
    optical: RGBArray, opacity: float = 0.275
) -> Tuple[GrayArray, GrayArray]:
    opacity = min(max(opacity, 0.0), 1.0)

    # Define azimuth, elevation, and depth
    ele = np.pi / 2.2  # radians
    azi = np.pi / 4.0  # radians
    dep = 15.0  # (0-100)

    # Get a B&W version of the image
    img = Image.fromarray(optical).convert("L")
    arr_embossment = np.array(img).astype(float)

    # Find the gradient and adjust the gradient by the depth factor
    grad_x, grad_y = np.gradient(arr_embossment)
    grad_x = grad_x * dep / 100.0
    grad_y = grad_y * dep / 100.0

    # Get the unit incident ray
    gd = np.cos(ele)  # length of projection of ray on ground plane
    dx = gd * np.cos(azi)
    dy = gd * np.sin(azi)
    dz = np.sin(ele)

    # Find the unit normal vectors for the image
    leng = np.sqrt(grad_x**2 + grad_y**2 + 1.0)
    uni_x = grad_x / leng
    uni_y = grad_y / leng
    uni_z = 1.0 / leng

    # Take the dot product, avoiding overflow
    arr_embossment = 255 * (dx * uni_x + dy * uni_y + dz * uni_z)
    arr_embossment = arr_embossment.clip(0, 255).astype(np.uint8)

    # Threshold non-edges to be white and see-through
    arr_embossment[arr_embossment > 180] = 255
    arr_embossment = arr_embossment.astype(np.uint8)
    arr_opacity = np.full_like(arr_embossment, fill_value=round(opacity * 255))
    arr_opacity[arr_embossment > 180] = 0
    return arr_embossment, arr_opacity


def _embossment_masks_pil(
    optical: RGBArray, opacity: float = 0.275
) -> Tuple[Image.Image, Image.Image]:
    e, o = _embossment_masks(optical, opacity)
    return Image.fromarray(e), Image.fromarray(o)


def _apply_mask_pil(image: Image.Image, mask: MaskArray):
    array_image = np.asarray(image)
    masked_image = _apply_mask(array_image, mask).filled(0)
    return Image.fromarray(masked_image)


def _apply_mask(image, mask: MaskArray):
    return np.ma.masked_array(image, np.invert(mask))
