import sys
from typing import Union, Counter, Tuple

import numpy as np
from sklearn.preprocessing import LabelBinarizer


def simple_majority_voting(array: np.ndarray):
    if array.size == 0:
        return None
    c = Counter(array)
    return c.most_common(1)[0][0]


def shuffle_in_unison(*arrays: np.ndarray, seed: int | None = None, axis: int = 0):
    """
    Shuffle any number of arrays with the same random permutation along the given axis.
    Returns shuffled copies (does not modify inputs) and the permutation used.

    All arrays must have the same length along `axis`.
    """
    if not arrays:
        raise ValueError("At least one array must be provided.")

    length = arrays[0].shape[axis]
    for i, arr in enumerate(arrays[1:], start=1):
        if arr.shape[axis] != length:
            raise ValueError(
                f"Length mismatch along axis {axis}: "
                f"array[0] has {length}, but array[{i}] has {arr.shape[axis]}"
            )

    rng = np.random.default_rng(seed)
    perm = rng.permutation(length)
    shuffled = [np.take(arr, perm, axis=axis) for arr in arrays]
    return *shuffled, perm


def _unison_shuffle_np_arrays_3(a: np.ndarray, b: np.ndarray, c: np.ndarray, seed: int | None = None):
    """Shuffle three arrays in unison along axis 0 and return the shuffled arrays.
    This is a thin wrapper around shuffle_in_unison to match existing expectations.
    """
    a_s, b_s, c_s, _ = shuffle_in_unison(a, b, c, seed=seed, axis=0)
    return a_s, b_s, c_s


