"""
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 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,
):
    """
    Training transforms.
    """
    return A.Compose([
        A.Resize(image_size, image_size),

        A.HorizontalFlip(p=0.5),

        A.ShiftScaleRotate(
            shift_limit=0.02,
            scale_limit=0.05,
            rotate_limit=7,
            border_mode=0,
            value=0,
            p=0.3,
        ),

        A.RandomBrightnessContrast(
            brightness_limit=0.15,
            contrast_limit=0.15,
            p=0.5,
        ),

        A.RandomGamma(
            gamma_limit=(85, 115),
            p=0.3,
        ),

        A.HueSaturationValue(
            hue_shift_limit=3,
            sat_shift_limit=10,
            val_shift_limit=10,
            p=0.25,
        ),

        A.OneOf([
            A.GaussianBlur(blur_limit=(3, 5)),
            A.Downscale(scale_min=0.80, scale_max=0.95),
            A.ImageCompression(quality_lower=75, quality_upper=100),
        ], p=0.2),

        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(),
    ])


# -------------------------------------------------------------------------
# Suggested CFP augmentation parameter sets
# -------------------------------------------------------------------------
#
# 1) DEFAULT / CONSERVATIVE
# Use this as a general starting point for CFP classification tasks.
#
# Rationale:
# - Simulates common camera/acquisition variability.
# - Keeps color and image-quality perturbations mild.
# - Good first choice when the disease signal may depend on subtle color,
#   contrast, texture, or anatomical context.
#
# brightness_limit = 0.15
# contrast_limit   = 0.15
# gamma_limit      = (85, 115)      # approximately gamma 0.85–1.15
# hue_shift_limit  = 3              # intentionally small for fundus color realism
# sat_shift_limit  = 10
# val_shift_limit  = 10
# rotate_limit     = 7
# shift_limit      = 0.02
# scale_limit      = 0.05
# blur_limit       = (3, 5)
# downscale_range  = (0.80, 0.95)
# jpeg_quality     = (75, 100)
#
#
# 2) MORE AGGRESSIVE / DOMAIN-ROBUSTNESS
# Use this when robustness across different CFP cameras, sites, image qualities,
# or acquisition pipelines is more important, and confirm using external or
# camera/site-held-out validation.
#
# Rationale:
# - Simulates broader variation across CFP devices and acquisition conditions.
# - May improve domain robustness.
# - Higher risk of altering disease-relevant appearance, so it should be
#   validated carefully for the target task.
#
# brightness_limit = 0.25
# contrast_limit   = 0.25
# gamma_limit      = (75, 130)      # approximately gamma 0.75–1.30
# hue_shift_limit  = 5              # still limited for fundus color realism
# sat_shift_limit  = 18
# val_shift_limit  = 18
# rotate_limit     = 12
# shift_limit      = 0.04
# scale_limit      = 0.10
# blur_limit       = (3, 7)
# downscale_range  = (0.65, 0.95)
# jpeg_quality     = (55, 100)
# -------------------------------------------------------------------------