import math
import numbers
import random
import numpy as np
import torchvision.transforms.functional as tf

from PIL import Image, ImageOps

class Compose(object):
    def __init__(self, augmentations):
        self.augmentations = augmentations
        self.PIL2Numpy = False

    def __call__(self, img, mask):
        if isinstance(img, np.ndarray):
            img = Image.fromarray(img, mode="I")
            mask = Image.fromarray(mask, mode="L")
            self.PIL2Numpy = True

        assert img.size == mask.size
        for a in self.augmentations:
            img, mask = a(img, mask)

        if self.PIL2Numpy:
            img, mask = np.array(img), np.array(mask, dtype=np.uint8)

        return img, mask

class ComposeTest(object):
    def __init__(self, augmentations):
        self.augmentations = augmentations
        self.PIL2Numpy = False

    def __call__(self, img):
        if isinstance(img, np.ndarray):
            img = Image.fromarray(img, mode="I")
            self.PIL2Numpy = True

        for a in self.augmentations:
            img = a(img)

        if self.PIL2Numpy:
            img = np.array(img)

        return img

class RandomErasing(object):
    def __init__(self, probability=0.5, sl=0.02, sh=0.4, r1=0.3, mean=[0.4914, 0.4822, 0.4465]):
        self.probability = probability
        self.mean = mean
        self.sl = sl
        self.sh = sh
        self.r1 = r1

    def __call__(self, img):
        
        if random.uniform(0, 1) > self.probability:
            return img
        
        for attempt in range(100):
            area = img.size()[1]*img.size()[2]

            target_area = random.uniform(self.sl, self.sh) * area
            aspect_ratio = random.uniform(self.r1, 1/self.r1)

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

            if w < img.size()[2] and h < img.size()[1]:
                x1 = random.randint(0, img.size()[1] - h)
                y1 = random.randint(0, img.size()[2] - w)
                if img.size()[0] == 3:
                    img[0, x1:x1+h, y1:y1+w] = self.mean[0]
                    img[1, x1:x1+h, y1:y1+w] = self.mean[1]
                    img[2, x1:x1+h, y1:y1+w] = self.mean[2]
                else:
                    img[0, x1:x1+h, y1:y1+w] = self.mean[0]
                return img

        return img

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

    def __call__(self, img, mask):
        if self.padding > 0:
            img = ImageOps.expand(img, border=self.padding, fill=0)
            mask = ImageOps.expand(mask, border=self.padding, fill=0)

        assert img.size == mask.size
        w, h = img.size
        th, tw = self.size
        if w == tw and h == th:
            return img, mask
        if w < tw or h < th:
            return (
                img.resize((tw, th), Image.BILINEAR),
                mask.resize((tw, th), Image.NEAREST),
            )

        x1 = random.randint(0, w - tw)
        y1 = random.randint(0, h - th)
        return (
            img.crop((x1, y1, x1 + tw, y1 + th)),
            mask.crop((x1, y1, x1 + tw, y1 + th)),
        )


class AdjustGamma(object):
    def __init__(self, gamma):
        self.gamma = gamma

    def __call__(self, img, mask):
        assert img.size == mask.size
        return tf.adjust_gamma(img, random.uniform(1, 1 + self.gamma)), mask


class AdjustSaturation(object):
    def __init__(self, saturation):
        self.saturation = saturation

    def __call__(self, img, mask):
        assert img.size == mask.size
        return tf.adjust_saturation(img,
                                    random.uniform(1 - self.saturation,
                                                   1 + self.saturation)), mask


class AdjustHue(object):
    def __init__(self, hue):
        self.hue = hue

    def __call__(self, img, mask):
        assert img.size == mask.size
        return tf.adjust_hue(img, random.uniform(-self.hue,
                                                 self.hue)), mask


class AdjustBrightness(object):
    def __init__(self, bf):
        self.bf = bf

    def __call__(self, img, mask):
        assert img.size == mask.size
        if isinstance(img, Image.Image):
            img = np.array(img)

        if img.ndim == 2:
            img = np.expand_dims(img, 2)
            img = np.concatenate((img, img, img), 2)
            img = Image.fromarray(img, "RGB")

        img = tf.adjust_brightness(img, random.uniform(1-self.bf, 1+self.bf))
        img = np.array(img)
        img = Image.fromarray(img, "RGB").convert("L")
        return img, mask
    
    def adjust_brightness(self, img, bf):
        if isinstance(img, Image.Image):
            img = np.array(img)
        
        if np.amax(img) > 1.0:
            img = img / 255

        img = np.clip(img*bf, 0.0, 1.0)
        img = img * 255
        return Image.fromarray(img.astype(np.uint32), mode="I")

