"""
This file is inherited from tsn-pytorch
"""

import torchvision
import random
from PIL import Image, ImageOps
import numpy as np
import numbers
import math
import torch
from torchvision.transforms import functional as F

# lstr
class Compose(object):

    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, data):
        for t in self.transforms:
            data = t(data)
        return data

#LSTR
class GroupCrop(object):
    """perfrom group crop'"""

    def __init__(self, crop_size, crop_pos='random'):
        super(GroupCrop, self).__init__()
        self.crop_size = crop_size
        self.crop_pos = crop_pos
        assert crop_pos in ['random', 'top_left', 'center', 'bottom_right'
                           ], f'crop_pos: {self.crop_pos} not implemented'

    def __call__(self, img_group):
        """
        Args:
            img_group (List(PIL.Image) or List(torch.Tensor)): a list of image to crop.
        Returns: The same type as input. cropped img_group.
            
        """
        crop_w, crop_h = self.crop_size
        if isinstance(img_group, torch.Tensor):
            assert img_group.dim(
            ) == 4, f'input img_group should be dim-4(TCHW). Got shape: {img_group.shape}'
            img_h, img_w = img_group.shape[-2:]
            offset_h, offset_w = self.get_crop_offset(img_w, img_h, crop_w,
                                                        crop_h)
            return F.crop(img_group, offset_h, offset_w, crop_h, crop_w)
        elif isinstance(img_group[0], Image):
            img_w, img_h = img_group[0].size
            offset_h, offset_w = self.get_crop_offset(img_w, img_h, crop_w,
                                                        crop_h)
            return [
                F.crop(img, offset_h, offset_w, crop_h, crop_w)
                for img in img_group
            ]
        else:
            raise NotImplementedError(
                f'img_group of type: {type(img_group)} is not supported')

    def get_crop_offset(self, img_w, img_h, crop_w, crop_h):
        """get the offset to crop.
        Returns: (offset_h, offset_w).
        """
        if self.crop_pos == 'random':
            offset_h = int(torch.randint(img_h - crop_h, []))
            offset_w = int(torch.randint(img_w - crop_w, []))
        elif self.crop_pos == 'top_left':
            offset_h = offset_w = 0
        elif self.crop_pos == 'center':
            offset_h = int((img_h - crop_h) // 2)
            offset_w = int((img_w - crop_w) // 2)
        elif self.crop_pos == 'bottom_right':
            offset_h = img_h - crop_h
            offset_w = img_w - crop_w
        else:
            raise NotImplementedError(
                f'crop_pos: {self.crop_pos} not implemented')
        return (offset_h, offset_w)

class GroupRandomCrop(object):
    def __init__(self, size):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size

    def __call__(self, img_group):

        w, h = img_group[0].size
        th, tw = self.size

        out_images = list()

        x1 = random.randint(0, w - tw)
        y1 = random.randint(0, h - th)

        for img in img_group:
            assert(img.size[0] == w and img.size[1] == h)
            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


class GroupCenterCrop(object):
    def __init__(self, size):
        self.worker = torchvision.transforms.CenterCrop(size)

    def __call__(self, img_group):
        return [self.worker(img) for img in img_group]


class GroupRandomHorizontalFlip(object):
    """Randomly horizontally flips the given PIL.Image with a probability of 0.5
    """
    def __init__(self, is_flow=False):
        self.is_flow = is_flow

    def __call__(self, img_group, is_flow=False):
        v = random.random()
        if v < 0.5:
            ret = [img.transpose(Image.FLIP_LEFT_RIGHT) for img in img_group]
            if self.is_flow:
                for i in range(0, len(ret), 2):
                    ret[i] = ImageOps.invert(ret[i])  # invert flow pixel values when flipping
            return ret
        else:
            return img_group


# class GroupNormalize(object):
#     def __init__(self, mean, std):
#         self.mean = mean
#         self.std = std

#     def __call__(self, tensor):
#         rep_mean = self.mean * (tensor.size()[0]//len(self.mean))
#         rep_std = self.std * (tensor.size()[0]//len(self.std))

#         # TODO: make efficient
#         for t, m, s in zip(tensor, rep_mean, rep_std):
#             t.sub_(m).div_(s)

#         return tensor

# lstr
class GroupNormalize(object):

    def __init__(self, mean, std):

        if len(mean) == 1:    # flow
            assert len(
                std
            ) == 1, f'flow input must have std of length 1. got {len(std)}'
            self.mean = mean[0]
            self.std = std[0]
        else:
            mean = np.array(mean, dtype=np.float32)
            self.mean = mean[np.newaxis, :, np.newaxis,
                                np.newaxis]    # shape: (C,) ->(T,C,H,W)
            std = np.array(std, dtype=np.float32)
            self.std = std[np.newaxis, :, np.newaxis,
                            np.newaxis]    # shape: (C,) ->(T,C,H,W)

    def __call__(self, tensor):
        """
        Args:
            tensor (torch.Tensor) : The stacked images. shape: (T, C, H, W).
        Returns: (torch.Tensor).  sub mean and div std along the channel dim.
            
        """
        return (tensor - self.mean) / self.std

class GroupScale(object):
    """ Rescales the input PIL.Image to the given 'size'.
    'size' will be the size of the smaller edge.
    For example, if height > width, then image will be
    rescaled to (size * height / width, size)
    size: size of the smaller edge
    interpolation: Default: PIL.Image.BILINEAR
    """

    def __init__(self, size, interpolation=Image.BILINEAR):
        self.worker = torchvision.transforms.Resize(size, interpolation)

    def __call__(self, img_group):
        return [self.worker(img) for img in img_group]

# lstr
class GroupResize(object):
    """perfrom group random resize"""

    def __init__(self, scale_range, keep_aspect_ratio=True):
        """
        Args:
            scale_range (tuple/list): The shorter side of the image will be resized to
                                        1) the range if `scale_range` is a tuple
                                        2) the specific size if `scale_range` contains only 1 element.
        """
        super(GroupResize, self).__init__()
        self.scale_range = scale_range
        self.keep_aspect_ratio = keep_aspect_ratio

    def __call__(self, img_group):
        """
        Args:
            img_group (List(PIL.Image) or Tensor([T,C,H,W])): a list of image to resize.
        Returns: The same type as input. resized img_group.
            
        """
        # resize
        size = self.scale_range[0] if len(
            self.scale_range) == 1 else torch.randint(self.scale_range[0],
                                                        self.scale_range[1], [])
        size = int(size)
        size = size if self.keep_aspect_ratio else (size, size)

        if isinstance(img_group, torch.Tensor):
            assert img_group.dim(
            ) == 4, f'input img_group should be dim-4(TCHW). Got shape: {img_group.shape}'
            return F.resize(img_group, size)
        elif isinstance(img_group[0], Image):
            return [F.resize(img, size) for img in img_group]
        else:
            raise NotImplementedError(
                f'img_group of type: {type(img_group)} is not supported')

class GroupOverSample(object):
    def __init__(self, crop_size, scale_size=None):
        self.crop_size = crop_size if not isinstance(crop_size, int) else (crop_size, crop_size)

        if scale_size is not None:
            self.scale_worker = GroupScale(scale_size)
        else:
            self.scale_worker = None

    def __call__(self, img_group):

        if self.scale_worker is not None:
            img_group = self.scale_worker(img_group)
        image_w, image_h = img_group[0].size
        crop_w, crop_h = self.crop_size

        offsets = GroupMultiScaleCrop.fill_fix_offset(False, image_w, image_h, crop_w, crop_h)
        oversample_group = list()
        for o_w, o_h in offsets:
            normal_group = list()
            flip_group = list()
            for i, img in enumerate(img_group):
                crop = img.crop((o_w, o_h, o_w + crop_w, o_h + crop_h))
                normal_group.append(crop)
                flip_crop = crop.copy().transpose(Image.FLIP_LEFT_RIGHT)

                if img.mode == 'L' and i % 2 == 0:
                    flip_group.append(ImageOps.invert(flip_crop))
                else:
                    flip_group.append(flip_crop)

            oversample_group.extend(normal_group)
            oversample_group.extend(flip_group)
        return oversample_group


class GroupMultiScaleCrop(object):

    def __init__(self, input_size, scales=None, max_distort=1, fix_crop=True, more_fix_crop=True):
        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
        self.input_size = input_size if not isinstance(input_size, int) else [input_size, input_size]
        self.interpolation = Image.BILINEAR

    def __call__(self, img_group):

        im_size = img_group[0].size

        crop_w, crop_h, offset_w, offset_h = self._sample_crop_size(im_size)
        crop_img_group = [img.crop((offset_w, offset_h, offset_w + crop_w, offset_h + crop_h)) for img in img_group]
        ret_img_group = [img.resize((self.input_size[0], self.input_size[1]), self.interpolation)
                         for img in crop_img_group]
        return ret_img_group

    def _sample_crop_size(self, im_size):
        image_w, image_h = im_size[0], im_size[1]

        # find a crop size
        base_size = min(image_w, image_h)
        crop_sizes = [int(base_size * x) for x in self.scales]
        crop_h = [self.input_size[1] if abs(x - self.input_size[1]) < 3 else x for x in crop_sizes]
        crop_w = [self.input_size[0] if abs(x - self.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 = random.randint(0, image_w - crop_pair[0])
            h_offset = random.randint(0, image_h - crop_pair[1])
        else:
            w_offset, h_offset = self._sample_fix_offset(image_w, image_h, crop_pair[0], crop_pair[1])

        return crop_pair[0], crop_pair[1], w_offset, h_offset

    def _sample_fix_offset(self, image_w, image_h, crop_w, crop_h):
        offsets = self.fill_fix_offset(self.more_fix_crop, image_w, image_h, crop_w, crop_h)
        return random.choice(offsets)

    @staticmethod
    def fill_fix_offset(more_fix_crop, image_w, image_h, crop_w, crop_h):
        w_step = (image_w - crop_w) // 4
        h_step = (image_h - crop_h) // 4

        ret = list()
        ret.append((0, 0))  # upper left
        ret.append((4 * w_step, 0))  # upper right
        ret.append((0, 4 * h_step))  # lower left
        ret.append((4 * w_step, 4 * h_step))  # lower right
        ret.append((2 * w_step, 2 * h_step))  # center

        if 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

        return ret


class GroupRandomSizedCrop(object):
    """Random crop the given PIL.Image to a random size of (0.08 to 1.0) of the original size
    and and a random aspect ratio of 3/4 to 4/3 of the original aspect ratio
    This is popularly used to train the Inception networks
    size: size of the smaller edge
    interpolation: Default: PIL.Image.BILINEAR
    """
    def __init__(self, size, interpolation=Image.BILINEAR):
        self.size = size
        self.interpolation = interpolation

    def __call__(self, img_group):
        for attempt in range(10):
            area = img_group[0].size(0) * img_group[0].size(1)
            target_area = random.uniform(0.08, 1.0) * area
            aspect_ratio = random.uniform(3. / 4, 4. / 3)

            w = int(round(math.sqrt(target_area * aspect_ratio)))
            h = int(round(math.sqrt(target_area / aspect_ratio)))

            if random.random() < 0.5:
                w, h = h, w

            if w <= img_group[0].size(0) and h <= img_group[0].size(1):
                x1 = random.randint(0, img_group[0].size(0) - w)
                y1 = random.randint(0, img_group[0].size(1) - h)
                found = True
                break
        else:
            found = False
            x1 = 0
            y1 = 0

        if found:
            out_group = list()
            for img in img_group:
                img = img.crop((x1, y1, x1 + w, y1 + h))
                assert(img.size == (w, h))
                out_group.append(img.resize((self.size, self.size), self.interpolation))
            return out_group
        else:
            # Fallback
            scale = GroupScale(self.size, interpolation=self.interpolation)
            crop = GroupRandomCrop(self.size)
            return crop(scale(img_group))


class Stack(object):

    def __init__(self, roll=False):
        self.roll = roll

    def __call__(self, img_group):
        if img_group[0].mode == 'L':
            return np.concatenate([np.expand_dims(x, 2) for x in img_group], axis=2)
        elif img_group[0].mode == 'RGB':
            if self.roll:
                return np.concatenate([np.array(x)[:, :, ::-1] for x in img_group], axis=2)
            else:
                return np.concatenate(img_group, axis=2)

# LSTR                
class GroupStack(object):
    """Stack to numpy array and convert  THWC to TCHW"""

    def __call__(self, img_group):
        """
        Args:
            img_group (List(PIL.Image)): a list of image. image shape: (H,W,C) for rgb, (H,W) for grayscale
        Returns: np.ndarray. stack the images along the first axis.
                shape: (T, C, H, W).
        """
        if isinstance(img_group, (list, tuple)):
            imgs = np.stack(img_group, axis=0)
        else:
            imgs = img_group
        if imgs.ndim == 3:    # THW
            imgs = imgs[:, np.newaxis, :, :]    # TCHW
        else:    # THWC
            imgs = np.transpose(imgs, [0, 3, 1, 2])    # THWC -> TCHW
        return imgs

class ToTorchFormatTensor(object):
    """ Converts a PIL.Image (RGB) or numpy.ndarray (H x W x C) in the range [0, 255]
    to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] """
    def __init__(self, div=True):
        self.div = div

    def __call__(self, pic):
        if isinstance(pic, np.ndarray):
            # handle numpy array
            img = torch.from_numpy(pic).permute(2, 0, 1).contiguous()
        else:
            # handle PIL Image
            img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
            img = img.view(pic.size[1], pic.size[0], len(pic.mode))
            # put it from HWC to CHW format
            # yikes, this transpose takes 80% of the loading time/CPU
            img = img.transpose(0, 1).transpose(0, 2).contiguous()
        return img.float().div(255) if self.div else img.float()

# LSTR
class GroupToTensor(object):
    
    def __init__(self, div=True):
        self.div = div

    def __call__(self, imgs):
        """
        Args:
            imgs (np.ndarray) : The stacked images. shape: (T, C, H, W). dtype: np.uint8
        Returns: (torch.Tensor). convert dtype to torch.float32.
                If `div` is True, the images are divided by 255 to be normalized to [0,1]
            
        """
        imgs = torch.from_numpy(imgs).float()
        if self.div:
            imgs = imgs.div(255.)
        return imgs

class IdentityTransform(object):

    def __call__(self, data):
        return data


if __name__ == "__main__":
    trans = torchvision.transforms.Compose([
        GroupScale(256),
        GroupRandomCrop(224),
        Stack(),
        ToTorchFormatTensor(),
        GroupNormalize(
            mean=[.485, .456, .406],
            std=[.229, .224, .225]
        )]
    )

    im = Image.open('../tensorflow-model-zoo.torch/lena_299.png')

    color_group = [im] * 3
    rst = trans(color_group)

    gray_group = [im.convert('L')] * 9
    gray_rst = trans(gray_group)

    trans2 = torchvision.transforms.Compose([
        GroupRandomSizedCrop(256),
        Stack(),
        ToTorchFormatTensor(),
        GroupNormalize(
            mean=[.485, .456, .406],
            std=[.229, .224, .225])
    ])
    print(trans2(color_group))