import numpy as np
from PIL import Image

import cv2

cv2.setNumThreads(0)
cv2.ocl.setUseOpenCL(False)

from torchvision.transforms import ColorJitter


class FlowAugmentor:
    def __init__(
        self,
        crop_size,
        min_scale=-0.2,
        max_scale=0.5,
        do_flip=True,
        pwc_aug=False,
        do_rotate=False,
        pre_scale=0.0,
    ):
        # spatial augmentation params
        self.crop_size = crop_size
        self.pre_scale = pre_scale
        self.min_scale = min_scale
        self.max_scale = max_scale
        self.spatial_aug_prob = 0.8
        self.stretch_prob = 0.8
        self.max_stretch = 0.2

        # flip augmentation params
        self.do_flip = do_flip
        self.do_rotate = do_rotate
        self.h_flip_prob = 0.5
        self.v_flip_prob = 0.1
        self.rotate_prob = 0.1

        # photometric augmentation params
        self.photo_aug = ColorJitter(
            brightness=0.4, contrast=0.4, saturation=0.4, hue=0.5 / 3.14
        )
        self.asymmetric_color_aug_prob = 0.2
        self.eraser_aug_prob = 0.5
        self.pwc_aug = pwc_aug
        if self.pwc_aug:
            print("[Using pwc-style spatial augmentation]")

    def color_transform(self, imgs):
        """Photometric augmentation"""

        # asymmetric
        if np.random.rand() < self.asymmetric_color_aug_prob:
            for i in range(len(imgs)):
                imgs[i] = np.array(
                    self.photo_aug(Image.fromarray(imgs[i])), dtype=np.uint8
                )
        # symmetric
        else:
            cnt_imgs = len(imgs)
            image_stack = np.concatenate(imgs, axis=0)
            image_stack = np.array(
                self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8
            )
            imgs = list(np.split(image_stack, cnt_imgs, axis=0))

        return imgs

    def eraser_transform(self, imgs, bounds=[50, 100]):
        """Occlusion augmentation (only for last frame)"""

        img_cnt = len(imgs)
        ht, wd, _ = imgs[0].shape

        if img_cnt >= 3:
            img_inds = [0, img_cnt - 1]
        else:
            img_inds = [img_cnt - 1]

        if np.random.rand() < self.eraser_aug_prob:
            for img_ind in img_inds:
                mean_color = np.mean(imgs[img_ind].reshape(-1, 3), axis=0)
                for _ in range(np.random.randint(1, 3)):
                    x0 = np.random.randint(0, wd)
                    y0 = np.random.randint(0, ht)
                    dx = np.random.randint(bounds[0], bounds[1])
                    dy = np.random.randint(bounds[0], bounds[1])
                    imgs[img_ind][y0 : y0 + dy, x0 : x0 + dx, :] = mean_color

        return imgs

    def spatial_transform(self, imgs, flows, valids):
        if self.pre_scale != 0:
            cur_scale = 2**self.pre_scale

            for i in range(len(imgs)):
                imgs[i] = cv2.resize(
                    imgs[i],
                    None,
                    fx=cur_scale,
                    fy=cur_scale,
                    interpolation=cv2.INTER_LINEAR,
                )

            for i in range(len(flows)):
                flows[i] = cv2.resize(
                    flows[i],
                    None,
                    fx=cur_scale,
                    fy=cur_scale,
                    interpolation=cv2.INTER_LINEAR,
                )
                flows[i] = flows[i] * [cur_scale, cur_scale]

            for i in range(len(valids)):
                valids[i] = cv2.resize(
                    valids[i],
                    None,
                    fx=cur_scale,
                    fy=cur_scale,
                    interpolation=cv2.INTER_LINEAR,
                )
                valids[i] = valids[i].round()

        if self.do_rotate:
            if np.random.rand() < self.rotate_prob:
                if np.random.rand() < 0.5:
                    num_rot = 1
                else:
                    num_rot = -1

                for i in range(len(imgs)):
                    imgs[i] = np.rot90(imgs[i], k=num_rot)

                for i in range(len(flows)):
                    if num_rot == 1:
                        flows[i] = np.rot90(flows[i], k=1)[:, :, ::-1] * [-1.0, 1.0]
                    elif num_rot == -1:
                        flows[i] = np.rot90(flows[i], k=-1)[:, :, ::-1] * [1.0, -1.0]

                for i in range(len(valids)):
                    valids[i] = np.rot90(valids[i], k=num_rot)

            pad_t = 0
            pad_b = 0
            pad_l = 0
            pad_r = 0
            if self.crop_size[0] > imgs[0].shape[0]:
                pad_b = self.crop_size[0] - imgs[0].shape[0]
            if self.crop_size[1] > imgs[0].shape[1]:
                pad_r = self.crop_size[1] - imgs[0].shape[1]
            if pad_b != 0 or pad_r != 0:
                for i in range(len(imgs)):
                    imgs[i] = np.pad(
                        imgs[i],
                        ((pad_t, pad_b), (pad_l, pad_r), (0, 0)),
                        "constant",
                        constant_values=((0, 0), (0, 0), (0, 0)),
                    )

                for i in range(len(flows)):
                    flows[i] = np.pad(
                        flows[i],
                        ((pad_t, pad_b), (pad_l, pad_r), (0, 0)),
                        "constant",
                        constant_values=((0, 0), (0, 0), (0, 0)),
                    )

                for i in range(len(valids)):
                    valids[i] = np.pad(
                        valids[i],
                        ((pad_t, pad_b), (pad_l, pad_r)),
                        "constant",
                        constant_values=((0, 0), (0, 0)),
                    )

        # randomly sample scale
        ht, wd = imgs[0].shape[:2]
        min_scale = np.maximum(
            (self.crop_size[0] + 8) / float(ht), (self.crop_size[1] + 8) / float(wd)
        )

        scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
        scale_x = scale
        scale_y = scale
        if np.random.rand() < self.stretch_prob:
            scale_x *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
            scale_y *= 2 ** np.random.uniform(-self.max_stretch, self.max_stretch)
        scale_x = np.clip(scale_x, min_scale, None)
        scale_y = np.clip(scale_y, min_scale, None)

        if np.random.rand() < self.spatial_aug_prob:
            # rescale the images
            for i in range(len(imgs)):
                imgs[i] = cv2.resize(
                    imgs[i],
                    None,
                    fx=scale_x,
                    fy=scale_y,
                    interpolation=cv2.INTER_LINEAR,
                )

            for i in range(len(flows)):
                flows[i] = cv2.resize(
                    flows[i],
                    None,
                    fx=scale_x,
                    fy=scale_y,
                    interpolation=cv2.INTER_LINEAR,
                )
                flows[i] = flows[i] * [scale_x, scale_y]

            for i in range(len(valids)):
                valids[i] = cv2.resize(
                    valids[i],
                    None,
                    fx=scale_x,
                    fy=scale_y,
                    interpolation=cv2.INTER_LINEAR,
                )
                valids[i] = valids[i].round()

        if self.do_flip:
            if np.random.rand() < self.h_flip_prob:  # h-flip
                for i in range(len(imgs)):
                    imgs[i] = imgs[i][:, ::-1]

                for i in range(len(flows)):
                    flows[i] = flows[i][:, ::-1] * [-1.0, 1.0]

                for i in range(len(valids)):
                    valids[i] = valids[i][:, ::-1]

            if np.random.rand() < self.v_flip_prob:  # v-flip
                for i in range(len(imgs)):
                    imgs[i] = imgs[i][::-1, :]

                for i in range(len(flows)):
                    flows[i] = flows[i][::-1, :] * [1.0, -1.0]

                for i in range(len(valids)):
                    valids[i] = valids[i][::-1, :]

        if imgs[0].shape[0] == self.crop_size[0]:
            y0 = 0
        else:
            y0 = np.random.randint(0, imgs[0].shape[0] - self.crop_size[0])
        if imgs[0].shape[1] == self.crop_size[1]:
            x0 = 0
        else:
            x0 = np.random.randint(0, imgs[0].shape[1] - self.crop_size[1])

        for i in range(len(imgs)):
            imgs[i] = imgs[i][y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]

        for i in range(len(flows)):
            flows[i] = flows[i][
                y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]
            ]

        for i in range(len(valids)):
            valids[i] = valids[i][
                y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]
            ]

        return imgs, flows, valids

    def __call__(self, imgs, flows, valids):
        imgs = self.color_transform(imgs)
        imgs = self.eraser_transform(imgs)

        if self.pwc_aug:
            raise NotImplementedError
        else:
            imgs, flows, valids = self.spatial_transform(imgs, flows, valids)

        imgs = [np.ascontiguousarray(imgs[i]) for i in range(len(imgs))]
        flows = [np.ascontiguousarray(flows[i]) for i in range(len(flows))]
        valids = [np.ascontiguousarray(valids[i]) for i in range(len(valids))]

        return imgs, flows, valids


