# -*- coding: utf-8 -*-
import math
import os
import random

import cv2
import numpy as np
from matplotlib import cm
from PIL import Image


def load_image(image_path):
    """Loading image"""
    img = Image.open(image_path)
    # Fix bug RGBA
    if img.mode != "RGB":
        img = img.convert("RGB")
    img = np.array(img)
    return img


def crop_by_margin(image, margin=[0, 0]):
    """Cropping images by margins as a step of preprocessing"""
    H, W = image.shape[:2]
    margin_x, margin_y = margin
    image = image[margin_y : H - margin_y, margin_x : W - margin_x, :]
    return image


def gaussian_radius(det_size, min_overlap=0.7):
    """Calculating gaussian radius to compute std for Unnormalized Gaussian Mask"""
    height, width = det_size

    a1 = 1
    b1 = height + width
    c1 = width * height * (1 - min_overlap) / (1 + min_overlap)
    sq1 = np.sqrt(b1**2 - 4 * a1 * c1)
    r1 = (b1 + sq1) / 2

    a2 = 4
    b2 = 2 * (height + width)
    c2 = (1 - min_overlap) * width * height
    sq2 = np.sqrt(b2**2 - 4 * a2 * c2)
    r2 = (b2 + sq2) / 2

    a3 = 4 * min_overlap
    b3 = -2 * min_overlap * (height + width)
    c3 = (min_overlap - 1) * width * height
    sq3 = np.sqrt(b3**2 - 4 * a3 * c3)
    r3 = (b3 + sq3) / 2
    return min(r1, r2, r3)


def cal_mask_wh(p, mask):
    """Adaptively calculating blending mask W, H at the most vulnerable points perspective"""
    cy, cx = p
    mask_h, mask_w = mask.shape
    w = 0
    h = 0

    for i in [-1, 1]:
        shift_y = 0
        while (
            (cy + shift_y > -mask_h)
            and (cy + shift_y < mask_h)
            and (mask[cy + shift_y, cx] > 128)
        ):
            w += 1
            shift_y += i

        shift_x = 0
        while (
            (cx + shift_x > -mask_w)
            and (cx + shift_x < mask_w)
            and (mask[cy, cx + shift_x] > 128)
        ):
            h += 1
            shift_x += i

    return w, h


def overlay_mask(
    img: Image.Image, mask: Image.Image, colormap: str = "jet", alpha: float = 0.7
) -> Image.Image:
    """Overlay a colormapped mask on a background image

    >>> from PIL import Image
    >>> import matplotlib.pyplot as plt
    >>> from torchcam.utils import overlay_mask
    >>> img = ...
    >>> cam = ...
    >>> overlay = overlay_mask(img, cam)

    Args:
        img: background image
        mask: mask to be overlayed in grayscale
        colormap: colormap to be applied on the mask
        alpha: transparency of the background image

    Returns:
        overlayed image

    Raises:
        TypeError: when the arguments have invalid types
        ValueError: when the alpha argument has an incorrect value
    """

    if not isinstance(img, Image.Image) or not isinstance(mask, Image.Image):
        raise TypeError("img and mask arguments need to be PIL.Image")

    if not isinstance(alpha, float) or alpha < 0 or alpha >= 1:
        raise ValueError(
            "alpha argument is expected to be of type float between 0 and 1"
        )

    cmap = cm.get_cmap(colormap)
    # Resize mask and apply colormap
    overlay = mask.resize(img.size, resample=Image.BICUBIC)
    overlay = (255 * cmap(np.asarray(overlay) ** 1)[:, :, :3]).astype(np.uint8)
    # Overlay the image with the mask
    overlayed_img = Image.fromarray(
        (alpha * np.asarray(img) + (1 - alpha) * overlay).astype(np.uint8)
    )

    return overlayed_img


def bgr2ycbcr(img_bgr):
    img_bgr = img_bgr.astype(np.float32)
    img_ycrcb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2YCR_CB)
    img_ycbcr = img_ycrcb[:, :, (0, 2, 1)].astype(np.float32)
    # to [16/255, 235/255]
    img_ycbcr[:, :, 0] = (img_ycbcr[:, :, 0] * (235 - 16) + 16) / 255.0
    # to [16/255, 240/255]
    img_ycbcr[:, :, 1:] = (img_ycbcr[:, :, 1:] * (240 - 16) + 16) / 255.0

    return img_ycbcr


def ycbcr2bgr(img_ycbcr):
    img_ycbcr = img_ycbcr.astype(np.float32)
    # to [0, 1]
    img_ycbcr[:, :, 0] = (img_ycbcr[:, :, 0] * 255.0 - 16) / (235 - 16)
    # to [0, 1]
    img_ycbcr[:, :, 1:] = (img_ycbcr[:, :, 1:] * 255.0 - 16) / (240 - 16)
    img_ycrcb = img_ycbcr[:, :, (0, 2, 1)].astype(np.float32)
    img_bgr = cv2.cvtColor(img_ycrcb, cv2.COLOR_YCR_CB2BGR)

    return img_bgr


def gaussian_noise_color(img, param=None):
    if param is None:
        param = [0.001, 0.002, 0.005, 0.01, 0.05]
    ycbcr = bgr2ycbcr(img) / 255
    size_a = ycbcr.shape
    b = (
        ycbcr + math.sqrt(param) * np.random.randn(size_a[0], size_a[1], size_a[2])
    ) * 255
    b = ycbcr2bgr(b)
    img = np.clip(b, 0, 255).astype(np.uint8)
    return img


def block_wise(img, param):
    width = 8
    block = np.ones((width, width, 3)).astype(int) * 128
    param = min(img.shape[0], img.shape[1]) // 256 * param
    for i in range(param):
        r_w = random.randint(0, img.shape[1] - 1 - width)
        r_h = random.randint(0, img.shape[0] - 1 - width)
        img[r_h : r_h + width, r_w : r_w + width, :] = block

    return img


def color_saturation(img, param):
    ycbcr = bgr2ycbcr(img)
    ycbcr[:, :, 1] = 0.5 + (ycbcr[:, :, 1] - 0.5) * param
    ycbcr[:, :, 2] = 0.5 + (ycbcr[:, :, 2] - 0.5) * param
    img = ycbcr2bgr(ycbcr).astype(np.uint8)

    return img


def color_contrast(img, param):
    img = img.astype(np.float32) * param
    img = img.astype(np.uint8)

    return img


def gaussian_blur(img, param):
    img = cv2.GaussianBlur(img, (param, param), param * 1.0 / 6)

    return img


def jpeg_compression(img, param):
    h, w, _ = img.shape
    s_h = h // param
    s_w = w // param
    img = cv2.resize(img, (s_w, s_h))
    img = cv2.resize(img, (w, h))

    return img


def video_compression(vid_in, vid_out, param):
    cmd = f"ffmpeg -i {vid_in} -crf {param} -y {vid_out}"
    os.system(cmd)

    return