class AdjustContrast(object):
    def __init__(self, cf):
        self.cf = cf

    def __call__(self, img, mask):
        assert img.size == mask.size 
        if isinstance(img, Image.Image):
            img = np.array(img)

        if img.ndim == 2:
            img = np.expand_dims(img, 2)
            img = np.concatenate((img, img, img), 2)
            img = Image.fromarray(img, "RGB")
        
        img = tf.adjust_contrast(img, random.uniform(1-self.cf, 1+self.cf))
        img = np.array(img)
        img = Image.fromarray(img, "RGB").convert("L")
        return img, mask

    def adjust_contrast(self, img, cf):
        if isinstance(img, Image.Image):
            img = np.array(img)
        
        if np.amax(img) > 1.0:
            img = img/255

        img = np.clip(0.5+cf*(img-0.5), 0.0, 1.0)
        img = img * 255
        return Image.fromarray(img.astype(np.uint32), mode="I")

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

    def __call__(self, img, mask):
        assert img.size == mask.size
        w, h = img.size
        th, tw = self.size
        x1 = int(round((w - tw) / 2.))
        y1 = int(round((h - th) / 2.))
        return (
            img.crop((x1, y1, x1 + tw, y1 + th)),
            mask.crop((x1, y1, x1 + tw, y1 + th)),
        )


class PaddingCenterCrop(object):
    def __init__(self, size):
        '''
        Argument crop_size is an integer for square cropping only.

        Performs padding and center cropping to a specified size.
        '''
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size

    def __call__(self, img, mask):
        assert img.size == mask.size
        w, h = img.size
        th, tw = self.size
        if w >= tw and h >= th:  # crop a center patch
            x1 = int(round((w - tw) / 2.))
            y1 = int(round((h - th) / 2.))
            return (img.crop((x1, y1, x1 + tw, y1 + th)),
                    mask.crop((x1, y1, x1 + tw, y1 + th)))
        else:  # pad zeros and do center crop
            pad_h = max(th - h, 0)
            pad_w = max(tw - w, 0)
            rem_h = pad_h % 2
            rem_w = pad_w % 2

            # padding lengh
            left = pad_w // 2
            top = pad_h // 2
            right = pad_w // 2 + rem_w
            bottom = pad_h // 2 + rem_h

            border = (left, top, right, bottom)  # left, top, right, bottom
            img = ImageOps.expand(img, border, fill=0)
            mask = ImageOps.expand(mask, border, fill=0)

            # do center crop
            x1 = max(int(round((w - tw) / 2.)), 0)
            y1 = max(int(round((h - th) / 2.)), 0)
            return (img.crop((x1, y1, x1 + tw, y1 + th)),
                    mask.crop((x1, y1, x1 + tw, y1 + th)))

class PaddingCenterCropTest(object):
    def __init__(self, size):
        '''
        Argument crop_size is an integer for square cropping only.

        Performs padding and center cropping to a specified size.
        '''
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size

    def __call__(self, img):
        w, h = img.size
        tw, th = self.size
        if w >= tw and h >= th:  # crop a center patch
            x1 = int(self.round_num((w - tw) / 2.))
            y1 = int(self.round_num((h - th) / 2.))
            return img.crop((x1, y1, x1 + tw, y1 + th))
        else:  # pad zeros and do center crop
            pad_h = max(th - h, 0)
            pad_w = max(tw - w, 0)
            rem_h = pad_h % 2
            rem_w = pad_w % 2

            # padding lengh
            left = pad_w // 2
            top = pad_h // 2
            right = pad_w // 2 + rem_w
            bottom = pad_h // 2 + rem_h

            border = (left, top, right, bottom)  # left, top, right, bottom
            img = ImageOps.expand(img, border, fill=0)

            # do center crop
            x1 = max(int(self.round_num((w - tw) / 2.)), 0)
            y1 = max(int(self.round_num((h - th) / 2.)), 0)
            return img.crop((x1, y1, x1 + tw, y1 + th))
    
    def round_num(self, x):
        return int(x) + 1 if (x-int(x)) >= 0.5 else int(x)

class RandomHorizontallyFlip(object):
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img, mask):
        if random.random() < self.p:
            return (
                img.transpose(Image.FLIP_LEFT_RIGHT),
                mask.transpose(Image.FLIP_LEFT_RIGHT),
            )
        return img, mask


class RandomVerticallyFlip(object):
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, img, mask):
        if random.random() < self.p:
            return (
                img.transpose(Image.FLIP_TOP_BOTTOM),
                mask.transpose(Image.FLIP_TOP_BOTTOM),
            )
        return img, mask


