from torchvision import transforms


def build_transform(split, input_size, auto_augment, interpolation, mean, std,
                    horizontal_flip_prob=0.5, vertical_flip_prob=0.5, rotation_range=10,
                    translate=(0.1, 0.1), scale=None, shear=None, erase_prob=0.1):
    if split == "train":
        transform_list = [
            transforms.RandomResizedCrop(input_size, interpolation=interpolation),
            transforms.RandomHorizontalFlip(horizontal_flip_prob),
            transforms.RandomVerticalFlip(vertical_flip_prob),
            transforms.RandomRotation(rotation_range),
            transforms.RandomAffine(rotation_range, translate, scale, shear),
            transforms.ToTensor(),
            transforms.Normalize(mean, std),
            transforms.RandomErasing(erase_prob),
        ]
        if auto_augment:
            transform_list.insert(1, transforms.AutoAugment())
        transform = transforms.Compose(transform_list)
    else:
        transform = transforms.Compose(
            [
                transforms.Resize(input_size[0] + 32),
                transforms.CenterCrop(input_size),
                transforms.ToTensor(),
                transforms.Normalize(mean, std),
            ]
        )
    return transform


mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
input_size=(224, 224)
train_transform = build_transform(split="train", input_size=input_size, auto_augment=False, interpolation=3,
                                  mean=mean, std=std,
                                  horizontal_flip_prob=0.5, vertical_flip_prob=0.5, rotation_range=10,
                                  translate=(0.1, 0.1), scale=None, shear=None, erase_prob=0.1)
val_transform = build_transform(split="val", input_size=input_size, auto_augment=False, interpolation=3,
                                mean=mean, std=std,

                                horizontal_flip_prob=0.5, vertical_flip_prob=0.5, rotation_range=10,
                                translate=(0.1, 0.1), scale=None, shear=None, erase_prob=0.1)

test_transform = build_transform(split="test", input_size=input_size, auto_augment=False, interpolation=3,
                                 mean=mean, std=std,
                                 horizontal_flip_prob=0.5, vertical_flip_prob=0.5, rotation_range=10,
                                 translate=(0.1, 0.1), scale=None, shear=None, erase_prob=0.1)
