"""  
Copyright (c) 2019-present NAVER Corp.
MIT License
"""
from typing import Tuple

import cv2
import numpy as np
from skimage import io


def loadImage(img_file: str) -> np.ndarray:
    """[summary] # TODO

    Args:
        img_file (str): [description]

    Returns:
        np.ndarray: [description]
    """
    img = io.imread(img_file)  # RGB order
    if img.shape[0] == 2:
        img = img[0]
    if len(img.shape) == 2:
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
    if img.shape[2] == 4:
        img = img[:, :, :3]
    img = np.array(img)

    return img


def normalizeMeanVariance(
    in_img: np.ndarray,
    mean: Tuple[float, float, float] = (0.485, 0.456, 0.406),
    variance: Tuple[float, float, float] = (0.229, 0.224, 0.225),
) -> np.ndarray:
    """[summary] # TODO

    Args:
        in_img (np.ndarray): [description]
        mean (Tuple[float, float, float], optional): [description]. Defaults to (0.485, 0.456, 0.406).
        variance (Tuple[float, float, float], optional): [description]. Defaults to (0.229, 0.224, 0.225).

    Returns:
        np.ndarray: [description]
    """
    # should be RGB order
    img = in_img.copy().astype(np.float32)

    img -= np.array([mean[0] * 255.0, mean[1] * 255.0, mean[2] * 255.0], dtype=np.float32)
    img /= np.array([variance[0] * 255.0, variance[1] * 255.0, variance[2] * 255.0], dtype=np.float32)
    return img


def denormalizeMeanVariance(
    in_img: np.ndarray,
    mean: Tuple[float, float, float] = (0.485, 0.456, 0.406),
    variance: Tuple[float, float, float] = (0.229, 0.224, 0.225),
) -> np.ndarray:
    """[summary] # TODO

    Args:
        in_img (np.ndarray): [description]
        mean (Tuple[float, float, float], optional): [description]. Defaults to (0.485, 0.456, 0.406).
        variance (Tuple[float, float, float], optional): [description]. Defaults to (0.229, 0.224, 0.225).

    Returns:
        np.ndarray: [description]
    """
    # should be RGB order
    img = in_img.copy()
    img *= variance
    img += mean
    img *= 255.0
    img = np.clip(img, 0, 255).astype(np.uint8)
    return img


def resize_aspect_ratio(
    img: np.ndarray, square_size: int, interpolation: int, mag_ratio: int = 1
) -> Tuple[np.ndarray, int, Tuple[int, int]]:
    """[summary] # TODO

    Args:
        img (np.ndarray): [description]
        square_size (int): [description]
        interpolation (int): [description]
        mag_ratio (int, optional): [description]. Defaults to 1.

    Returns:
        Tuple[np.ndarray, int, Tuple[int, int]]: [description]
    """
    height, width, channel = img.shape

    # magnify image size
    target_size = mag_ratio * max(height, width)

    # set original image size
    if target_size > square_size:
        target_size = square_size

    ratio = target_size / max(height, width)

    target_h, target_w = int(height * ratio), int(width * ratio)
    proc = cv2.resize(img, (target_w, target_h), interpolation=interpolation)

    # make canvas and paste image
    target_h32, target_w32 = target_h, target_w
    if target_h % 32 != 0:
        target_h32 = target_h + (32 - target_h % 32)
    if target_w % 32 != 0:
        target_w32 = target_w + (32 - target_w % 32)
    resized = np.zeros((target_h32, target_w32, channel), dtype=np.float32)
    resized[0:target_h, 0:target_w, :] = proc
    target_h, target_w = target_h32, target_w32

    size_heatmap = (int(target_w / 2), int(target_h / 2))

    return resized, ratio, size_heatmap


def cvt2HeatmapImg(img: np.ndarray) -> np.ndarray:
    """[summary] # TODO

    Args:
        img (np.ndarray): [description]

    Returns:
        np.ndarray: [description]
    """
    img = (np.clip(img, 0, 1) * 255).astype(np.uint8)
    img = cv2.applyColorMap(img, cv2.COLORMAP_JET)
    return img
