# Copyright (c) OpenMMLab. All rights reserved.
import collections

import cv2
from mmcv.utils import build_from_cfg

from ..builder import PIPELINES
from .transforms import Albu, get_albu_transform

@PIPELINES.register_module()
class Compose:
    """Compose multiple transforms sequentially.

    Args:
        transforms (Sequence[dict | callable]): Sequence of transform object or
            config dict to be composed.
    """

    def __init__(self, transforms):
        assert isinstance(transforms, collections.abc.Sequence)
        self.transforms = []
        self.albu_prob = 0        # abcdefghijk 自定义AlbuAug
        for transform in transforms:
            # if transform.type == 'Albu':
            #     self.albu_prob = transform.pop('probability')
            if isinstance(transform, dict):
                transform = build_from_cfg(transform, PIPELINES)
                self.transforms.append(transform)
            elif callable(transform):
                self.transforms.append(transform)
            else:
                raise TypeError('transform must be callable or a dict')

    def __call__(self, data):
        """Call function to apply transforms sequentially.

        Args:
            data (dict): A result dict contains the data to transform.

        Returns:
           dict: Transformed data.
        """

        for t in self.transforms:
            if isinstance(t, Albu) and self.albu_prob != 0:
                albrans = get_albu_transform(p=self.albu_prob)
                data['image'] = albrans(image=data['img'])['image']
                # cv2.imshow('albrans', data['image'])
                # cv2.waitKey(-1)
                continue
            data = t(data)
            if data is None:
                return None
        return data

    def __repr__(self):
        format_string = self.__class__.__name__ + '('
        for t in self.transforms:
            str_ = t.__repr__()
            if 'Compose(' in str_:
                str_ = str_.replace('\n', '\n    ')
            format_string += '\n'
            format_string += f'    {str_}'
        format_string += '\n)'
        return format_string
