from albumentations import HorizontalFlip, OneOf, IAAPiecewiseAffine, IAAAffine, OpticalDistortion, GridDistortion
from albumentations.core.composition import Transforms
import pytest


def empty_aug1():
    return [
        HorizontalFlip(p=0.001),
        # IAAPiecewiseAffine(p=1.0),
        OneOf(
            [
                # OpticalDistortion(p=0.1),
                # GridDistortion(p=0.1),
                # IAAPerspective(p=1.0),
                # IAAAffine(p=1.0),
                IAAPiecewiseAffine(p=1.0)
            ],
            p=0.0,
        ),
    ]


def empty_aug2():
    return [
        HorizontalFlip(p=0.001),
        IAAPiecewiseAffine(p=1.0),
        OneOf(
            [
                # OpticalDistortion(p=0.1),
                # GridDistortion(p=0.1),
                # IAAPerspective(p=1.0),
                IAAAffine(p=1.0),
                # IAAPiecewiseAffine(p=1.0),
            ],
            p=0.0,
        ),
    ]


def empty_aug3():
    return [
        # HorizontalFlip(p=0.001),
        # IAAPiecewiseAffine(p=1.0),
        OneOf(
            [
                OpticalDistortion(p=0.1),
                GridDistortion(p=0.1),
                # IAAPerspective(p=1.0),
                # IAAAffine(p=1.0),
                # IAAPiecewiseAffine(p=1.0),
            ],
            p=0.0,
        )
    ]


@pytest.mark.parametrize(["aug", "start_end"], [[empty_aug1, [0, 1]], [empty_aug2, [0, 2]], [empty_aug3, [0, 0]]])
def test_strong_aug(aug, start_end):
    assert Transforms(aug()).start_end == start_end
