"""Definition of Visual Critial Discourse Analysis dataloader."""

import random

from minio import Minio
from PIL import Image
from torch import Tensor
from torch.utils.data import Dataset
from torchvision.transforms import ColorJitter, InterpolationMode, Normalize
from torchvision.transforms.functional import (
    hflip,
    pad,
    resize,
    rotate,
    to_pil_image,
    to_tensor,
)

from shared.data_store import get_image

# resnet18 original normalization values
RESNET_NORMALIZE_MEAN = [0.485, 0.456, 0.406]
RESNET_NORMALIZE_STD = [0.229, 0.224, 0.225]


class VCDADataset(Dataset):
    """VCDA dataset class."""

    def __init__(
        self,
        minio_client: Minio,
        data_name_list: list[str],
        do_augment: bool = False,
        random_annotations: bool = False,
    ):
        super().__init__()
        self.minio_client = minio_client
        self.data_name_list = data_name_list
        self.do_augment = do_augment
        self.random_annotations = random_annotations
        self.normalize_mean = RESNET_NORMALIZE_MEAN
        self.normalize_std = RESNET_NORMALIZE_STD
        # prepare augmentation functions
        self.normalize = Normalize(
            mean=self.normalize_mean,
            std=self.normalize_std,
        )
        self.color_jitter = ColorJitter(
            brightness=1e-1,
            contrast=8e-2,
            saturation=8e-2,
        )

    def __len__(self):
        return len(self.data_name_list)

    def __getitem__(self, idx):
        # get image from database
        object_name = self.data_name_list[idx]
        image = get_image(
            client=self.minio_client,
            object_name=object_name,
        )
        tensor = self.image_to_tensor(image)
        if self.do_augment:
            tensor = self.augment(tensor)
        return tensor

    def image_to_tensor(
        self,
        image: Image.Image,
    ) -> Tensor:
        """Load images tensor from bytes."""
        tensor = to_tensor(image)
        tensor /= 255  # normalize 8-bit image
        tensor = self.square_pad(tensor=tensor)
        tensor = resize(
            img=tensor,
            size=(512, 512),
            interpolation=InterpolationMode.BICUBIC,
        )
        return tensor

    @staticmethod
    def square_pad(
        tensor: Tensor,
    ) -> Tensor:
        """Pads image to a square with side length equal to the largest side of
        the input image."""
        assert isinstance(tensor, Tensor)
        # B, nc, w, h = img.shape
        h = tensor.shape[-2]
        w = tensor.shape[-1]
        if h == w:
            return tensor
        max_wh = max([h, w])
        hp = int((max_wh - w) / 2)
        vp = int((max_wh - h) / 2)
        padding = (hp, vp, hp, vp)
        tensor = pad(tensor, padding, 0, 'constant')
        return tensor

    def augment(
        self,
        tensor: Tensor,
    ) -> Tensor:
        """Augment image with random horizontal flips, rotations and color
        jitter."""
        assert isinstance(tensor, Tensor)
        # left-right flip
        if random.random() >= 0.5:
            tensor = hflip(tensor)
        # rotation
        rnd = random.random()
        if rnd < 0.25:
            tensor = rotate(tensor, angle=90)
        if rnd < 0.5:
            tensor = rotate(tensor, angle=180)
        if rnd < 0.75:
            tensor = rotate(tensor, angle=270)
        # color jitter
        tensor = self.color_jitter(tensor)
        return tensor

    def reverse_normalise(
        self,
        tensor: Tensor,
    ) -> Image.Image:
        """Reverse normalization to get an image that can be interpreted by
        humans."""
        assert isinstance(tensor, Tensor)
        tensor *= Tensor(self.normalize_std).reshape((3, 1, 1))
        tensor += Tensor(self.normalize_mean).reshape((3, 1, 1))
        image = to_pil_image(
            pic=tensor,
            mode='RGB',
        )
        return image