import cv2
import torch
import torchvision

from typing import Literal
from abc import ABCMeta, abstractmethod
import numpy as np
import random

try:
    import albumentations as A
except ImportError:
    # albumentations is not installed, training with augmentations will not be possible
    A = None


def get_augmentations(split: Literal["train", "dev", "test"], args):
    if split == "train":
        augmentations = [
            Scale_2d(args, {}),
            Rotate_Range(args, {"deg": 20}),
            ToTensor(),
            Force_Num_Chan_Tensor_2d(args, {}),
            Normalize_Tensor_2d(args, {}),
        ]
    else:
        augmentations = [
            Scale_2d(args, {}),
            ToTensor(),
            Force_Num_Chan_Tensor_2d(args, {}),
            Normalize_Tensor_2d(args, {}),
        ]

    return augmentations


class Abstract_augmentation(object):
    """
    Abstract-transformer.
    Default - non cachable
    """

    __metaclass__ = ABCMeta

    def __init__(self):
        self._is_cachable = False
        self._trans_sep = "@"
        self._attr_sep = "#"
        self.name = (
            self.__str__().split("sybil.augmentations.")[-1].split(" ")[0].lower()
        )

    @abstractmethod
    def __call__(self, input_dict):
        pass

    def set_seed(self, seed):
        random.seed(seed)
        np.random.seed(seed)
        torch.random.manual_seed(seed)

    def cachable(self):
        return self._is_cachable

    def set_cachable(self, *keys):
        """
        Sets the transformer as cachable
        and sets the _caching_keys according to the input variables.
        """
        self._is_cachable = True
        name_str = "{}{}".format(self._trans_sep, self.name)
        keys_str = "".join(self._attr_sep + str(k) for k in keys)
        self._caching_keys = "{}{}".format(name_str, keys_str)
        return

    def caching_keys(self):
        return self._caching_keys


class ComposeAug(Abstract_augmentation):
    """
    Composes multiple augmentations
    """

    def __init__(self, augmentations):
        super(ComposeAug, self).__init__()
        self.augmentations = augmentations

    def __call__(self, input_dict, sample=None):
        for transformer in self.augmentations:
            input_dict = transformer(input_dict, sample)

        return input_dict


class ToTensor(Abstract_augmentation):
    """
    torchvision.transforms.ToTensor wrapper.
    """

    def __init__(self):
        super(ToTensor, self).__init__()
        self.name = "totensor"

    def __call__(self, input_dict, sample=None):
        input_dict["input"] = torch.from_numpy(input_dict["input"]).float()
        if input_dict.get("mask", None) is not None:
            input_dict["mask"] = torch.from_numpy(input_dict["mask"]).float()
        return input_dict


class ResizeTransform:
    def __init__(self, width, height):
        self.width = width
        self.height = height

    def __call__(self, image=None, mask=None):
        out = {"image": None, "mask": None}
        if image is not None:
            out["image"] = cv2.resize(image, dsize=(self.width, self.height), interpolation=cv2.INTER_LINEAR)
        if mask is not None:
            out["mask"] = cv2.resize(mask, dsize=(self.width, self.height), interpolation=cv2.INTER_NEAREST)
        return out


class Scale_2d(Abstract_augmentation):
    """
    Given PIL image, enforce its some set size
    (can use for down sampling / keep full res)
    """

    def __init__(self, args, kwargs):
        super(Scale_2d, self).__init__()
        assert len(kwargs.keys()) == 0
        width, height = args.img_size
        self.set_cachable(width, height)
        self.transform = ResizeTransform(width, height)

    def __call__(self, input_dict, sample=None):
        out = self.transform(
            image=input_dict["input"], mask=input_dict.get("mask", None)
        )
        input_dict["input"] = out["image"]
        input_dict["mask"] = out["mask"]
        return input_dict


class Rotate_Range(Abstract_augmentation):
    """
    Rotate image counter clockwise by random degree https://albumentations.ai/docs/api_reference/augmentations/geometric/rotate/#albumentations.augmentations.geometric.rotate.Rotate

        kwargs
            deg: max degrees to rotate
    """

    def __init__(self, args, kwargs):
        super(Rotate_Range, self).__init__()
        assert len(kwargs.keys()) == 1
        self.max_angle = int(kwargs["deg"])
        assert A is not None, "albumentations is not installed"
        self.transform = A.Rotate(limit=self.max_angle, p=0.5)

    def __call__(self, input_dict, sample=None):
        if sample and "seed" in sample:
            self.set_seed(sample["seed"])
        out = self.transform(
            image=input_dict["input"], mask=input_dict.get("mask", None)
        )
        input_dict["input"] = out["image"]
        input_dict["mask"] = out["mask"]
        return input_dict


class Normalize_Tensor_2d(Abstract_augmentation):
    """
    Normalizes input by channel
    wrapper for torchvision.transforms.Normalize wrapper.
    """

    def __init__(self, args, kwargs):
        super(Normalize_Tensor_2d, self).__init__()
        assert len(kwargs) == 0
        channel_means = [args.img_mean] if len(args.img_mean) == 1 else args.img_mean
        channel_stds = [args.img_std] if len(args.img_std) == 1 else args.img_std

        self.transform = torchvision.transforms.Normalize(
            torch.Tensor(channel_means), torch.Tensor(channel_stds)
        )

        self.permute = args.img_file_type in [
            "png",
        ]

    def __call__(self, input_dict, sample=None):
        img = input_dict["input"]
        if len(img.size()) == 2:
            img = img.unsqueeze(0)

        if self.permute:
            img = img.permute(2, 0, 1)
            input_dict["input"] = self.transform(img).permute(1, 2, 0)
        else:
            input_dict["input"] = self.transform(img)

        return input_dict


class Force_Num_Chan_Tensor_2d(Abstract_augmentation):
    """
    Convert gray scale images to image with args.num_chan num channels.
    """

    def __init__(self, args, kwargs):
        super(Force_Num_Chan_Tensor_2d, self).__init__()
        assert len(kwargs) == 0
        self.args = args

    def __call__(self, input_dict, sample=None):
        img = input_dict["input"]
        mask = input_dict.get("mask", None)
        if mask is not None:
            input_dict["mask"] = mask.unsqueeze(0)

        num_dims = len(img.shape)
        if num_dims == 2:
            img = img.unsqueeze(0)
        existing_chan = img.size()[0]
        if not existing_chan == self.args.num_chan:
            input_dict["input"] = img.expand(self.args.num_chan, *img.size()[1:])

        return input_dict