from __future__ import annotations
import random

import albumentations as A
import numpy as np
import PIL.Image
from albumentations.pytorch.transforms import ToTensorV2
from torch.utils.data import Dataset, DataLoader

from realfake.config import LABELS

IMG_RESIZE = 256
IMG_CROP = 224


class DictDataset(Dataset):
    def __init__(self, records: list[dict], transform_x=None):
        self.records = records
        self.transform_x = transform_x

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

    def __getitem__(self, idx):
        record = self.records[idx]
        image = np.asarray(PIL.Image.open(record["path"]))
        if self.transform_x is not None:
            image = self.transform_x(image=image)["image"]
        item = {"image": image}
        if "label" in record:
            item["label"] = LABELS[record["label"]]
        return item


def get_augs(train: bool = True) -> A.Compose:
    if train:
        return A.Compose([
            A.Resize(IMG_RESIZE, IMG_RESIZE),
            A.RandomCrop(IMG_CROP, IMG_CROP),
            A.HorizontalFlip(),
            A.VerticalFlip(),
            A.RandomBrightnessContrast(),
            A.Affine(),
            A.Rotate(),
            A.CoarseDropout(),
            ExpandChannels(),
            RGBAtoRGB(),
            A.Normalize(),
            ToTensorV2(),
        ])
    else:
        return A.Compose([
            A.Resize(IMG_RESIZE, IMG_RESIZE),
            A.CenterCrop(IMG_CROP, IMG_CROP),
            ExpandChannels(),
            RGBAtoRGB(),
            A.Normalize(),
            ToTensorV2(),
        ])


class ExpandChannels(A.ImageOnlyTransform):
    """Expands image up to three channes if the image is grayscale."""

    def __init__(self, always_apply: bool = False, p: float = 0.5):
        super().__init__(True, 1.0)

    def apply(self, image, **params):
        if image.ndim == 2:
            image = np.repeat(image[..., None], 3, axis=2)
        elif image.shape[2] == 1:
            image = np.repeat(image, 3, axis=2)
        return image


class RGBAtoRGB(A.ImageOnlyTransform):
    """Converts RGBA image to RGB."""

    def __init__(self, always_apply: bool = False, p: float = 0.5):
        super().__init__(True, 1.0)

    def apply(self, image, **params):
        if image.shape[2] == 4:
            image = image[:, :, :3]
        return image


def get_dss(records: list) -> tuple[DictDataset, DictDataset]:
    train_records = [x for x in records if not x["valid"]]
    valid_records = [x for x in records if x["valid"]]
    assert len(train_records) + len(valid_records) == len(records)
    random.shuffle(train_records)
    train_ds = DictDataset(train_records, transform_x=get_augs(train=True))
    valid_ds = DictDataset(valid_records, transform_x=get_augs(train=False))
    return train_ds, valid_ds


def get_dls(train_ds: DictDataset, valid_ds: DictDataset, bs: int, num_workers: int) -> tuple[DataLoader, DataLoader]:
    train_dl = DataLoader(train_ds, batch_size=bs, num_workers=num_workers)
    valid_dl = DataLoader(valid_ds, batch_size=bs, num_workers=num_workers, shuffle=False)
    return train_dl, valid_dl