# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import random
import traceback
import numpy as np
from PIL import Image

import logging
logger = logging.getLogger(__name__)

__all__ = ['GroupScale', 'GroupMultiScaleCrop', 'GroupRandomCrop',
           'GroupRandomFlip', 'GroupCenterCrop', 'NormalizeImage',
           'Compose']


class Compose(object):
    def __init__(self, transforms=[]):
        self.transforms = transforms

    def __call__(self, *data):
        for f in self.transforms:
            try:
                data = f(*data)
            except Exception as e:
                stack_info = traceback.format_exc()
                logger.info("fail to perform transform [{}] with error: "
                        "{} and stack:\n{}".format(f, e, str(stack_info)))
                raise e
        return data


class GroupScale(object):
    """
    Group scale image

    Args:
        target_size (int): image resize target size
    """
    def __init__(self, target_size=224):
        self.target_size = target_size

    def __call__(self, imgs, label):
        resized_imgs = []
        for i in range(len(imgs)):
            img = imgs[i]
            w, h = img.size
            if (w <= h and w == self.target_size) or \
                    (h <= w and h == self.target_size):
                resized_imgs.append(img)
                continue

            if w < h:
                ow = self.target_size
                oh = int(self.target_size * 4.0 / 3.0)
                resized_imgs.append(img.resize((ow, oh), Image.BILINEAR))
            else:
                oh = self.target_size
                ow = int(self.target_size * 4.0 / 3.0)
                resized_imgs.append(img.resize((ow, oh), Image.BILINEAR))
        
        return resized_imgs, label


class GroupMultiScaleCrop(object):
    """
    FIXME: add comments
    """
    def __init__(self,
                 short_size=256,
                 scales=None,
                 max_distort=1,
                 fix_crop=True,
                 more_fix_crop=True):
        self.short_size = short_size
        self.scales = scales if scales is not None \
                        else [1, .875, .75, .66]
        self.max_distort = max_distort
        self.fix_crop = fix_crop
        self.more_fix_crop = more_fix_crop

    def __call__(self, imgs, label):
        input_size = [self.short_size, self.short_size]
        im_size = imgs[0].size
        
        # get random crop offset
        def _sample_crop_size(im_size):
            image_w, image_h = im_size[0], im_size[1]
            
            base_size = min(image_w, image_h)
            crop_sizes = [int(base_size * x) for x in self.scales]
            crop_h = [
                input_size[1] if abs(x - input_size[1]) < 3 else x
                for x in crop_sizes
            ]
            crop_w = [
                input_size[0] if abs(x - input_size[0]) < 3 else x
                for x in crop_sizes
            ]
            
            pairs = []
            for i, h in enumerate(crop_h):
                for j, w in enumerate(crop_w):
                    if abs(i - j) <= self.max_distort:
                        pairs.append((w, h))
            crop_pair = random.choice(pairs)
            if not self.fix_crop:
                w_offset = np.random.randint(0, image_w - crop_pair[0])
                h_offset = np.random.randint(0, image_h - crop_pair[1])
            else:
                w_step = (image_w - crop_pair[0]) / 4
                h_step = (image_h - crop_pair[1]) / 4
                
                ret = list()
                ret.append((0, 0))  # upper left
                if w_step != 0:
                    ret.append((4 * w_step, 0))  # upper right
                if h_step != 0:
                    ret.append((0, 4 * h_step))  # lower left
                if h_step != 0 and w_step != 0:
                    ret.append((4 * w_step, 4 * h_step))  # lower right
                if h_step != 0 or w_step != 0:
                    ret.append((2 * w_step, 2 * h_step))  # center
                
                if self.more_fix_crop:
                    ret.append((0, 2 * h_step))  # center left
                    ret.append((4 * w_step, 2 * h_step))  # center right
                    ret.append((2 * w_step, 4 * h_step))  # lower center
                    ret.append((2 * w_step, 0 * h_step))  # upper center
                    
                    ret.append((1 * w_step, 1 * h_step))  # upper left quarter
                    ret.append((3 * w_step, 1 * h_step))  # upper right quarter
                    ret.append((1 * w_step, 3 * h_step))  # lower left quarter
                    ret.append((3 * w_step, 3 * h_step))  # lower righ quarter
                
                w_offset, h_offset = random.choice(ret)
            
            return crop_pair[0], crop_pair[1], w_offset, h_offset
        
        crop_w, crop_h, offset_w, offset_h = _sample_crop_size(im_size)
        crop_imgs = [
            img.crop((offset_w, offset_h, offset_w + crop_w, offset_h + crop_h))
            for img in imgs
        ]
        ret_imgs = [
            img.resize((input_size[0], input_size[1]), Image.BILINEAR)
            for img in crop_imgs
        ]

        return ret_imgs, label


class GroupRandomCrop(object):
    def __init__(self, target_size=224):
        self.target_size = target_size

    def __call__(self, imgs, label):
        w, h = imgs[0].size
        th, tw = self.target_size, self.target_size
        
        assert (w >= self.target_size) and (h >= self.target_size), \
            "image width({}) and height({}) should be larger than " \
            "crop size".format(w, h, self.target_size)
        
        out_images = []
        x1 = np.random.randint(0, w - tw)
        y1 = np.random.randint(0, h - th)
        
        for img in imgs:
            if w == tw and h == th:
                out_images.append(img)
            else:
                out_images.append(img.crop((x1, y1, x1 + tw, y1 + th)))
        
        return out_images, label


class GroupRandomFlip(object):
    def __call__(self, imgs, label):
        v = np.random.random()
        if v < 0.5:
            ret = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in imgs]
            return ret, label
        else:
            return imgs, label 


class GroupCenterCrop(object):
    def __init__(self, target_size=224):
        self.target_size = target_size

    def __call__(self, imgs, label):
        crop_imgs = []
        for img in imgs:
            w, h = img.size
            th, tw = self.target_size, self.target_size
            assert (w >= self.target_size) and (h >= self.target_size), \
                "image width({}) and height({}) should be larger " \
                "than crop size".format(w, h, self.target_size)
            x1 = int(round((w - tw) / 2.))
            y1 = int(round((h - th) / 2.))
            crop_imgs.append(img.crop((x1, y1, x1 + tw, y1 + th)))
        
        return crop_imgs, label 


class NormalizeImage(object):
    def __init__(self,
                 target_size=224,
                 img_mean=[0.485, 0.456, 0.406],
                 img_std=[0.229, 0.224, 0.225],
                 seg_num=8,
                 seg_len=1):
        self.target_size = target_size
        self.img_mean = np.array(img_mean).reshape((3, 1, 1)).astype('float32')
        self.img_std = np.array(img_std).reshape((3, 1, 1)).astype('float32')
        self.seg_num = seg_num
        self.seg_len = seg_len

    def __call__(self, imgs, label):
        np_imgs = (np.array(imgs[0]).astype('float32').transpose(
            (2, 0, 1))).reshape(1, 3, self.target_size,
            self.target_size) / 255
        for i in range(len(imgs) - 1):
            img = (np.array(imgs[i + 1]).astype('float32').transpose(
                (2, 0, 1))).reshape(1, 3, self.target_size,
                self.target_size) / 255
            np_imgs = np.concatenate((np_imgs, img))
        
        np_imgs -= self.img_mean
        np_imgs /= self.img_std
        np_imgs = np.reshape(np_imgs, (self.seg_num, self.seg_len * 3,
                             self.target_size, self.target_size))
        
        return np_imgs, label