from __future__ import annotations

import random
from io import BytesIO

from torch import Tensor
from torch.utils.data import Dataset
from torchvision.io import read_image
from torchvision.transforms import ColorJitter
from torchvision.transforms import InterpolationMode
from torchvision.transforms import Normalize
from torchvision.transforms.functional import hflip
from torchvision.transforms.functional import pad
from torchvision.transforms.functional import resize
from torchvision.transforms.functional import rotate

from core.database import connect

# 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):
    def __init__(
        self,
        data_name_list: list[str],
        do_augment: bool = False,
        random_annotations: bool = False,
        normalize_mean: list[float] = RESNET_NORMALIZE_MEAN,
        normalize_std: list[float] = RESNET_NORMALIZE_STD,
    ):
        super().__init__()
        self.data_name_list = data_name_list
        self.do_augment = do_augment
        self.random_annotations = random_annotations
        self.normalize_mean = normalize_mean
        self.normalize_std = normalize_std
        # prepare augmentation functions
        self.normalize = Normalize(
            mean=normalize_mean,
            std=normalize_std,
        )
        self.color_jitter = ColorJitter(
            brightness=1e-1,
            contrast=8e-2,
            saturation=8e-2,
        )
        # connect to database
        collection, _, _ = connect()
        self.collection = collection

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

    def __getitem__(self, idx):
        # get image from database
        name = self.data_name_list[idx]
        query = {
            'name': name,
        }
        projection = {
            '_id': False,
            'image': True,
        }
        img_bytes = self.collection.find_one(
            filter=query,
            projection=projection,
        )
        img = self.load_image(img_bytes)
        if self.do_augment:
            img = self.augment(img)
        return img

    def load_image(
        self,
        data: bytes,
    ) -> Tensor:
        """Load images tensor from bytes."""
        assert isinstance(data, bytes)
        img = read_image(BytesIO(data))
        img /= 255  # normalize 8-bit image
        img = self.square_pad(img)
        img = resize(
            img=img,
            size=(512, 512),
            interpolation=InterpolationMode.BICUBIC,
        )
        return img

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

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

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