# Copyright (C) 2022-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).

# --------------------------------------------------------
# Data augmentation for training stereo and flow
# --------------------------------------------------------

# References
# https://github.com/autonomousvision/unimatch/blob/master/dataloader/stereo/transforms.py
# https://github.com/autonomousvision/unimatch/blob/master/dataloader/flow/transforms.py


import random

import cv2
import numpy as np
from PIL import Image

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

import torch
import torchvision.transforms.functional as FF
from torchvision.transforms import ColorJitter


class StereoAugmentor(object):
    def __init__(
        self,
        crop_size,
        scale_prob=0.5,
        scale_xonly=True,
        lhth=800.0,
        lminscale=0.0,
        lmaxscale=1.0,
        hminscale=-0.2,
        hmaxscale=0.4,
        scale_interp_nearest=True,
        rightjitterprob=0.5,
        v_flip_prob=0.5,
        color_aug_asym=True,
        color_choice_prob=0.5,
    ):
        self.crop_size = crop_size
        self.scale_prob = scale_prob
        self.scale_xonly = scale_xonly
        self.lhth = lhth
        self.lminscale = lminscale
        self.lmaxscale = lmaxscale
        self.hminscale = hminscale
        self.hmaxscale = hmaxscale
        self.scale_interp_nearest = scale_interp_nearest
        self.rightjitterprob = rightjitterprob
        self.v_flip_prob = v_flip_prob
        self.color_aug_asym = color_aug_asym
        self.color_choice_prob = color_choice_prob

    def _random_scale(self, img1, img2, disp):
        ch, cw = self.crop_size
        h, w = img1.shape[:2]
        if self.scale_prob > 0.0 and np.random.rand() < self.scale_prob:
            min_scale, max_scale = (
                (self.lminscale, self.lmaxscale)
                if min(h, w) < self.lhth
                else (self.hminscale, self.hmaxscale)
            )
            scale_x = 2.0 ** np.random.uniform(min_scale, max_scale)
            scale_x = np.clip(scale_x, (cw + 8) / float(w), None)
            scale_y = 1.0
            if not self.scale_xonly:
                scale_y = scale_x
                scale_y = np.clip(scale_y, (ch + 8) / float(h), None)
            img1 = cv2.resize(
                img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR
            )
            img2 = cv2.resize(
                img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR
            )
            disp = (
                cv2.resize(
                    disp,
                    None,
                    fx=scale_x,
                    fy=scale_y,
                    interpolation=cv2.INTER_LINEAR
                    if not self.scale_interp_nearest
                    else cv2.INTER_NEAREST,
                )
                * scale_x
            )
        else:  # check if we need to resize to be able to crop
            h, w = img1.shape[:2]
            clip_scale = (cw + 8) / float(w)
            if clip_scale > 1.0:
                scale_x = clip_scale
                scale_y = scale_x if not self.scale_xonly else 1.0
                img1 = cv2.resize(
                    img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR
                )
                img2 = cv2.resize(
                    img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR
                )
                disp = (
                    cv2.resize(
                        disp,
                        None,
                        fx=scale_x,
                        fy=scale_y,
                        interpolation=cv2.INTER_LINEAR
                        if not self.scale_interp_nearest
                        else cv2.INTER_NEAREST,
                    )
                    * scale_x
                )
        return img1, img2, disp

    def _random_crop(self, img1, img2, disp):
        h, w = img1.shape[:2]
        ch, cw = self.crop_size
        assert ch <= h and cw <= w, (img1.shape, h, w, ch, cw)
        offset_x = np.random.randint(w - cw + 1)
        offset_y = np.random.randint(h - ch + 1)
        img1 = img1[offset_y : offset_y + ch, offset_x : offset_x + cw]
        img2 = img2[offset_y : offset_y + ch, offset_x : offset_x + cw]
        disp = disp[offset_y : offset_y + ch, offset_x : offset_x + cw]
        return img1, img2, disp

    def _random_vflip(self, img1, img2, disp):
        # vertical flip
        if self.v_flip_prob > 0 and np.random.rand() < self.v_flip_prob:
            img1 = np.copy(np.flipud(img1))
            img2 = np.copy(np.flipud(img2))
            disp = np.copy(np.flipud(disp))
        return img1, img2, disp

    def _random_rotate_shift_right(self, img2):
        if self.rightjitterprob > 0.0 and np.random.rand() < self.rightjitterprob:
            angle, pixel = 0.1, 2
            px = np.random.uniform(-pixel, pixel)
            ag = np.random.uniform(-angle, angle)
            image_center = (
                np.random.uniform(0, img2.shape[0]),
                np.random.uniform(0, img2.shape[1]),
            )
            rot_mat = cv2.getRotationMatrix2D(image_center, ag, 1.0)
            img2 = cv2.warpAffine(
                img2, rot_mat, img2.shape[1::-1], flags=cv2.INTER_LINEAR
            )
            trans_mat = np.float32([[1, 0, 0], [0, 1, px]])
            img2 = cv2.warpAffine(
                img2, trans_mat, img2.shape[1::-1], flags=cv2.INTER_LINEAR
            )
        return img2

    def _random_color_contrast(self, img1, img2):
        if np.random.random() < 0.5:
            contrast_factor = np.random.uniform(0.8, 1.2)
            img1 = FF.adjust_contrast(img1, contrast_factor)
            if self.color_aug_asym and np.random.random() < 0.5:
                contrast_factor = np.random.uniform(0.8, 1.2)
            img2 = FF.adjust_contrast(img2, contrast_factor)
        return img1, img2

    def _random_color_gamma(self, img1, img2):
        if np.random.random() < 0.5:
            gamma = np.random.uniform(0.7, 1.5)
            img1 = FF.adjust_gamma(img1, gamma)
            if self.color_aug_asym and np.random.random() < 0.5:
                gamma = np.random.uniform(0.7, 1.5)
            img2 = FF.adjust_gamma(img2, gamma)
        return img1, img2

    def _random_color_brightness(self, img1, img2):
        if np.random.random() < 0.5:
            brightness = np.random.uniform(0.5, 2.0)
            img1 = FF.adjust_brightness(img1, brightness)
            if self.color_aug_asym and np.random.random() < 0.5:
                brightness = np.random.uniform(0.5, 2.0)
            img2 = FF.adjust_brightness(img2, brightness)
        return img1, img2

    def _random_color_hue(self, img1, img2):
        if np.random.random() < 0.5:
            hue = np.random.uniform(-0.1, 0.1)
            img1 = FF.adjust_hue(img1, hue)
            if self.color_aug_asym and np.random.random() < 0.5:
                hue = np.random.uniform(-0.1, 0.1)
            img2 = FF.adjust_hue(img2, hue)
        return img1, img2

    def _random_color_saturation(self, img1, img2):
        if np.random.random() < 0.5:
            saturation = np.random.uniform(0.8, 1.2)
            img1 = FF.adjust_saturation(img1, saturation)
            if self.color_aug_asym and np.random.random() < 0.5:
                saturation = np.random.uniform(-0.8, 1.2)
            img2 = FF.adjust_saturation(img2, saturation)
        return img1, img2

    def _random_color(self, img1, img2):
        trfs = [
            self._random_color_contrast,
            self._random_color_gamma,
            self._random_color_brightness,
            self._random_color_hue,
            self._random_color_saturation,
        ]
        img1 = Image.fromarray(img1.astype("uint8"))
        img2 = Image.fromarray(img2.astype("uint8"))
        if np.random.random() < self.color_choice_prob:
            # A single transform
            t = random.choice(trfs)
            img1, img2 = t(img1, img2)
        else:
            # Combination of trfs
            # Random order
            random.shuffle(trfs)
            for t in trfs:
                img1, img2 = t(img1, img2)
        img1 = np.array(img1).astype(np.float32)
        img2 = np.array(img2).astype(np.float32)
        return img1, img2

    def __call__(self, img1, img2, disp, dataset_name):
        img1, img2, disp = self._random_scale(img1, img2, disp)
        img1, img2, disp = self._random_crop(img1, img2, disp)
        img1, img2, disp = self._random_vflip(img1, img2, disp)
        img2 = self._random_rotate_shift_right(img2)
        img1, img2 = self._random_color(img1, img2)
        return img1, img2, disp


