# encoding: utf-8
"""
@author: ccj
@contact:
"""

import numpy as np

from typing import List, Dict, Tuple, Any

import torch
import torch.nn.functional as F


def crop_white(image: np.ndarray, value: int = 255) -> np.ndarray:
    """
        Crop white border from image
    :param image: Type: np.ndarray, image to be processed
    :param value: Type: int, default value is 255
    :return:
        Cropped image after removing the white border
    """
    assert (image.shape[2] == 3), "image shape should be [W, H, 3]"
    assert (image.dtype == np.uint8), "image type should be np.uint8"
    ys, = (image.min((1, 2)) < value).nonzero()
    xs, = (image.min(0).min(1) < value).nonzero()
    if len(xs) == 0 or len(ys) == 0:
        return image
    return image[ys.min(): ys.max()+1, xs.min(): xs.max()+1]


def get_tiles(image: np.ndarray, tile_size: int = 256, n_tiles: int = 36,
              mode: int = 0) -> Tuple[Dict[str, Any], bool]:
    """
        Crop big image to multiple pieces of small patches
    :param image: Type np.ndarray, image to be cropped
    :param tile_size: Type int, size of small patches
    :param n_tiles: Type int, number of small patches
    :param mode: Type int, pad type for cropping
    :return:
        dict includes small pacthes and its responding index, bool flag indicates if the
        image can get enough small patches
    """
    result = []
    h, w, c = image.shape
    pad_h = (tile_size - h % tile_size) % tile_size + ((tile_size * mode) // 2)
    pad_w = (tile_size - w % tile_size) % tile_size + ((tile_size * mode) // 2)

    img2 = np.pad(image, [[pad_h // 2, pad_h - pad_h // 2], [pad_w // 2, pad_w - pad_w//2], [0, 0]], constant_values=255)
    img3 = img2.reshape(
        img2.shape[0] // tile_size,
        tile_size,
        img2.shape[1] // tile_size,
        tile_size,
        3
    )

    img3 = img3.transpose(0, 2, 1, 3, 4).reshape(-1, tile_size, tile_size,3)
    n_tiles_with_info = (img3.reshape(img3.shape[0], -1).sum(1) < tile_size ** 2 * 3 * 255).sum()
    if len(img3) < n_tiles:
        img3 = np.pad(img3, [[0, n_tiles-len(img3)], [0, 0], [0, 0], [0, 0]], constant_values=255)
    idxs = np.argsort(img3.reshape(img3.shape[0], -1).sum(-1))[:n_tiles]
    img3 = img3[idxs]
    for i in range(len(img3)):
        result.append({'img': img3[i], 'idx': i})
    return result, n_tiles_with_info >= n_tiles


def glue_to_one_picture(image: np.ndarray, tile_size: int = 256, n_tiles: int = 36,
                        random_idx: bool = True) -> np.ndarray:
    """
        reorganize the distribution of images
    :param image: Type: np.ndarray, image to be processed
    :param tile_size: Type int, size of small patches
    :param n_tiles: Type int, number of small patches
    :param random_idx: Type bool, determine if the small patches are randomly organized
    :return:
        a image that is generated by reorganizing cropped patches from original image
    """
    tiles, _ = get_tiles(image, tile_size=tile_size, n_tiles=n_tiles, mode=0)

    if random_idx:
        patch_idxes = np.random.choice(list(range(n_tiles)), n_tiles, replace=False)
    else:
        patch_idxes = list(range(n_tiles))

    n_row_tiles = int(np.sqrt(n_tiles))
    images = np.zeros((tile_size * n_row_tiles, tile_size * n_row_tiles, 3)).astype(np.uint8)
    index = 0
    for h in range(n_row_tiles):
        for w in range(n_row_tiles):
            if len(tiles) > patch_idxes[index]:
                this_img = tiles[patch_idxes[index]]["img"]
                index = index + 1
            else:
                this_img = np.zeros((tile_size, tile_size, 3)).astype(np.uint8)
            images[h * tile_size:(h + 1) * tile_size, w * tile_size:(w + 1) * tile_size, :] = this_img

    return images


def cutmix(batch: List[Dict[str, Any]], hparams: Dict[str, Any]) -> Dict[str, Any]:
    """
        Apply cutmix transform for one batch of images
    :param batch: Type: dict, batch of dataset (default: {"image": image, "target": target}
    :param hparams: Type: config, config file
    :return:
        batch of dataset
    """
    image, target = batch["image"], batch["target"]
    batch_size = image.shape[0]
    img_h, img_w = image.shape[2:]
    imgs, labs = [], []

    for j in range(batch_size):
        p = np.random.uniform(0., 1.)
        if p >= hparams["cutmix_prob"]:
            idx = int(np.random.uniform(0, batch_size))
            # choose x, y and beta dist
            x = np.random.uniform(0, img_w)
            y = np.random.uniform(0, img_h)
            b = np.random.uniform(0., 1.)

            w = img_w * np.sqrt(1 - b)
            h = img_h * np.sqrt(1 - b)

            x0 = int(np.round(max(0, x - w/2)))
            x1 = int(np.round(min(img_w, x + w/2)))
            y0 = int(np.round(max(0, y - h/2)))
            y1 = int(np.round(min(img_h, y + h/2)))

            one = image[j, :, y0:y1, 0:x0]
            two = image[idx, :, y0:y1, x0:x1]
            three = image[j, :, y0:y1, x1: img_w]
            middle = torch.cat((one, two, three), dim=2)
            img = torch.cat((image[j, :, 0:y0, :], middle, image[j, :, y1:img_h, :]), dim=1)
            imgs.append(img)

            a = w * h / img_w / img_h

            if len(target.shape) < 2:
                if hparams.ohe_mode:
                    lab1 = F.one_hot(target[j], num_classes=hparams.num_class)
                    lab2 = F.one_hot(target[idx], num_classes=hparams.num_class)
                else:
                    lab1 = target[j]
                    lab2 = target[idx]
            else:
                lab1 = target[j, :]
                lab2 = target[idx, :]

            labs.append((1 - a) * lab1 + a * lab2)
        else:
            imgs.append(image[j, :, :, :])
            if len(target.shape) < 2:
                if hparams.ohe_mode:
                    labs.append(F.one_hot(target[j], num_classes=hparams.num_class).float())
                else:
                    labs.append(target[j])
            else:
                labs.append(target[j, :])

    image2 = torch.stack(imgs)
    label2 = torch.stack(labs)
    return {
        "image": image2,
        "target": label2,
    }


def mixup(batch: List[Dict[str, Any]], hparams: Dict[str, Any]) -> Dict[str, Any]:
    """
        Apply mixup transform for one batch of images
    :param batch: Type: dict, batch of dataset (default: {"image": image, "target": target}
    :param hparams: Type: config, config file
    :return:
        batch of dataset
    """
    image, target = batch["image"], batch["target"]
    batch_size = image.shape[0]
    imgs, labs = [], []

    for j in range(batch_size):
        p = np.random.uniform(0., 1.)
        if p >= hparams["mixup_prob"]:
            idx = int(np.random.uniform(0, batch_size))
            # choose beta dist
            b = np.random.uniform(0., 1.)

            img = (1 - b) * image[j, :, :, :] + b * image[idx, :, :, :]
            imgs.append(img)

            if len(target.shape) < 2:
                if hparams.ohe_mode:
                    lab1 = F.one_hot(target[j], num_classes=hparams.num_class)
                    lab2 = F.one_hot(target[idx], num_classes=hparams.num_class)
                else:
                    lab1 = target[j]
                    lab2 = target[idx]
            else:
                lab1 = target[j, :]
                lab2 = target[idx, :]

            labs.append((1 - b) * lab1 + b * lab2)
        else:
            imgs.append(image[j, :, :, :])
            if len(target.shape) < 2:
                if hparams.ohe_mode:
                    labs.append(F.one_hot(target[j], num_classes=hparams.num_class).float())
                else:
                    labs.append(target[j])
            else:
                labs.append(target[j, :])

    image2 = torch.stack(imgs)
    label2 = torch.stack(labs)
    return {
        "image": image2,
        "target": label2,
    }


class MixCollator:
    def __init__(self, hparams: Dict[str, Any]):
        super(MixCollator, self).__init__()
        self.hparams = hparams

    def __call__(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
        batch = torch.utils.data.dataloader.default_collate(batch)
        if self.hparams["mix_aug"]["cutmix"]:
            batch = cutmix(batch, self.hparams)
        if self.hparams["mix_aug"]["mixup"]:
            batch = mixup(batch, self.hparams)

        return batch





