#!/usr/bin/env python3
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
import copy
from typing import Optional, Sequence

import cv2
import megengine.data as data
import megengine.data.transform as T
import numpy as np
from basecore.config import ConfigDict
from loguru import logger

from basecls.utils import registers

from .augment import WARP_PARAMS, TorchAutoAugment, TorchRandAugment
from .const import CV2_INTERP, PIL_INTERP
from .mixup import MixupCutmixCollator
from .rand_erase import RandomErasing

__all__ = [
    "build_transform",
    "AutoAugment",
    "SimpleAugment",
    "ColorAugment",
    "RandAugment",
    "build_mixup",
]


def build_transform(
    cfg: ConfigDict, train: bool = True, augments: T.Transform = None
) -> T.Transform:
    """Build function for MegEngine transform.

    Args:
        cfg: config for building transform.
        train: train set or test set. Default: ``True``
        augments: augments for building transform.

    Returns:
        A transform.
    """
    if train:
        assert augments is not None
        bgr_mean = copy.deepcopy(cfg.preprocess.img_mean)
        bgr_std = copy.deepcopy(cfg.preprocess.img_std)
        if cfg.preprocess.img_color_space == "RGB":
            bgr_mean = bgr_mean[::-1]
            bgr_std = bgr_std[::-1]
        WARP_PARAMS["fillcolor"] = tuple(round(v) for v in bgr_mean[::-1])  # need RGB
        WARP_PARAMS["resample"] = PIL_INTERP[cfg.augments.resize.interpolation]

        transforms = [
            T.RandomResizedCrop(
                cfg.preprocess.img_size,
                cfg.augments.resize.scale_range,
                cfg.augments.resize.ratio_range,
                CV2_INTERP[cfg.augments.resize.interpolation],
            ),
            
T.RandomHorizontalFlip()