import numpy as np
from lazy_dataset.core import DynamicTimeSeriesBucket
from padertorch.utils import to_list


class DynamicExtendedTimeSeriesBucket(DynamicTimeSeriesBucket):
    def __init__(
            self,
            init_example,
            min_label_diversity=0,
            label_key=None,
            multi_hot_encoded_labels=True,
            min_dataset_examples=None,
            bucket_id=None,
            **kwargs
    ):
        """
        Extension of the DynamicTimeSeriesBucket such that examples are
        balanced with respect to the dataset they originate from

        Args:
            init_example: first example in the bucket
            dataset_sizes:
            **kwargs: kwargs of DynamicTimeSeriesBucket
        """
        super().__init__(init_example, **kwargs)
        self.example_ids = set()
        self.min_label_diversity = min_label_diversity
        if min_label_diversity > 0:
            assert label_key is not None
            assert min_label_diversity <= self.batch_size
        self.label_key = label_key
        self.multi_hot_encoded_labels = multi_hot_encoded_labels
        self.label_classes = set()
        if min_dataset_examples is not None:
            self.missing_dataset_examples = {
                key: value for key, value in min_dataset_examples.items()
            }
            assert sum(min_dataset_examples.values()) <= self.batch_size
        else:
            self.missing_dataset_examples = None
        self.bucket_id = bucket_id

    def assess(self, example):
        if example["example_id"] in self.example_ids:
            return False
        if self.bucket_id is not None and example[self.bucket_id] != self.data[0][self.bucket_id]:
            return False
        if not super().assess(example):
            return False
        if self.missing_dataset_examples is not None:
            dataset_names = example['dataset'].split('+')  # '+' indicates mixtures
            assert all([name in self.missing_dataset_examples for name in dataset_names]), (
                dataset_names, sorted(self.missing_dataset_examples.keys())
            )
            if not (
                (self.batch_size - len(self.data)) > sum(self.missing_dataset_examples.values())
                or
                any([self.missing_dataset_examples[name] > 0 for name in dataset_names])
            ):
                return False
        if self.label_key is not None:
            labels = self._get_labels(example)
            if not (
                (self.batch_size - len(self.data)) > (self.min_label_diversity - len(self.label_classes))
                or
                any([label not in self.label_classes for label in labels])
            ):
                return False
        return True

    def _append(self, example):
        super()._append(example)
        self.example_ids.add(example["example_id"])
        if self.missing_dataset_examples is not None:
            for name in example['dataset'].split('+'):
                if self.missing_dataset_examples[name] > 0:
                    self.missing_dataset_examples[name] -= 1
        if self.label_key is not None and self.label_key in example:
            self.label_classes.update(self._get_labels(example))

    def _get_labels(self, example):
        labels = example[self.label_key]
        if self.multi_hot_encoded_labels:
            assert labels.ndim >= 1, labels.shape
            if labels.ndim > 1:
                assert labels.ndim == 2, labels.shape
                labels = labels.sum(-1)
            labels = np.argwhere(labels > 0).flatten()
        if isinstance(labels, np.ndarray):
            labels = labels.tolist()
        return to_list(labels)


def split_dataset(dataset, fold, nfolds=5, seed=0):
    """

    Args:
        dataset:
        fold:
        nfolfds:
        seed:

    Returns:

    >>> split_dataset(np.array([1,2,3,4,5]), 0, nfolds=2)
    [array([2, 4, 5]), array([1, 3])]
    >>> split_dataset(np.array([1,2,3,4,5]), 1, nfolds=2)
    [array([1, 3]), array([2, 4, 5])]
    """
    indices = np.arange(len(dataset))
    if seed is not None:
        np.random.RandomState(seed).shuffle(indices)
    folds = np.split(
        indices,
        np.linspace(0, len(dataset), nfolds + 1)[1:-1].astype(np.int64)
    )
    validation_indices = folds.pop(fold)
    training_indices = np.concatenate(folds)
    return [
        dataset[sorted(indices.tolist())]
        for indices in (training_indices, validation_indices)
    ]
