import albumentations as alb
from albumentations.pytorch import ToTensorV2
import cv2

def get_augs(name):
    IMG_SIZE = 380  
    if name == "REAlbu":
        return alb.Compose([
            alb.HorizontalFlip(),
            alb.CoarseDropout(max_holes = 1, min_height=int(IMG_SIZE*0.02), max_height=int(IMG_SIZE*0.2), min_width=int(IMG_SIZE*0.02), max_width=int(IMG_SIZE*0.2), p=1),
        ])
    elif name == "RandCropAlbu":
        return alb.Compose([
            alb.HorizontalFlip(),
            alb.RandomResizedCrop(height = IMG_SIZE, width = IMG_SIZE, scale=(1/1.3, 1.0), ratio=(0.9,1.1)),
        ])
    elif name == "DFDCAlbu":
        return alb.Compose([
            alb.ImageCompression(quality_lower=60, quality_upper=100, p=0.5),
            alb.GaussNoise(p=0.1),
            alb.GaussianBlur(blur_limit=3, p=0.05),
            alb.HorizontalFlip(),
            alb.OneOf([
                alb.LongestMaxSize(max_size=IMG_SIZE, interpolation=cv2.INTER_CUBIC),
                alb.LongestMaxSize(max_size=IMG_SIZE, interpolation=cv2.INTER_AREA),
                alb.LongestMaxSize(max_size=IMG_SIZE, interpolation=cv2.INTER_LINEAR)
            ], p=1.0),
            alb.PadIfNeeded(min_height=IMG_SIZE, min_width=IMG_SIZE, border_mode=cv2.BORDER_CONSTANT),
            alb.OneOf([alb.RandomBrightnessContrast(), alb.FancyPCA(), alb.HueSaturationValue()], p=0.7),
            alb.ToGray(p=0.2),
            alb.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=10, border_mode=cv2.BORDER_CONSTANT, p=0.5),
        ])