from PIL import Image

import mxnet as mx

from mxnet.gluon import Block

from ..transforms import *

class RandAugment(object):

    def __init__(self, n, m):

        self.n = n

        self.m = m

        self.augment_list = rand_augment_list()

        self.topil = ToPIL()

    def __call__(self, img):

        img = self.topil(img)

        ops = random.choices(self.augment_list, k=self.n)

        for op, minval, maxval in ops:

            if random.random() > random.uniform(0.2, 0.8):

                continue

            val = (float(self.m) / 30) * float(maxval - minval) + minval

            img = op(img, val)

        return img

class ToPIL(object):

    """Convert image from ndarray format to PIL

    """

    def __call__(self, img):

        x = Image.fromarray(img.asnumpy())

        return x

class ToNDArray(object):

    def __call__(self, img):

        x = mx.nd.array(np.array(img), mx.cpu(0))

        return x

class AugmentationBlock(Block):

    r"""

    AutoAugment Block

    Example

    -------

    >>> from autogluon.utils.augment import AugmentationBlock, autoaug_imagenet_policies

    >>> aa_transform = AugmentationBlock(autoaug_imagenet_policies())

    """

    def __init__(self, policies):

        """

        plicies : list of (name, pr, level)

        """

        super().__init__()

        self.policies = policies

        self.topil = ToPIL()

        self.tond = ToNDArray()

    def forward(self, img):

        img = self.topil(img)

        policy = random.choice(self.policies)

        for name, pr, level in policy:

            if random.random() > pr:

                continue

            img = apply_augment(img, name, level)

        img = self.tond(img)

        return img