class FlowAugmentor:
    def __init__(
        self,
        crop_size,
        min_scale=-0.2,
        max_scale=0.5,
        spatial_aug_prob=0.8,
        stretch_prob=0.8,
        max_stretch=0.2,
        h_flip_prob=0.5,
        v_flip_prob=0.1,
        asymmetric_color_aug_prob=0.2,
    ):
        # spatial augmentation params
        self.crop_size = crop_size
        self.min_scale = min_scale
        self.max_scale = max_scale
        self.spatial_aug_prob = spatial_aug_prob
        self.stretch_prob = stretch_prob
        self.max_stretch = max_stretch

        # flip augmentation params
        self.h_flip_prob = h_flip_prob
        self.v_flip_prob = v_flip_prob

        # 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 = asymmetric_color_aug_prob

    def color_transform(self, img1, img2):
        """Photometric augmentation"""

        # asymmetric
        if np.random.rand() < self.asymmetric_color_aug_prob:
            img1 = np.array(self.photo_aug(Image.fromarray(img1)), dtype=np.uint8)
            img2 = np.array(self.photo_aug(Image.fromarray(img2)), dtype=np.uint8)

        # symmetric
        else:
            image_stack = np.concatenate([img1, img2], axis=0)
            image_stack = np.array(
                self.photo_aug(Image.fromarray(image_stack)), dtype=np.uint8
            )
            img1, img2 = np.split(image_stack, 2, axis=0)

        return img1, img2

    def _resize_flow(self, flow, scale_x, scale_y, factor=1.0):
        if np.all(np.isfinite(flow)):
            flow = cv2.resize(
                flow,
                None,
                fx=scale_x / factor,
                fy=scale_y / factor,
                interpolation=cv2.INTER_LINEAR,
            )
            flow = flow * [scale_x, scale_y]
        else:  # sparse version
            fx, fy = scale_x, scale_y
            ht, wd = flow.shape[:2]
            coords = np.meshgrid(np.arange(wd), np.arange(ht))
            coords = np.stack(coords, axis=-1)

            coords = coords.reshape(-1, 2).astype(np.float32)
            flow = flow.reshape(-1, 2).astype(np.float32)
            valid = np.isfinite(flow[:, 0])

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

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

            rescale = np.expand_dims(np.array([fx, fy]), axis=0)
            coords1 = coords0 * rescale / factor
            flow1 = flow0 * rescale

            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 = np.inf * np.ones(
                [ht1, wd1, 2], dtype=np.float32
            )  # invalid value every where, before we fill it with the correct ones
            flow[yy, xx] = flow1
        return flow

    def spatial_transform(self, img1, img2, flow, dname):
        if np.random.rand() < self.spatial_aug_prob:
            # randomly sample scale
            ht, wd = img1.shape[:2]
            clip_min_scale = np.maximum(
                (self.crop_size[0] + 8) / float(ht), (self.crop_size[1] + 8) / float(wd)
            )
            min_scale, max_scale = self.min_scale, self.max_scale
            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, clip_min_scale, None)
            scale_y = np.clip(scale_y, clip_min_scale, None)
            # rescale the images
            img1 = cv2.resize(
                img1, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR
            )
            img2 = cv2.resize(
                img2, None, fx=scale_x, fy=scale_y, interpolation=cv2.INTER_LINEAR
            )
            flow = self._resize_flow(
                flow, scale_x, scale_y, factor=2.0 if dname == "Spring" else 1.0
            )
        elif dname == "Spring":
            flow = self._resize_flow(flow, 1.0, 1.0, factor=2.0)

        if self.h_flip_prob > 0.0 and np.random.rand() < self.h_flip_prob:  # h-flip
            img1 = img1[:, ::-1]
            img2 = img2[:, ::-1]
            flow = flow[:, ::-1] * [-1.0, 1.0]

        if self.v_flip_prob > 0.0 and np.random.rand() < self.v_flip_prob:  # v-flip
            img1 = img1[::-1, :]
            img2 = img2[::-1, :]
            flow = flow[::-1, :] * [1.0, -1.0]

        # In case no cropping
        if img1.shape[0] - self.crop_size[0] > 0:
            y0 = np.random.randint(0, img1.shape[0] - self.crop_size[0])
        else:
            y0 = 0
        if img1.shape[1] - self.crop_size[1] > 0:
            x0 = np.random.randint(0, img1.shape[1] - self.crop_size[1])
        else:
            x0 = 0

        img1 = img1[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]
        img2 = img2[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]
        flow = flow[y0 : y0 + self.crop_size[0], x0 : x0 + self.crop_size[1]]

        return img1, img2, flow

    def __call__(self, img1, img2, flow, dname):
        img1, img2, flow = self.spatial_transform(img1, img2, flow, dname)
        img1, img2 = self.color_transform(img1, img2)
        img1 = np.ascontiguousarray(img1)
        img2 = np.ascontiguousarray(img2)
        flow = np.ascontiguousarray(flow)
        return img1, img2, flow