class Dataset:
    # Data to train on. In many cases this is considered to be equivalent to the enrolment data, but for distance
    # learning, these sets can be different.
    _trainset_data: np.ndarray = None
    _trainset_labels: np.ndarray = None

    # Data to validate on. Will be set to the same as the training data if not specified otherwise.
    _validationset_data: np.ndarray = None
    _validationset_labels: np.ndarray = None

    # Actual biometric enrollment data. Will be set to the training dawta if not specified otherwise.
    _testset_enrollment_data: np.ndarray = None
    _testset_enrollment_labels: np.ndarray = None

    # ACtual biometric matching (testset) data.
    _testset_matching_data: np.ndarray = None
    _testset_matching_labels: np.ndarray = None

    def __init__(self,
                 trainset_data: np.ndarray,
                 trainset_labels: np.ndarray,
                 validationset_data: np.ndarray = None,
                 validationset_labels: np.ndarray = None,
                 testset_enrollment_data: np.ndarray = None,
                 testset_enrollment_labels: np.ndarray = None,
                 testset_matching_data: np.ndarray = None,
                 testset_matching_labels: np.ndarray = None,
                 normalization_check: bool = False):
        self._trainset_data = self._check_and_convert_type_data(trainset_data, name="trainset_data")
        self._trainset_labels = self._check_and_convert_type_label(trainset_labels, name="trainset_labels")
        self._validationset_data = self._check_and_convert_type_data(validationset_data, name="validationset_data")
        self._validationset_labels = self._check_and_convert_type_label(validationset_labels,
                                                                       name="validationset_labels")
        self._testset_enrollment_data = self._check_and_convert_type_data(testset_enrollment_data,
                                                                         name="testset_enrollment_data")
        self._testset_enrollment_labels = self._check_and_convert_type_label(testset_enrollment_labels,
                                                                            name="testset_enrollment_labels")
        self._testset_matching_data = self._check_and_convert_type_data(testset_matching_data,
                                                                       name="testset_matching_data")
        self._testset_matching_labels = self._check_and_convert_type_label(testset_matching_labels,
                                                                          name="testset_matching_labels")

        # perform checks of data
        self._check_data_pairs_provided()  # check if corresponding pairs were provided,
        # no array is empty and corresponding pairs are of the same length
        self._check_nan()  # check if any NaNs exists
        self._check_dataset_is_balanced()
        if normalization_check:
            self._normalisation_check()  # check if data is normalized in [-1, 1]

        self._setup_missing_sets()
        self._setup_reject_label()

        # OHC Encoding
        self.ohc_encoder = LabelBinarizer()
        self.ohc_encoder.fit(self._gather_labels())

    def _setup_missing_sets(self):
        if self._validationset_data is None:
            if self._testset_enrollment_data is None:
                raise ValueError(f"Critical: a validation set needs to be present.")
            else:
                self._validationset_data = self._testset_enrollment_data
                self._validationset_labels = self._testset_enrollment_labels
                print(f"Warning: because the validation set is None, the supplied testset enrollment data will "
                      f"be used instead.", file=sys.stderr)
        if self._testset_enrollment_data is None:
            if self._testset_matching_data is None:
                print(f"Warning: testset enrollment and matching data are None. This will cause all evaluation "
                      f"methods to fail.", file=sys.stderr)
            else:
                self._testset_enrollment_data = self._testset_matching_data
                self._testset_enrollment_labels = self._testset_matching_labels
                print(f"Warning: because the enrollment testset is None, the matching testset will be used "
                      f"instead.", file=sys.stderr)
        else:
            if self._testset_matching_data is None:
                print(f"Warning: testset matching data is None. This will cause all evaluation "
                      f"methods to fail.", file=sys.stderr)

    def _setup_reject_label(self):
        self.reject_label = 'reject'  # In some cases this will be changed further down in the method
        if self._testset_matching_labels is not None:
            try:
                match_label_set = set(self._testset_matching_labels)
                enroll_label_set = set(self._testset_enrollment_labels)
            except TypeError:  # np.ndarray is not hashable
                match_label_set = set(self._testset_matching_labels.tolist())
                enroll_label_set = set(self._testset_enrollment_labels.tolist())
            diff = match_label_set.difference(enroll_label_set)
            if len(diff) == 1:
                self.reject_label = diff.pop()
            else:
                for user_to_reject in diff:  # This does nothing if len(diff) == 0, but it also doesn't hurt
                    self._testset_matching_labels[self._testset_matching_labels == user_to_reject] = self.reject_label

    def view_one_hot_encoded_labels(self):
        def convert_to_onehot(labels: np.ndarray):
            return None if labels is None else self._label_binarizer.transform(labels)

        return convert_to_onehot(self._trainset_labels), convert_to_onehot(self._validationset_labels), \
               convert_to_onehot(self._testset_enrollment_labels), convert_to_onehot(self._testset_matching_labels)

    @staticmethod
    def _check_and_convert_type_data(obj, name) -> np.array:
        """Check the types of objects, print warnings and perform conversions."""
        return Dataset._check_and_convert_type(obj, name, 'float32')

    @staticmethod
    def _check_and_convert_type_label(obj, name) -> np.array:
        """Check the types of objects, print warnings and perform conversions."""
        result = Dataset._check_and_convert_type(obj, name, str)
        if result is not None and len(result.shape) != 1:
            raise ValueError("Critical: label arrays must be one-dimensional.")
        return result

    @staticmethod
    def _check_and_convert_type(obj, name, dtype):
        if obj is None:
            return None
        if not isinstance(obj, np.ndarray):
            print(f"Warning: provided object {name} is of type {str(type(obj))} (expected: np.ndarray). "
                  f"Attempting a cast to np.ndarray.", file=sys.stderr)
            result = np.array(obj, dtype=dtype)  # this fails if dtype is float and the resulting array would be ragged
        elif obj.dtype != dtype:
            print(f"Warning: provided object {name} contains values of type {obj.dtype} (expected: {dtype}). "
                  f"Attempting a cast to expected datatype!", file=sys.stderr)
            result = obj.astype(dtype=dtype, copy=False)
        else:
            result = obj
        return result

    def _check_data_pairs_provided(self):
        """Checks that if training data is provided also labels are provided. Raises exception if not."""

        # check for None
        if self._trainset_data is None or len(self._trainset_data) == 0:
            raise ValueError("Critical: _trainset_data must not be None or empty.")

        if self._trainset_labels is None:
            raise ValueError("Critical: _trainset_labels must not be None.")

        def _check(data, labels, name):
            if data is not None:
                if labels is None:
                    raise ValueError(f"Critical: if {name}_data is not None, {name}_labels must also be not None"
                                     f"However, {name}_labels is None! You rascal!")
                if len(data) == 0:
                    raise ValueError(f"Critical: length of {name}_data is {len(self._trainset_labels)}, "
                                     f"but must be greater than 0!")
                if len(data) != len(labels):
                    raise ValueError(f"Critical: length of {name}_data ({len(self._trainset_data)}) is not equal to the "
                                     f"length of {name}_labels ({len(self._trainset_labels)})!")

        _check(self._validationset_data, self._validationset_labels, 'validationset')
        _check(self._testset_enrollment_data, self._testset_enrollment_labels, 'testset_enrollment')
        _check(self._testset_matching_data, self._testset_matching_labels, 'testset_matching')

    def _check_nan(self):
        def _check_nan(array: np.ndarray, name):
            try:
                if array is not None and np.isnan(array).any():
                    raise ValueError(f'Critical: provided array {name} contains NaN values!')
            except TypeError:
                print(f'Warning: could not check if NaNs exist in array {name}, as its dtype is "{array.dtype}". '
                      f'A numerical dtype (e.g., int or float) is expected. This might lead to subsequent '
                      f'critical errors and wrong calculations.', file=sys.stderr)

        _check_nan(self._trainset_data, name="_trainset_data")
        _check_nan(self._testset_matching_data, name="_testset_matching_data")
        _check_nan(self._testset_enrollment_data, name="_testset_enrollment_data")
        _check_nan(self._validationset_data, name="_validationset_data")

    def _check_dataset_is_balanced(self):
        return True  # TODO implement

    def _normalisation_check(self):
        def _normalisation_check(array: np.ndarray):
            if array is not None:
                _max, _min = np.amax(array), np.amin(array)
                if _max > 1:
                    raise ValueError(f"Critical: dataset contains maximum value {_max} which exceeds [-1, 1].")
                if _min < -1:
                    raise ValueError(f"Critical: dataset contains minimum value {_min} which exceeds [-1, 1].")

        _normalisation_check(self._trainset_data)
        _normalisation_check(self._validationset_data)
        _normalisation_check(self._testset_enrollment_data)
        _normalisation_check(self._testset_matching_data)

    def _gather_labels(self) -> np.ndarray:
        """Concatenates all labels."""
        labels = [self._trainset_labels]
        if self._validationset_labels is not None:
            labels.append(self._validationset_labels)
        if self._testset_enrollment_labels is not None:
            labels.append(self._testset_enrollment_labels)
        if self._testset_matching_labels is not None:
            labels.append(self._testset_matching_labels)
        return np.concatenate(labels)

    def _get_unique_labels(self):
        return np.unique(self._gather_labels())

    def get_unique_classes(self):
        return np.unique(self._gather_labels())

    def get_unique_enrollment_classes(self):
        return np.unique(self._testset_enrollment_labels)

    def get_unique_matching_classes(self):
        return np.unique(self._testset_matching_labels)

    def _get_enrollment_data_filter(self, filter):
        """
        Return enrollment data filtered by the given filter interacting with enrollment labels.

        The filter can be one of the following:
          - callable: predicate taking a label and returning True/False
          - scalar label value: selects items equal to this label
          - iterable of label values: selects items whose label is in the iterable
        """
        if self._testset_enrollment_data is None or self._testset_enrollment_labels is None:
            raise ValueError("Enrollment data or labels are not available to filter.")

        labels = self._testset_enrollment_labels

        # Build boolean mask according to provided filter
        mask: np.ndarray
        if callable(filter):
            # Apply predicate per label
            mask = np.array([bool(filter(lbl)) for lbl in labels])
        else:
            try:
                # If filter is iterable (e.g., list, set), treat as membership filter
                values = set(list(filter))  # may raise TypeError if not iterable
                mask = np.array([lbl in values for lbl in labels])
            except TypeError:
                # Treat as scalar equality
                mask = (labels == filter)

        return self._testset_enrollment_data[mask]

    def _get_matching_data_filter(self, filter):
        """
        Return matching data filtered by the given filter interacting with matching labels.

        The filter can be one of the following:
          - callable: predicate taking a label and returning True/False
          - scalar label value: selects items equal to this label
          - iterable of label values: selects items whose label is in the iterable
        """
        if self._testset_matching_data is None or self._testset_matching_labels is None:
            raise ValueError("Matching data or labels are not available to filter.")

        labels = self._testset_matching_labels

        # Build boolean mask according to provided filter
        if callable(filter):
            mask = np.array([bool(filter(lbl)) for lbl in labels])
        else:
            try:
                values = set(list(filter))
                mask = np.array([lbl in values for lbl in labels])
            except TypeError:
                mask = (labels == filter)

        return self._testset_matching_data[mask]

    def get_dl_training_shape(self):
        return self._testset_enrollment_data.shape[1:]

    def get_enrollment_labels(self, ohc=False):
        if ohc:
            return self.ohc_encoder.transform(self._testset_enrollment_labels)
        else:
            return self._testset_enrollment_labels

    def get_matching_labels(self, ohc=False):
        if ohc:
            return self.ohc_encoder.transform(self._testset_matching_labels)
        else:
            return self._testset_matching_labels

    def get_enrollment_data(self, filter=None):
        if filter is not None:
            return self._get_enrollment_data_filter(filter)
        else:
            return self._testset_enrollment_data

    def get_trainset_data(self):
        return self._trainset_data

    def get_trainset_labels(self):
        return self._trainset_labels

    def get_matching_data(self, filter=None):
        if filter is not None:
            return self._get_matching_data_filter(filter)
        else:
            return self._testset_matching_data

    def ohc_labels_to_classes(self, ohc_labels):
        return self.ohc_encoder.inverse_transform(ohc_labels)