class FreeScale(object):
    def __init__(self, size):
        self.size = tuple(reversed(size))  # size: (h, w)

    def __call__(self, img, mask):
        assert img.size == mask.size
        return (
            img.resize(self.size, Image.BILINEAR),
            mask.resize(self.size, Image.NEAREST),
        )


class RandomTranslate(object):
    def __init__(self, offset):
        self.offset = offset  # tuple (delta_x, delta_y)

    def __call__(self, img, mask):
        assert img.size == mask.size
        x_offset = int(2 * (random.random() - 0.5) * self.offset[0])
        y_offset = int(2 * (random.random() - 0.5) * self.offset[1])

        x_crop_offset = x_offset
        y_crop_offset = y_offset
        if x_offset < 0:
            x_crop_offset = 0
        if y_offset < 0:
            y_crop_offset = 0

        cropped_img = tf.crop(img,
                              y_crop_offset,
                              x_crop_offset,
                              img.size[1] - abs(y_offset),
                              img.size[0] - abs(x_offset))

        if x_offset >= 0 and y_offset >= 0:
            padding_tuple = (0, 0, x_offset, y_offset)

        elif x_offset >= 0 and y_offset < 0:
            padding_tuple = (0, abs(y_offset), x_offset, 0)

        elif x_offset < 0 and y_offset >= 0:
            padding_tuple = (abs(x_offset), 0, 0, y_offset)

        elif x_offset < 0 and y_offset < 0:
            padding_tuple = (abs(x_offset), abs(y_offset), 0, 0)

        return (
            tf.pad(cropped_img,
                   padding_tuple,
                   padding_mode='reflect'),
            tf.affine(mask,
                      translate=(-x_offset, -y_offset),
                      scale=1.0,
                      angle=0.0,
                      shear=0.0,
                      fillcolor=250))


class RandomRotate(object):
    def __init__(self, degree):
        self.degree = degree

    def __call__(self, img, mask):
        rotate_degree = random.random() * 2 * self.degree - self.degree
        return (
            tf.affine(img,
                      translate=(0, 0),
                      scale=1.0,
                      angle=rotate_degree,
                      resample=Image.BILINEAR,
                      fillcolor=0,
                      shear=0.0),
            tf.affine(mask,
                      translate=(0, 0),
                      scale=1.0,
                      angle=rotate_degree,
                      resample=Image.NEAREST,
                      fillcolor=0,
                      shear=0.0))


class Scale(object):
    def __init__(self, size=512, random_scale=False):
        self.size = size
        self.random_scale = random_scale

    def __call__(self, img, mask):
        assert img.size == mask.size
        w, h = img.size
        if self.random_scale:
            new_w = int(random.uniform(0.5, 2) * w)
            new_h = int(random.uniform(0.5, 2) * h)
            return (
                img.resize((new_w, new_h), Image.BILINEAR),
                mask.resize((new_w, new_h), Image.NEAREST),
            )
        if (w >= h and w == self.size) or (h >= w and h == self.size):
            return img, mask
        if w > h:
            ow = self.size
            oh = int(self.size * h / w)
            return (
                img.resize((ow, oh), Image.BILINEAR),
                mask.resize((ow, oh), Image.NEAREST),
            )
        else:
            oh = self.size
            ow = int(self.size * w / h)
            return (
                img.resize((ow, oh), Image.BILINEAR),
                mask.resize((ow, oh), Image.NEAREST),
            )


class RandomSizedCrop(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, img, mask):
        assert img.size == mask.size
        for attempt in range(10):
            area = img.size[0] * img.size[1]
            target_area = random.uniform(0.45, 1.0) * area
            aspect_ratio = random.uniform(0.5, 2)

            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.size[0] and h <= img.size[1]:
                x1 = random.randint(0, img.size[0] - w)
                y1 = random.randint(0, img.size[1] - h)

                img = img.crop((x1, y1, x1 + w, y1 + h))
                mask = mask.crop((x1, y1, x1 + w, y1 + h))
                assert img.size == (w, h)

                return (
                    img.resize((self.size, self.size), Image.BILINEAR),
                    mask.resize((self.size, self.size), Image.NEAREST),
                )

        # Fallback
        scale = Scale(self.size)
        crop = CenterCrop(self.size)
        return crop(*scale(img, mask))


class RandomSized(object):
    def __init__(self, size):
        self.size = size
        self.scale = Scale(self.size)
        self.crop = RandomCrop(self.size)

    def __call__(self, img, mask):
        assert img.size == mask.size

        w = int(random.uniform(0.5, 2) * img.size[0])
        h = int(random.uniform(0.5, 2) * img.size[1])

        img, mask = (
            img.resize((w, h), Image.BILINEAR),
            mask.resize((w, h), Image.NEAREST),
        )

        return self.crop(*self.scale(img, mask))
