import logging

import numpy as np
import scipy.sparse
import tensorflow as tf

from sklearn.model_selection import train_test_split
from typing import Optional, Dict, Text, List, Tuple, Any, Union, Generator, NamedTuple
from collections import defaultdict
from rasa.utils.tensorflow.constants import BALANCED, SEQUENCE

logger = logging.getLogger(__name__)


# Mapping of feature name to a list of numpy arrays representing the actual features
# For example:
# "text_features" -> [
#   "numpy array containing dense features for every training example",
#   "numpy array containing sparse features for every training example"
# ]
Data = Dict[Text, List[np.ndarray]]


class FeatureSignature(NamedTuple):
    """Stores the shape and the type (sparse vs dense) of features."""

    is_sparse: bool
    shape: List[int]


class RasaModelData:
    """Data object used for all RasaModels.

    It contains all features needed to train the models.
    """

    def __init__(
        self, label_key: Optional[Text] = None, data: Optional[Data] = None
    ) -> None:
        """
        Initializes the RasaModelData object.

        Args:
            label_key: the label_key used for balancing, etc.
            data: the data holding the features
        """

        self.data = data or {}
        self.label_key = label_key
        # should be updated when features are added
        self.num_examples = self.number_of_examples()

    def get_only(self, key: Text) -> Optional[np.ndarray]:
        if key in self.data:
            return self.data[key][0]
        else:
            return None

    def get(self, key: Text) -> List[np.ndarray]:
        if key in self.data:
            return self.data[key]
        else:
            return []

    def items(self):
        return self.data.items()

    def values(self):
        return self.data.values()

    def keys(self):
        return self.data.keys()

    def first_data_example(self) -> Data:
        return {
            feature_name: [feature[:1] for feature in features]
            for feature_name, features in self.data.items()
        }

    def feature_not_exist(self, key: Text) -> bool:
        """Check if feature key is present and features are available."""

        return key not in self.data or not self.data[key]

    def is_empty(self) -> bool:
        """Checks if data is set."""

        return not self.data

    def number_of_examples(self, data: Optional[Data] = None) -> int:
        """Obtain number of examples in data.

        Raises: A ValueError if number of examples differ for different features.
        """

        if not data:
            data = self.data

        if not data:
            return 0

        example_lengths = [v.shape[0] for values in data.values() for v in values]

        # check if number of examples is the same for all values
        if not all(length == example_lengths[0] for length in example_lengths):
            raise ValueError(
                f"Number of examples differs for keys '{data.keys()}'. Number of "
                f"examples should be the same for all data."
            )

        return example_lengths[0]

    def feature_dimension(self, key: Text) -> int:
        """Get the feature dimension of the given key."""

        number_of_features = 0
        for data in self.data[key]:
            if data.size > 0:
                number_of_features += data[0].shape[-1]

        return number_of_features

    def add_features(self, key: Text, features: List[np.ndarray]):
        """Add list of features to data under specified key.

        Should update number of examples.
        """

        if not features:
            return

        if key in self.data:
            raise ValueError(f"Key '{key}' already exists in RasaModelData.")

        self.data[key] = []

        for data in features:
            if data.size > 0:
                self.data[key].append(data)

        if not self.data[key]:
            del self.data[key]

        # update number of examples
        self.num_examples = self.number_of_examples()

    def add_lengths(self, key: Text, from_key: Text) -> None:
        """Adds np.array of lengths of sequences to data under given key."""
        if not self.data.get(from_key):
            return

        self.data[key] = []

        for data in self.data[from_key]:
            if data.size > 0:
                lengths = np.array([x.shape[0] for x in data])
                self.data[key].append(lengths)
                break

    def split(
        self, number_of_test_examples: int, random_seed: int
    ) -> Tuple["RasaModelData", "RasaModelData"]:
        """Create random hold out test set using stratified split."""

        self._check_label_key()

        if self.label_key is None:
            # randomly split data as no label key is split
            multi_values = [v for values in self.data.values() for v in values]
            solo_values = [[] for values in self.data.values() for v in values]
            stratify = None
        else:
            # make sure that examples for each label value are in both split sets
            label_ids = self._create_label_ids(self.data[self.label_key][0])
            label_counts = dict(zip(*np.unique(label_ids, return_counts=True, axis=0)))

            self._check_train_test_sizes(number_of_test_examples, label_counts)

            counts = np.array([label_counts[label] for label in label_ids])
            # we perform stratified train test split,
            # which insures every label is present in the train and test data
            # this operation can be performed only for labels
            # that contain several data points
            multi_values = [
                v[counts > 1] for values in self.data.values() for v in values
            ]
            # collect data points that are unique for their label
            solo_values = [
                v[counts == 1] for values in self.data.values() for v in values
            ]

            stratify = label_ids[counts > 1]

        output_values = train_test_split(
            *multi_values,
            test_size=number_of_test_examples,
            random_state=random_seed,
            stratify=stratify,
        )

        return self._convert_train_test_split(output_values, solo_values)

    def get_signature(self) -> Dict[Text, List[FeatureSignature]]:
        """Get signature of RasaModelData.

        Signature stores the shape and whether features are sparse or not for every key.
        """

        return {
            key: [
                FeatureSignature(
                    True if isinstance(v[0], scipy.sparse.spmatrix) else False,
                    v[0].shape,
                )
                for v in values
            ]
            for key, values in self.data.items()
        }

    def as_tf_dataset(
        self, batch_size: int, batch_strategy: Text = SEQUENCE, shuffle: bool = False
    ) -> tf.data.Dataset:
        """Create tf dataset."""

        shapes, types = self._get_shapes_types()

        return tf.data.Dataset.from_generator(
            lambda batch_size_: self._gen_batch(batch_size_, batch_strategy, shuffle),
            output_types=types,
            output_shapes=shapes,
            args=([batch_size]),
        )

    def prepare_batch(
        self,
        data: Optional[Data] = None,
        start: Optional[int] = None,
        end: Optional[int] = None,
        tuple_sizes: Optional[Dict[Text, int]] = None,
    ) -> Tuple[Optional[np.ndarray]]:
        """Slices model data into batch using given start and end value."""

        if not data:
            data = self.data

        batch_data = []

        for key, values in data.items():
            # add None for not present values during processing
            if not values:
                if tuple_sizes:
                    batch_data += [None] * tuple_sizes[key]
                else:
                    batch_data.append(None)
                continue

            for v in values:
                if start is not None and end is not None:
                    _data = v[start:end]
                elif start is not None:
                    _data = v[start:]
                elif end is not None:
                    _data = v[:end]
                else:
                    _data = v[:]

                if isinstance(_data[0], scipy.sparse.spmatrix):
                    batch_data.extend(self._scipy_matrix_to_values(_data))
                else:
                    batch_data.append(self._pad_dense_data(_data))

        # len of batch_data is equal to the number of keys in model data
        return tuple(batch_data)

    def _get_shapes_types(self) -> Tuple:
        """Extract shapes and types from model data."""

        types = []
        shapes = []

        def append_shape(features: np.ndarray) -> None:
            if isinstance(features[0], scipy.sparse.spmatrix):
                # scipy matrix is converted into indices, data, shape
                shapes.append((None, features[0].ndim + 1))
                shapes.append((None,))
                shapes.append((features[0].ndim + 1))
            elif features[0].ndim == 0:
                shapes.append((None,))
            elif features[0].ndim == 1:
                shapes.append((None, features[0].shape[-1]))
            else:
                shapes.append((None, None, features[0].shape[-1]))

        def append_type(features: np.ndarray) -> None:
            if isinstance(features[0], scipy.sparse.spmatrix):
                # scipy matrix is converted into indices, data, shape
                types.append(tf.int64)
                types.append(tf.float32)
                types.append(tf.int64)
            else:
                types.append(tf.float32)

        for values in self.data.values():
            for v in values:
                append_shape(v)
                append_type(v)

        return tuple(shapes), tuple(types)

    def _shuffled_data(self, data: Data) -> Data:
        """Shuffle model data."""

        ids = np.random.permutation(self.num_examples)
        return self._data_for_ids(data, ids)

    def _balanced_data(self, data: Data, batch_size: int, shuffle: bool) -> Data:
        """Mix model data to account for class imbalance.

        This batching strategy puts rare classes approximately in every other batch,
        by repeating them. Mimics stratified batching, but also takes into account
        that more populated classes should appear more often.
        """

        self._check_label_key()

        # skip balancing if labels are token based
        if self.label_key is None or data[self.label_key][0][0].size > 1:
            return data

        label_ids = self._create_label_ids(data[self.label_key][0])

        unique_label_ids, counts_label_ids = np.unique(
            label_ids, return_counts=True, axis=0
        )
        num_label_ids = len(unique_label_ids)

        # group data points by their label
        # need to call every time, so that the data is shuffled inside each class
        data_by_label = self._split_by_label_ids(data, label_ids, unique_label_ids)

        # running index inside each data grouped by labels
        data_idx = [0] * num_label_ids
        # number of cycles each label was passed
        num_data_cycles = [0] * num_label_ids
        # if a label was skipped in current batch
        skipped = [False] * num_label_ids

        new_data = defaultdict(list)

        while min(num_data_cycles) == 0:
            if shuffle:
                indices_of_labels = np.random.permutation(num_label_ids)
            else:
                indices_of_labels = range(num_label_ids)

            for index in indices_of_labels:
                if num_data_cycles[index] > 0 and not skipped[index]:
                    skipped[index] = True
                    continue
                else:
                    skipped[index] = False

                index_batch_size = (
                    int(counts_label_ids[index] / self.num_examples * batch_size) + 1
                )

                for k, values in data_by_label[index].items():
                    for i, v in enumerate(values):
                        if len(new_data[k]) < i + 1:
                            new_data[k].append([])
                        new_data[k][i].append(
                            v[data_idx[index] : data_idx[index] + index_batch_size]
                        )

                data_idx[index] += index_batch_size
                if data_idx[index] >= counts_label_ids[index]:
                    num_data_cycles[index] += 1
                    data_idx[index] = 0

                if min(num_data_cycles) > 0:
                    break

        final_data = defaultdict(list)
        for k, values in new_data.items():
            for v in values:
                final_data[k].append(np.concatenate(np.array(v)))

        return final_data

    def _gen_batch(
        self, batch_size: int, batch_strategy: Text = SEQUENCE, shuffle: bool = False
    ) -> Generator[Tuple[Optional[np.ndarray]], None, None]:
        """Generate batches."""

        data = self.data
        num_examples = self.num_examples

        if shuffle:
            data = self._shuffled_data(data)

        if batch_strategy == BALANCED:
            data = self._balanced_data(data, batch_size, shuffle)
            # after balancing, number of examples increased
            num_examples = self.number_of_examples(data)

        num_batches = num_examples // batch_size + int(num_examples % batch_size > 0)

        for batch_num in range(num_batches):
            start = batch_num * batch_size
            end = start + batch_size

            yield self.prepare_batch(data, start, end)

    def _check_train_test_sizes(
        self, number_of_test_examples: int, label_counts: Dict[Any, int]
    ):
        """Check whether the test data set is too large or too small."""

        if number_of_test_examples >= self.num_examples - len(label_counts):
            raise ValueError(
                f"Test set of {number_of_test_examples} is too large. Remaining "
                f"train set should be at least equal to number of classes "
                f"{len(label_counts)}."
            )
        elif number_of_test_examples < len(label_counts):
            raise ValueError(
                f"Test set of {number_of_test_examples} is too small. It should "
                f"be at least equal to number of classes {label_counts}."
            )

    @staticmethod
    def _data_for_ids(data: Optional[Data], ids: np.ndarray) -> Data:
        """Filter model data by ids."""

        new_data = defaultdict(list)

        if data is None:
            return new_data

        for k, values in data.items():
            for v in values:
                new_data[k].append(v[ids])
        return new_data

    def _split_by_label_ids(
        self, data: Optional[Data], label_ids: np.ndarray, unique_label_ids: np.ndarray
    ) -> List["RasaModelData"]:
        """Reorganize model data into a list of model data with the same labels."""

        label_data = []
        for label_id in unique_label_ids:
            matching_ids = label_ids == label_id
            label_data.append(
                RasaModelData(self.label_key, self._data_for_ids(data, matching_ids))
            )
        return label_data

    def _check_label_key(self):
        if self.label_key is not None and (
            self.label_key not in self.data or len(self.data[self.label_key]) > 1
        ):
            raise ValueError(f"Key '{self.label_key}' not in RasaModelData.")

    def _convert_train_test_split(
        self, output_values: List[Any], solo_values: List[Any]
    ) -> Tuple["RasaModelData", "RasaModelData"]:
        """Converts the output of sklearn's train_test_split into model data."""

        data_train = defaultdict(list)
        data_val = defaultdict(list)

        # output_values = x_train, x_val, y_train, y_val, z_train, z_val, etc.
        # order is kept, e.g. same order as model data keys

        # train datasets have an even index
        index = 0
        for key, values in self.data.items():
            for _ in values:
                data_train[key].append(
                    self._combine_features(output_values[index * 2], solo_values[index])
                )
                index += 1

        # val datasets have an odd index
        index = 0
        for key, values in self.data.items():
            for _ in range(len(values)):
                data_val[key].append(output_values[(index * 2) + 1])
                index += 1

        return (
            RasaModelData(self.label_key, data_train),
            RasaModelData(self.label_key, data_val),
        )

    @staticmethod
    def _combine_features(
        feature_1: Union[np.ndarray, scipy.sparse.spmatrix],
        feature_2: Union[np.ndarray, scipy.sparse.spmatrix],
    ) -> Union[np.ndarray, scipy.sparse.spmatrix]:
        """Concatenate features."""

        if isinstance(feature_1, scipy.sparse.spmatrix) and isinstance(
            feature_2, scipy.sparse.spmatrix
        ):
            if feature_2.shape[0] == 0:
                return feature_1
            if feature_1.shape[0] == 0:
                return feature_2
            return scipy.sparse.vstack([feature_1, feature_2])

        return np.concatenate([feature_1, feature_2])

    @staticmethod
    def _create_label_ids(label_ids: np.ndarray) -> np.ndarray:
        """Convert various size label_ids into single dim array.

        For multi-label y, map each distinct row to a string representation
        using join because str(row) uses an ellipsis if len(row) > 1000.
        Idea taken from sklearn's stratify split.
        """

        if label_ids.ndim == 1:
            return label_ids

        if label_ids.ndim == 2 and label_ids.shape[-1] == 1:
            return label_ids[:, 0]

        if label_ids.ndim == 2:
            return np.array([" ".join(row.astype("str")) for row in label_ids])

        if label_ids.ndim == 3 and label_ids.shape[-1] == 1:
            return np.array([" ".join(row.astype("str")) for row in label_ids[:, :, 0]])

        raise ValueError("Unsupported label_ids dimensions")

    @staticmethod
    def _pad_dense_data(array_of_dense: np.ndarray) -> np.ndarray:
        """Pad data of different lengths.

        Sequential data is padded with zeros. Zeros are added to the end of data.
        """

        if array_of_dense[0].ndim < 2:
            # data doesn't contain a sequence
            return array_of_dense.astype(np.float32)

        data_size = len(array_of_dense)
        max_seq_len = max([x.shape[0] for x in array_of_dense])

        data_padded = np.zeros(
            [data_size, max_seq_len, array_of_dense[0].shape[-1]],
            dtype=array_of_dense[0].dtype,
        )
        for i in range(data_size):
            data_padded[i, : array_of_dense[i].shape[0], :] = array_of_dense[i]

        return data_padded.astype(np.float32)

    @staticmethod
    def _scipy_matrix_to_values(array_of_sparse: np.ndarray) -> List[np.ndarray]:
        """Convert a scipy matrix into indices, data, and shape."""

        # we need to make sure that the matrices are coo_matrices otherwise the
        # transformation does not work (e.g. you cannot access x.row, x.col)
        if not isinstance(array_of_sparse[0], scipy.sparse.coo_matrix):
            array_of_sparse = [x.tocoo() for x in array_of_sparse]

        max_seq_len = max([x.shape[0] for x in array_of_sparse])

        # get the indices of values
        indices = np.hstack(
            [
                np.vstack([i * np.ones_like(x.row), x.row, x.col])
                for i, x in enumerate(array_of_sparse)
            ]
        ).T

        data = np.hstack([x.data for x in array_of_sparse])

        number_of_features = array_of_sparse[0].shape[-1]
        shape = np.array((len(array_of_sparse), max_seq_len, number_of_features))

        return [
            indices.astype(np.int64),
            data.astype(np.float32),
            shape.astype(np.int64),
        ]
