from .augmentation import Augmentation #find the imageAugmentor equilavent in dectron2 
from fvcore.transforms.transform import Transform

__all__ = ['Albumentations']


class AlbumentationsTransform(Transform):
    def __init__(self, aug, param):
        self.aug = aug
        self.param = param

    def apply_image(self, img):
        return self.aug.apply(img, **self.param)

    def apply_coords(self, coords):
        return self.coords

class Albumentations(Augmentation):
    """
    Wrap an augmentor form the albumentations library: https://github.com/albu/albumentations.
    Coordinate augmentation is not supported by the library.
    Example:
    .. code-block:: python
        from tensorpack import imgaug
        # or from dataflow import imgaug  # if you're using the standalone version of dataflow
        import albumentations as AB
        myaug = imgaug.Albumentations(AB.RandomRotate90(p=1))
    """

    def __init__(self, augmentor):
        """
        Args:
            augmentor (albumentations.BasicTransform):
        """
        #super(Albumentations, self).__init__() - using python > 3.7 no need to call rng
        self._aug = augmentor
    def get_transform(self, img):
        return AlbumentationsTransform(self._aug, self._aug.get_params())
'''
    def __init__(self, augmentor):
        """
        Args:
            augs (list[Augmentation or Transform]):
        """
        super().__init__()
        self.aug = super._transform_to_aug(augmentor)
    def __call__(self, aug_input) -> Transform:
        tfm = self.aug(aug_input)
        return Transform(tfm)
    def __repr__(self):
        msgs = str(self.aug)
        return "Augmentation[{}]".format(", ".join(msgs))
    __str__ = __repr__
'''