import random
from typing import Union

import cv2
import numpy as np
from albumentations import DualTransform, to_tuple


def _parse_border_mode(border: Union[str, int]):
    if isinstance(border, str):
        mode = {
            'reflect': cv2.BORDER_REFLECT101,
        }[border]
        value = None
    else:
        mode = cv2.BORDER_CONSTANT
        value = border
    return dict(borderMode=mode, borderValue=value)


def _parse_interpolation(interpolation: str):
    return {
        'lanczos': cv2.INTER_LANCZOS4,
    }[interpolation]


class RandomPerspective(DualTransform):
    def __init__(
        self,
        corner_shift_limit=.2,
        clip_coords=True,
        border: Union[str, int] = 0,
        interpolation: str = 'lanczos',
        always_apply=False, p=0.5
    ):
        super().__init__(always_apply, p)
        self.clip_coords = clip_coords
        self.corner_shift_limit = to_tuple(corner_shift_limit)
        self.border_mode = _parse_border_mode(border)
        self.interpolation = _parse_interpolation(interpolation)

    def get_params(self):
        """
        Sample 4 corner relative shifts (x, y).
        Returns: np.array of shape (4, 2)
        """
        shifts = np.array([
            random.uniform(
                *self.corner_shift_limit,
            ) for _ in range(8)
        ]).reshape(4, 2).astype(np.float32)

        coords: np.ndarray = shifts + [
            [0, 0],
            [1, 0],
            [1, 1],
            [0, 1]
        ]

        if self.clip_coords:
            coords = coords.clip(0, 1)

        return dict(coords=coords)

    def update_params(self, params, **kwargs):
        params = super().update_params(params, **kwargs)
        h, w = params["rows"], params["cols"]
        target_coords = params.pop("coords")

        target_coords = (target_coords *  [[w, h]]).astype(np.float32)

        source_coords = np.array([
            [0, 0],
            [w, 0],
            [w, h],
            [0, h]
        ], dtype=np.float32)

        params["matrix"] = \
            cv2.getPerspectiveTransform(source_coords, target_coords)

        return params

    # noinspection PyMethodOverriding
    def apply(self, image, matrix, interpolation, rows, cols, **kwargs):
        return cv2.warpPerspective(
            image, matrix, (cols, rows),
            flags=interpolation, **self.border_mode
        )