"""
"""

from random import sample
from abc import ABC, abstractmethod
from typing import Optional, Tuple, Union, List, Any, Sequence, NoReturn

import numpy as np
import torch
from torch import Tensor

from .base import Augmenter
from .baseline_wander import BaselineWanderAugmenter
from .label_smooth import LabelSmooth
from .mixup import Mixup
from .random_flip import RandomFlip
from .random_masking import RandomMasking
from .random_renormalize import RandomRenormalize
from .stretch_compress import StretchCompress
from ..utils.misc import default_class_repr, ReprMixin


__all__ = [
    "AugmenterManager",
]


class AugmenterManager(torch.nn.Module):
    """

    The `Module` to manage the augmenters

    Examples
    --------
    ```python
    import torch
    from easydict import EasyDict as ED
    from torch_ecg.augmenters import AugmenterManager

    config = ED(
        random=False,
        fs=500,
        baseline_wander={},
        label_smooth={},
        mixup={},
        random_flip={},
        random_masking={},
        random_renormalize={},
        stretch_compress={},
    )
    am = AugmenterManager.from_config(config)
    sig, label, mask = torch.rand(2,12,5000), torch.rand(2,26), torch.rand(2,5000,1)
    sig, label, mask = am(sig, label, mask)
    ```
    """

    __name__ = "AugmenterManager"

    def __init__(
        self, *augs: Optional[Tuple[Augmenter, ...]], random: bool = False
    ) -> NoReturn:
        """finished, checked,

        Parameters
        ----------
        aug: tuple of `Augmenter`, optional,
            the augmenters to be added to the manager
        random: bool, default False,
            whether to apply the augmenters in random order
        """
        super().__init__()
        self.random = random
        self._augmenters = list(augs)

    def _add_baseline_wander(self, **config: dict) -> NoReturn:
        self._augmenters.append(BaselineWanderAugmenter(**config))

    def _add_label_smooth(self, **config: dict) -> NoReturn:
        self._augmenters.append(LabelSmooth(**config))

    def _add_mixup(self, **config: dict) -> NoReturn:
        self._augmenters.append(Mixup(**config))

    def _add_random_flip(self, **config: dict) -> NoReturn:
        self._augmenters.append(RandomFlip(**config))

    def _add_random_masking(self, **config: dict) -> NoReturn:
        self._augmenters.append(RandomMasking(**config))

    def _add_random_renormalize(self, **config: dict) -> NoReturn:
        self._augmenters.append(RandomRenormalize(**config))

    def _add_stretch_compress(self, **config: dict) -> NoReturn:
        self._augmenters.append(StretchCompress(**config))

    def forward(
        self,
        sig: Tensor,
        label: Optional[Tensor],
        *extra_tensors: Sequence[Tensor],
        **kwargs: Any,
    ) -> Union[Tensor, Tuple[Tensor]]:
        """finished, checked,

        Parameters
        ----------
        sig: Tensor,
            the ECGs to be augmented, of shape (batch, lead, siglen)
        label: Tensor, optional,
            labels of the ECGs
        extra_tensors: Tensor(s), optional,
            extra tensors to be augmented, e.g. masks for custom loss functions, etc.
        kwargs: keyword arguments

        Returns
        -------
        Tensor(s), the augmented ECGs, labels, and optional extra tensors
        """
        if len(self.augmenters) == 0:
            # raise ValueError("No augmenters added to the manager.")
            return (sig, label, *extra_tensors)
        ordering = list(range(len(self.augmenters)))
        if self.random:
            ordering = sample(ordering, len(ordering))
        for idx in ordering:
            sig, label, *extra_tensors = self.augmenters[idx](
                sig, label, *extra_tensors, **kwargs
            )
        return (sig, label, *extra_tensors)

    @property
    def augmenters(self) -> List[Augmenter]:
        return self._augmenters

    def extra_repr(self) -> str:
        indent = 4 * " "
        s = (
            f"augmenters = [\n{indent}"
            + f",\n{2*indent}".join(
                default_class_repr(aug, depth=2) for aug in self.augmenters
            )
            + f"{indent}\n]"
        )
        return s

    @classmethod
    def from_config(cls, config: dict) -> "AugmenterManager":
        """finished, checked,

        Parameters
        ----------
        config: dict,
            the configuration of the augmenters,
            better to be an `OrderedDict`

        Returns
        -------
        am: AugmenterManager,
            a new instance of `AugmenterManager`
        """
        am = cls(random=config.get("random", False))
        _mapping = {
            "baseline_wander": am._add_baseline_wander,
            "label_smooth": am._add_label_smooth,
            "mixup": am._add_mixup,
            "random_flip": am._add_random_flip,
            "random_masking": am._add_random_masking,
            "random_renormalize": am._add_random_renormalize,
            "stretch_compress": am._add_stretch_compress,
        }
        for aug_name, aug_config in config.items():
            if aug_name in [
                "fs",
                "random",
            ]:
                continue
            elif aug_name in _mapping and isinstance(aug_config, dict):
                _mapping[aug_name](fs=config["fs"], **aug_config)
            else:
                # just ignore the other items
                pass
                # raise ValueError(f"Unknown augmenter name: {aug_name}")
        return am

    def rearrange(self, new_ordering: List[str]) -> NoReturn:
        """finished, checked,

        Parameters
        ----------
        new_ordering: list of str,
            the list of augmenter names in the new order
        """
        _mapping = {
            "".join([w.capitalize() for w in k.split("_")]): k
            for k in "label_smooth,mixup,random_flip,random_masking,random_renormalize,stretch_compress".split(
                ","
            )
        }
        _mapping.update({"BaselineWanderAugmenter": "baseline_wander"})
        for k in new_ordering:
            if k not in _mapping:
                _mapping.update({k: k})
        self._augmenters.sort(
            key=lambda aug: new_ordering.index(_mapping[aug.__name__])
        )
