"""Image Augmentation and Transformation Pipelines using Albumentations."""

from typing import Tuple
import albumentations as A
from albumentations.pytorch import ToTensorV2

IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


def _make_random_resized_crop(image_size: int, scale: Tuple[float, float] = (0.8, 1.0)):
    """Version-agnostic RandomResizedCrop creator."""
    try:
        return A.RandomResizedCrop(size=(image_size, image_size), scale=scale)
    except Exception:
        return A.RandomResizedCrop(height=image_size, width=image_size, scale=scale)


def _make_resize(image_size: int):
    """Version-agnostic Resize creator."""
    try:
        return A.Resize(size=(image_size, image_size))
    except Exception:
        return A.Resize(height=image_size, width=image_size)


def _make_coarse_dropout(p: float = 0.3):
    """Version-agnostic CoarseDropout creator."""
    try:
        return A.CoarseDropout(
            num_holes_range=(1, 8),
            hole_height_range=(8, 32),
            hole_width_range=(8, 32),
            p=p,
        )
    except Exception:
        try:
            return A.CoarseDropout(max_holes=8, max_height=32, max_width=32, min_holes=1, p=p)
        except Exception:
            return A.Dropout(p=p)


def get_train_transforms(image_size: int = 380) -> A.Compose:
    return A.Compose(
        [
            _make_random_resized_crop(image_size=image_size, scale=(0.8, 1.0)),
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.RandomRotate90(p=0.5),
            A.ColorJitter(
                brightness=0.2,
                contrast=0.2,
                saturation=0.1,
                hue=0.05,
                p=0.5,
            ),
            _make_coarse_dropout(p=0.3),
            A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
            ToTensorV2(),
        ]
    )


def get_val_transforms(image_size: int = 380) -> A.Compose:
    return A.Compose(
        [
            _make_resize(image_size=image_size),
            A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
            ToTensorV2(),
        ]
    )