class WindowslicedDataset(Dataset):
    def __init__(self,
                 window_slicing_stride: int,
                 window_slicing_size: int,
                 trainset_data: Union[list, np.ndarray],
                 trainset_labels: Union[list, np.ndarray],
                 validationset_data: Union[list, np.ndarray] = None,
                 validationset_labels: Union[list, np.ndarray] = None,
                 testset_enrollment_data: Union[list, np.ndarray] = None,
                 testset_enrollment_labels: Union[list, np.ndarray] = None,
                 testset_matching_data: Union[list, np.ndarray] = None,
                 testset_matching_labels: Union[list, np.ndarray] = None,
                 normalization_check: bool = False,
                 shuffle_window_slices: bool = False,
                 voting_function = None,
                 ):
        if voting_function is None:
            self.voting_function = simple_majority_voting

        self._trainset_data: np.ndarray = trainset_data  # type hints here are just for IDE support
        self._trainset_labels: np.ndarray = trainset_labels
        self._validationset_data: np.ndarray = validationset_data
        self._validationset_labels: np.ndarray = validationset_labels
        self._testset_enrollment_data: np.ndarray = testset_enrollment_data
        self._testset_enrollment_labels: np.ndarray = testset_enrollment_labels
        self._testset_matching_data: np.ndarray = testset_matching_data
        self._testset_matching_labels: np.ndarray = testset_matching_labels
        self._setup_missing_sets()  # applying the replacement rules for missing sets in this way seems a little


        def _len_or_none(data: np.ndarray):
            return None if data is None else len(data)

        trainset_origin_len = _len_or_none(self._trainset_data)
        validationset_origin_len = _len_or_none(self._validationset_data)
        testset_enrollment_origin_len = _len_or_none(self._testset_enrollment_data)
        testset_matching_origin_len = _len_or_none(self._testset_matching_data)

        _trainset_data, _trainset_labels, self.trainset_slicedsample_origin = self._transform_for_windowslicing(
            self._trainset_data, self._trainset_labels, window_slicing_stride, window_slicing_size, shuffle_window_slices)

        _validationset_data, _validationset_labels, self.validationset_slicedsample_origin = \
            self._transform_for_windowslicing(self._validationset_data, self._validationset_labels, window_slicing_stride,
                                              window_slicing_size, shuffle_window_slices)

        _testset_enrollment_data, _testset_enrollment_labels, self.testset_enrollment_slicedsample_origin = \
            self._transform_for_windowslicing(self._testset_enrollment_data, self._testset_enrollment_labels,
                                              window_slicing_stride, window_slicing_size, shuffle_window_slices)

        _testset_matching_data, _testset_matching_labels, self.testset_matching_slicedsample_origin = \
            self._transform_for_windowslicing(self._testset_matching_data, self._testset_matching_labels,
                                              window_slicing_stride, window_slicing_size, shuffle_window_slices)

        self._create_origin_masks(trainset_origin_len, validationset_origin_len, testset_enrollment_origin_len,
                                  testset_matching_origin_len)

        super().__init__(_trainset_data, _trainset_labels, _validationset_data, _validationset_labels,
                         _testset_enrollment_data, _testset_enrollment_labels, _testset_matching_data,
                         _testset_matching_labels, normalization_check)

    def _create_origin_masks(self, trainset_origin_len, validationset_origin_len, testset_enrollment_origin_len,
                             testset_matching_origin_len):
        def _create_origin_mask(origin_len, slicedsample_origin):
            if slicedsample_origin is None:
                return None
            else:
                result_origin_mask = []
                for origin_id in range(origin_len):
                    result_origin_mask.append(slicedsample_origin == origin_id)
                return np.array(result_origin_mask, dtype=bool)

        self.trainset_sliceorigin_mask = _create_origin_mask(trainset_origin_len, self.trainset_slicedsample_origin)
        self.validationset_sliceorigin_mask = _create_origin_mask(validationset_origin_len,
                                                                  self.validationset_slicedsample_origin)
        self.testset_enrollment_sliceorigin_mask = _create_origin_mask(testset_enrollment_origin_len,
                                                                       self.testset_enrollment_slicedsample_origin)
        self.testset_matching_sliceorigin_mask = _create_origin_mask(testset_matching_origin_len,
                                                                     self.testset_matching_slicedsample_origin)

    @staticmethod
    def _transform_for_windowslicing(data: Union[list, np.ndarray],
                                     labels: Union[list, np.ndarray],
                                     window_slicing_stride: int,
                                     window_slicing_size: int,
                                     shuffle_window_slices: bool):
        if data is None:
            return None, None, None

        if not len(data) == len(labels):
            raise ValueError("Critical: data and labels differ in size.")

        ret_sliced_data, ret_sliced_labels, ret_sample_origin_ids = [], [], []

        for origin_id, (sample, label) in enumerate(zip(data, labels)):
            array_slices, label_slices, sample_origin_ids = [], [], []

            lost_samples_cnt, added_cnt = 0, 0

            for step_idx in range(0, len(sample) - window_slicing_size + 1, window_slicing_stride):
                arr = np.array(sample[step_idx:step_idx + window_slicing_size],
                               dtype=float)  # this fails if the resulting array would be ragged

                if not arr.any():
                    pass  # do not append if slice consists only of zero
                else:
                    array_slices.append(arr)
                    label_slices.append(label)
                    sample_origin_ids.append(origin_id)

            for i in range(len(array_slices)):
                if len(array_slices[i]) == window_slicing_size:
                    ret_sliced_data.append(array_slices[i])
                    ret_sliced_labels.append(label_slices[i])
                    ret_sample_origin_ids.append(sample_origin_ids[i])
                    added_cnt += 1
                else:
                    lost_samples_cnt += 1

            if lost_samples_cnt > 0:
                print(f'Warning: during slicing  {lost_samples_cnt} malformed slices needed to be discarded. '
                      f'Check slicing parameters.', file=sys.stderr)

        try:
            ret_sliced_data = np.array(ret_sliced_data, dtype=float)
        except ValueError:
            print("Critical: slicing the data produced slices differing in size.", file=sys.stderr)
            raise

        ret_sliced_labels = np.array(ret_sliced_labels, dtype=str)
        ret_sample_origin_ids = np.array(ret_sample_origin_ids, dtype=int)

        print(f'Info: sliced window data takes up {round(ret_sliced_data.nbytes / 1024 / 1024, 2)}mb.')

        if shuffle_window_slices:
            return _unison_shuffle_np_arrays_3(ret_sliced_data, ret_sliced_labels, ret_sample_origin_ids)
        else:
            return ret_sliced_data, ret_sliced_labels, ret_sample_origin_ids

    def apply_voting_for_testset_matching_slices(self, predictions_slices: np.ndarray,
                                                 ground_truth_slices: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        ground_truth = np.zeros((len(self.testset_matching_sliceorigin_mask),), dtype=ground_truth_slices.dtype)
        predictions = np.zeros((len(self.testset_matching_sliceorigin_mask),), dtype=predictions_slices.dtype)

        for sample_id, sample_mask in enumerate(self.testset_matching_sliceorigin_mask):
            if np.any(sample_mask):
                if predictions_slices.ndim == 2:  # ohc tabular predictions
                    predictions[sample_id] = self.voting_function(np.argmax(predictions_slices[sample_mask], axis=1))
                    ground_truth[sample_id] = np.argmax(ground_truth_slices[sample_mask], axis=1)[0]
                else:
                    predictions[sample_id] = self.voting_function(predictions_slices[sample_mask])
                    ground_truth[sample_id] = ground_truth_slices[sample_mask][0]  # assuming consistency within slices
            else:
                predictions[sample_id] = None  # Default value for predictions when no slices match
                ground_truth[sample_id] = None  # Default value for ground truth when no slices match

        return predictions, ground_truth
