"""
augmentations.py

Simple camera-style augmentations for color fundus photography (CFP)
classification.

Expected input:
    RGB NumPy image, shape (H, W, 3)

Dependencies:
    pip install albumentations opencv-python
"""

import cv2
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 get_train_transforms(
    image_size=1024,
    mean=IMAGENET_MEAN,
    std=IMAGENET_STD,
):
    return A.Compose([
        A.Resize(image_size, image_size),

        # Geometry is safe
        A.HorizontalFlip(p=0.5),

        A.ShiftScaleRotate(
            shift_limit=0.02,
            scale_limit=0.03,   # slightly reduced
            rotate_limit=5,     # slightly reduced
            border_mode=0,
            value=0,
            p=0.3,
        ),

        # MUCH weaker photometric changes
        A.RandomBrightnessContrast(
            brightness_limit=0.08,   # ↓ from 0.15
            contrast_limit=0.08,
            p=0.3,
        ),

        # Remove or reduce gamma
        A.RandomGamma(
            gamma_limit=(95, 105),   # very mild
            p=0.2,
        ),

        # Remove hue shift entirely (important)
        # Hue shifts are not realistic for fundus physiology
        # -> comment this out or reduce heavily
        # A.HueSaturationValue(...)

        # Keep mild quality perturbation
        A.OneOf([
            A.GaussianBlur(blur_limit=(3, 5)),
            A.Downscale(scale_min=0.85, scale_max=0.95, interpolation=cv2.INTER_LINEAR),
            A.ImageCompression(quality_lower=80, quality_upper=100),
        ], p=0.15),

        A.Normalize(mean=mean, std=std),
        ToTensorV2(),
    ])


def get_val_transforms(
    image_size=1024,
    mean=IMAGENET_MEAN,
    std=IMAGENET_STD,
):
    """
    Validation/test transforms.
    """
    return A.Compose([
        A.Resize(image_size, image_size),
        A.Normalize(mean=mean, std=std),
        ToTensorV2(),
    ])