class SparseFlowAugmentor:
    def __init__(
        self,
        crop_size,
        min_scale=-0.2,
        max_scale=0.5,
        do_flip=False,
        do_rotate=False,
        pre_scale=0.0,
    ):
        # spatial augmentation params
        self.crop_size = crop_size
        self.pre_scale = pre_scale
        self.min_scale = min_scale
        self.max_scale = max_scale
        self.spatial_aug_prob = 0.8
        self.stretch_prob = 0.8
        self.max_stretch = 0.2

        # flip augmentation params
        self.do_flip = do_flip
        self.do_rotate = do_rotate
        self.h_flip_prob = 0.5
        self.v_flip_prob = 0.1
        self.rotate_prob = 0.1

        # photometric augmentation params
        self.photo_aug = ColorJitter(
            brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3 / 3.14
        )
        self.asymmetric_color_aug_prob = 0.2
        self.eraser_aug_prob = 0.5

    def color_transform(self, imgs):
        cnt_imgs = len(imgs)
        image_stack = np.concatenate(imgs, axis=0)
        image_stack = np.array(
            self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8
        )
        imgs = list(np.split(image_stack, cnt_imgs, axis=0))
        return imgs

    def eraser_transform(self, imgs):
        ht, wd = imgs[0].shape[:2]
        if np.random.rand() < self.eraser_aug_prob:
            mean_color = np.mean(imgs[-1].reshape(-1, 3), axis=0)
            for _ in range(np.random.randint(1, 3)):
                x0 = np.random.randint(0, wd)
                y0 = np.random.randint(0, ht)
                dx = np.random.randint(50, 100)
                dy = np.random.randint(50, 100)
                imgs[-1][y0 : y0 + dy, x0 : x0 + dx, :] = mean_color

        return imgs

    def resize_sparse_flow_map(self, flow, valid, fx=1.0, fy=1.0):
        ht, wd = flow.shape[:2]
        coords = np.meshgrid(np.arange(wd), np.arange(ht), indexing="ij")
        coords = np.stack(coords, axis=-1)

        coords = coords.reshape(-1, 2).astype(np.float32)
        flow = flow.reshape(-1, 2).astype(np.float32)
        valid = valid.reshape(-1).astype(np.float32)

        coords0 = coords[valid >= 1]
        flow0 = flow[valid >= 1]

        ht1 = int(round(ht * fy))
        wd1 = int(round(wd * fx))

        coords1 = coords0 * [fx, fy]
        flow1 = flow0 * [fx, fy]

        xx = np.round(coords1[:, 0]).astype(np.int32)
        yy = np.round(coords1[:, 1]).astype(np.int32)

        v = (xx > 0) & (xx < wd1) & (yy > 0) & (yy < ht1)
        xx = xx[v]
        yy = yy[v]
        flow1 = flow1[v]

        flow_img = np.zeros([ht1, wd1, 2], dtype=np.float32)
        valid_img = np.zeros([ht1, wd1], dtype=np.int32)

        flow_img[yy, xx] = flow1
        valid_img[yy, xx] = 1

        return flow_img, valid_img

    def spatial_transform(self, imgs, flows, valids):
        if self.pre_scale != 0:
            cur_scale = 2**self.pre_scale

            for i in range(len(imgs)):
                imgs[i] = cv2.resize(
                    imgs[i],
                    None,
                    fx=cur_scale,
                    fy=cur_scale,
                    interpolation=cv2.INTER_LINEAR,
                )

            for i in range(len(flows)):
                flows[i] = cv2.resize(
                    flows[i],
                    None,
                    fx=cur_scale,
                    fy=cur_scale,
                    interpolation=cv2.INTER_NEAREST,
                )
                flows[i] = flows[i] * [cur_scale, cur_scale]

            for i in range(len(valids)):
                valids[i] = cv2.resize(
                    valids[i],
                    None,
                    fx=cur_scale,
                    fy=cur_scale,
                    interpolation=cv2.INTER_NEAREST,
                )
                valids[i] = valids[i].round()

        if self.do_rotate:
            if np.random.rand() < self.rotate_prob:
                if np.random.rand() < 0.5:
                    num_rot = 1
                else:
                    num_rot = -1

                for i in range(len(imgs)):
                    imgs[i] = np.rot90(imgs[i], k=num_rot)

                for i in range(len(flows)):
                    if num_rot == 1:
                        flows[i] = np.rot90(flows[i], k=1)[:, :, ::-1] * [-1.0, 1.0]
                    elif num_rot == -1:
                        flows[i] = np.rot90(flows[i], k=-1)[:, :, ::-1] * [1.0, -1.0]

                for i in range(len(valids)):
                    valids[i] = np.rot90(valids[i], k=num_rot)

        pad_t = 0
        pad_b = 0
        pad_l = 0
        pad_r = 0
        if self.crop_size[0] > imgs[0].shape[0]:
            pad_b = self.crop_size[0] - imgs[0].shape[0]
        if self.crop_size[1] > imgs[0].shape[1]:
            pad_r = self.crop_size[1] - imgs[0].shape[1]
        if pad_b != 0 or pad_r != 0:
            for i in range(len(imgs)):
                imgs[i] = np.pad(
                    imgs[i],
                    ((pad_t, pad_b), (pad_l, pad_r), (0, 0)),
                    "constant",
                    constant_values=((0, 0), (0, 0), (0, 0)),
                )

            for i in range(len(flows)):
                flows[i] = np.pad(
                    flows[i],
                    ((pad_t, pad_b), (pad_l, pad_r), (0, 0)),
                    "constant",
                    constant_values=((0, 0), (0, 0), (0, 0)),
                )

            for i in range(len(valids)):
                valids[i] = np.pad(
                    valids[i],
                    ((pad_t, pad_b), (pad_l, pad_r)),
                    "constant",
                    constant_values=((0, 0), (0, 0)),
                )
        # randomly sample scale

        ht, wd = imgs[0].shape[:2]
        min_scale = np.maximum(
            (self.crop_size[0] + 1) / float(ht), (self.crop_size[1] + 1) / float(wd)
        )

        scale = 2 ** np.random.uniform(self.min_scale, self.max_scale)
        scale_x = np.clip(scale, min_scale, None)
        scale_y = np.clip(scale, min_scale, None)

        if np.random.rand() < self.spatial_aug_prob:
            # rescale the images

            for i in range(len(imgs)):
                imgs[i] = cv2.resize(
                    imgs[i],
                    None,
                    fx=scale_x,
                    fy=scale_y,
                    interpolation=cv2.INTER_LINEAR,
                )

            for i in range(len(flows)):
                flow, valid = self.resize_sparse_flow_map(
                    flows[i], valids[i], fx=scale_x, fy=scale_y
                )
                flows[i] = flow
                valids[i] = valid

        if self.do_flip:
            if np.random.rand() < 0.5:  # h-flip
                for i in range(len(imgs)):
                    imgs[i] = imgs[i][:, ::-1]

                for i in range(len(flows)):
                    flows[i] = flows[i][:, ::-1] * [-1.0, 1.0]

                for i in range(len(valids)):
                    valids[i] = valids[i][:, ::-1]

        margin_y = 20
        margin_x = 50

        y0 = np.random.randint(0, imgs[0].shape[0] - self.crop_size[0] + margin_y)
        x0 = np.random.randint(
            -margin_x, imgs[0].shape[1] - self.crop_size[1] + margin_x
        )

        y0 = np.clip(y0, 0, imgs[0].shape[0] - self.crop_size[0])
        x0 = np.clip(x0, 0, imgs[0].shape[1] - self.crop_size[1])

        for i in range(len(imgs)):
            imgs[i] = imgs[i][y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]

        for i in range(len(flows)):
            flows[i] = flows[i][
                y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]
            ]

        for i in range(len(valids)):
            valids[i] = valids[i][
                y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]
            ]

        return imgs, flows, valids

    def __call__(self, imgs, flows, valids):
        imgs = self.color_transform(imgs)
        imgs = self.eraser_transform(imgs)

        imgs, flows, valids = self.spatial_transform(imgs, flows, valids)

        imgs = [np.ascontiguousarray(imgs[i]) for i in range(len(imgs))]
        flows = [np.ascontiguousarray(flows[i]) for i in range(len(flows))]
        valids = [np.ascontiguousarray(valids[i]) for i in range(len(valids))]

        return imgs, flows, valids