# Import external functionalities
import numpy as np
import random
import pandas as pd
from scipy.stats import random_correlation
from scipy.linalg import cholesky
# Import internal functionalities
from .utils.generator_utils import dist_constant, dist_periodical, dist_gauss, dist_uniform, \
    helper_get_feature, helper_get_sensor, linear_trans_func
from .concept.regression_concept import RegressionConceptModel


class DataGenerator:
    """
    Generator used to make a non-stationary regression problem.

    Initialization defines behaviour of the non-stationary concept and sampling defines behaviour of the input data.
    """

    def __init__(self, number_of_data_points: int = 100, number_of_features: int = 2, number_of_outputs: int = 1,
                 feature_min: float = -100.0, feature_max: float = 100.0,
                 number_of_models: int = 5, noise_var: float = 0, rand_seed: int = -1,
                 concept_drifts: int = 0, data_drifts: int = 0, concept_drift_class=None,
                 data_drift_class=None, transition_func=linear_trans_func,
                 continous_time: bool = True, max_time_sparsity: int = 10, drift_blocking_mode: bool = False,
                 max_severity=1, min_severity=0):
        """
        Initialization of the Generator, defines behaviour of the non stationary concept.
        Default values are one possible useage of the generator. Reassigning is highly encouraged to get data streams
        that can be used for testing.


        Parameters
        ----------
        :param number_of_data_points: int, default = 100
            Number of instances that will be included in the dataset generated by get_data().
            As drifts should occur in a timeframe where data is sampled, the timeframe has to be defined to define when
            drift may occur.

        :param number_of_features: int, default = 2
            Number of input features generated and used in the data stream.

        :param number_of_outputs: int, default = 1
            Number of output targets produced by the artificial non-stationary concept, based on the generated input
            features.

        :param feature_min: float, default = -100.0
            The initial set of generated input features will not fall below this value. Note that data drifts and noise
            may lead to feature values that to not comply to this limitation.

        :param feature_max: float, default = 100.0
            The initial set of generated input features will not surpass this value. Note that data drifts and noise
            may lead to feature values that to not comply to this limitation.

        :param number_of_models: int, default = 5
            Number of random models used to build the artificial concept. More models enable more complex concepts.
            Note that an increasing number of models does increase the maximal complexity but also may reduce variance
            in the generated concept, as the random models are linearly combined to form the concept.

        :param noise_var: float, default = 0.0
            Variance of the normal distribution used to sample the scaling factor for noise in features and targets.
            As input and output values may have very different value ranges we opted for a noisy sclaing factor instead
            of an additive noise.

        :param rand_seed: int, default = -1
            Seed used for random choices like sampling. If set to -1 no specific seed will be set, leading to random
            results in every execution

        :param concept_drifts: int, default=0
            Number of concept drifts introduced in the generated dataset of get_data().

        :param data_drifts: int, default=0
            Number of data drifts introduced in the generated dataset of get_data().

        :param concept_drift_class: str, default=None
            What type ('sudden','gradual','incremental',...) the generated concept drifts should have.
            If set to None, the type of will be decided for each drift at random.

        :param data_drift_class: str, default=None
            What type ('sudden','gradual','incremental','faulty_sensor',...) the generated data drifts should have.
            If set to None, the type of will be decided for each drift at random.

        :param transition_func: callable, default=linear_transition_func
            Hoe two concepts are to be blended. Used for incremental and gradual features. Custom transition functions
            may be provided, but must implement the api of the following exemplary call:
                transition_func(current_time_stamp:int, shift_centre:int, shift_radius:float)->float
            Note that the transition function is only called if there is a transition ongoing, i.e. :
                 shift_centre - shift_radius < current_time_stamp <= shift_centre + shift_radius

            See the exemplary linear_transition function used as default:
                def linear_trans_func(current_time, shift_centre, shift_radius):
                    return (current_time - (shift_centre - shift_radius)) / (shift_radius * 2)"

        :param continous_time: bool, default=True
            Whether to use complete timeseries or timeseries with missing timestamps.
            If set to False behaviour of missing timestamps can be set by max_time_sparsity.

        :param max_time_sparsity: int, default=10
            Only used when continous_time=False.
            Sets te maximal number of timestamps skipped after a timestamp is evaluated.

        :param drift_blocking_mode: bool, default=False
            Whether drifts may overlap or be confined to equally sized timeframes in the generated timeseries.
            If set to True concept drifts will not overlap with other concept drifts, same goes for data drifts.
            Note that concept and data drifts are still able to overlap and influence each other.

        :param max_severity: float, default=1
            Internal scaling factor on how drastic drifts are allowed to be.
            Can be used to control relative intensity of drifts.
            If max_severity=min_severity , then all drifts will have comparable impacts on data distributions and
            concept.

        :param min_severity: float, default=0
            Internal scaling factor in how drastic drifts have to be.
        """
        if rand_seed != -1:
            random.seed(rand_seed)
            np.random.seed(rand_seed)
        if continous_time:
            time_stamp = np.arange(0, number_of_data_points)
        else:
            temp_time_stamp = []
            current_time_stamp = 0
            for i in range(number_of_data_points):
                temp_time_stamp.append(current_time_stamp)
                current_time_stamp += random.randint(1, max_time_sparsity)
            time_stamp = np.array(temp_time_stamp)

        self.continous_time = continous_time
        self.number_of_data_points = number_of_data_points
        self.time_stamp = time_stamp
        self.transition_func = transition_func

        if drift_blocking_mode and concept_drifts != 0:
            step = (time_stamp[-1] - time_stamp[0]) / concept_drifts
            shift_stamps = np.array([0.5 * step + step * i for i in range(concept_drifts)])
        else:
            shift_stamps = np.random.rand(concept_drifts) * (time_stamp[-1] - time_stamp[0])
        if rand_seed != -1:
            random.seed(rand_seed)
            np.random.seed(rand_seed)
        self.concept = RegressionConceptModel(number_of_models, number_of_features, number_of_outputs,
                                              shift_stamps=shift_stamps, drift_blocking_mode=drift_blocking_mode,
                                              drift_spacing=(time_stamp[-1] - time_stamp[
                                                  0]) / concept_drifts if concept_drifts > 0 else None,
                                              max_severity=max_severity, min_severity=min_severity,
                                              drift_class=concept_drift_class)

        if drift_blocking_mode and data_drifts != 0:
            step = (time_stamp[-1] - time_stamp[0]) / data_drifts
            self.data_shifts = np.array([0.5 * step + step * i for i in range(data_drifts)])
        else:
            self.data_shifts = np.random.rand(data_drifts) * (time_stamp[-1] - time_stamp[0])
        if rand_seed != -1:
            random.seed(rand_seed)
            np.random.seed(rand_seed)
        self.data_shift_info = []
        for i in range(data_drifts):

            if drift_blocking_mode:
                self.data_shift_info.append(
                    [random.randint(0, number_of_features - 1),
                     random.random() / 2 * (time_stamp[-1] - time_stamp[0]) / data_drifts,
                     np.minimum(max_severity, np.maximum(min_severity, np.random.rand(10))),
                     random.choice(('sudden', 'gradual', 'incremental', 'reoccuring_concept', 'faulty_sensor'))])
            else:
                self.data_shift_info.append(
                    [random.randint(0, number_of_features - 1), random.random() * (np.max(self.data_shifts) / 2),
                     np.minimum(max_severity, np.maximum(min_severity, np.random.rand(10))),
                     random.choice(('sudden', 'gradual', 'incremental', 'reoccuring_concept', 'faulty_sensor'))])
            if data_drift_class is not None:
                temp_list = list(self.data_shift_info[-1])
                temp_list[-1] = data_drift_class
                self.data_shift_info[-1] = temp_list

        self.min = feature_min
        self.max = feature_max
        self.n_models = number_of_models
        self.n_features = number_of_features
        self.n_outputs = number_of_outputs
        self.noise_var = noise_var

    def get_data(self, n_uniform_feat=0, n_gauss_feat=0, n_constant_feat=0, n_periodical_feat=0,
                 correlated_features=0, n_sensor_features=0):
        """
        Use the initialized Generator to sample datapoints in the non-stationary environment.
        Data behaviour can be reassigned with each calling on the function.

        Data behaviour is defined by used and parameterized data distributions.
        The input arguments can be used to set the minimal number of features that are sampled from a particular kind
        of distribution. You may only assign as many features as where defined in generator initialization.

        If there are fewer given assignments then features, the remaining features will be sorted into categories by
        random.

        Correlated and sensor features are special as they are generated based on already generated features.
        Correlated features are restructurings of features and therefore do not increase the number of assigned
        features.
        Sensor features are features that are dependent on other features.
        To produce these features there have to be at least two features in the basic categories.
        Correlated and sensor features will NOT be generated at random.

        :param n_uniform_feat: int, default=0
            Number of features that follow a uniform distribution

        :param n_gauss_feat: int, default=0
            Number of features that follow a gaussian distribution

        :param n_constant_feat: int, default=0
            Number of features that are constant

        :param n_periodical_feat: int, default=0
            Number of features that follow a periodical distribution.
            Periodical distributions are for example cosine of a time dependent feature that resets after a certain
            number of days.

        :param correlated_features: int, default=0
            Number of features that are restructured to be correlated.

        :param n_sensor_features:  int, default=0
            Number of features generated based on other features.

        :return: (Features, Targets)
        :rtype: (pandas.DataFrame of shape (number_of_data_points,number_of_features),
                pandas.DataFrame of shape (number_of_data_points,number_of_outputs))

        """
        if n_uniform_feat + n_gauss_feat + n_constant_feat + n_periodical_feat + n_sensor_features > self.n_features:
            n_set_feat = n_uniform_feat + n_gauss_feat + n_constant_feat + n_periodical_feat + n_sensor_features
            raise TypeError(
                f"Too many features set staticaly: {n_set_feat} of {self.n_features}")
        if self.n_features - 2 < n_sensor_features and n_sensor_features > 0:
            raise TypeError("Too many features set to be sensors. At least two features have to be no sensors.")

        indices_uniform = []
        indices_gauss = []
        indices_constant = []
        indices_periodical = []
        indices_sensors = []
        indices = np.arange(0, self.n_features)
        indices = list(indices)
        for i in range(n_sensor_features):
            indices_sensors.append(indices.pop())
        indices = np.asarray(indices)
        np.random.shuffle(indices)
        indices = list(indices)

        # set feature behaviour
        for i in range(n_uniform_feat):
            indices_uniform.append(indices.pop())
        for i in range(n_gauss_feat):
            indices_gauss.append(indices.pop())
        for i in range(n_constant_feat):
            indices_constant.append(indices.pop())
        for i in range(n_periodical_feat):
            indices_periodical.append(indices.pop())

        for i in range(len(indices)):
            random.choice((indices_uniform, indices_gauss, indices_constant)).append(indices.pop())

        # compute feature behaviour
        samples = np.zeros((self.number_of_data_points, self.n_features))
        samples = np.transpose(samples)
        list_all_dist = [dist_uniform, dist_gauss, dist_constant, dist_periodical]

        for i in range(samples.shape[0]):
            if i in indices_uniform:
                samples[i] = helper_get_feature(i, self.min, self.max, dist_uniform, list_all_dist, self.data_shifts,
                                                self.data_shift_info, self.time_stamp, self.transition_func)
            if i in indices_gauss:
                samples[i] = helper_get_feature(i, self.min, self.max, dist_gauss, list_all_dist, self.data_shifts,
                                                self.data_shift_info, self.time_stamp, self.transition_func)
            if i in indices_constant:
                samples[i] = helper_get_feature(i, self.min, self.max, dist_constant, list_all_dist, self.data_shifts,
                                                self.data_shift_info, self.time_stamp, self.transition_func)
            if i in indices_periodical:
                samples[i] = helper_get_feature(i, self.min, self.max, dist_periodical, list_all_dist, self.data_shifts,
                                                self.data_shift_info, self.time_stamp, self.transition_func)
            if i in indices_sensors:
                sensor_input = np.arange(0, self.n_features - n_sensor_features)
                np.random.shuffle(sensor_input)
                sensor_input = list(sensor_input)
                samples[i] = helper_get_sensor(i, samples[sensor_input.pop()], samples[sensor_input.pop()],
                                               self.data_shifts, self.data_shift_info, self.time_stamp,
                                               self.transition_func).flatten()
        if correlated_features != 0:
            # recompute to correlate
            eigs = (np.random.rand(correlated_features)) * 20
            eigs *= correlated_features / np.sum(eigs)

            C = random_correlation.rvs(eigs)

            U = cholesky(C)  # Cholesky decomposition
            R = samples[:correlated_features]  # Three uncorrelated sequences
            Rc = R.T @ U  # Array of correlated random sequences
            Rc = Rc.T
            for i in range(correlated_features):
                samples[i] = Rc[i]

        samples = np.transpose(samples)

        # labels and noise
        labels = self.concept.label(samples, self.time_stamp)
        samples = np.transpose(samples)
        for i in range(samples.shape[0]):
            if i not in indices_constant:
                samples[i] = samples[i] * (np.random.normal(0, self.noise_var, samples[i].shape) + 1)
        samples = np.transpose(samples)
        labels = labels * (np.random.normal(0, self.noise_var, labels.shape) + 1)
        stamp_frame = pd.DataFrame(self.time_stamp, columns=['time_stamp'])
        sample_frame = pd.DataFrame(samples, columns=[f'feat_{i}' for i in range(samples.shape[1])])
        label_frame = pd.DataFrame(labels, columns=[f'label_{i}' for i in range(labels.shape[1])])
        return pd.concat([stamp_frame, sample_frame], axis=1), pd.concat([stamp_frame, label_frame], axis=1)

    def get_concept(self):
        """
        Provides the drit induced concept function that is used to map the generated inputs to the targets.
        Arbitrary inputs may be labeled with get_concept().label(input,(int)timestamp)

        :return: Concept used in the drifted datastream
        :rtype: RegressionConceptModel()
        """
        return self.concept

    def get_shift_information(self):
        """
        Provides two pandas dataframes, that describe the concept and data drifts.
        For the concept drifts (time,duration,weight shift,class) is provided.
        For data drifts (time,affected feature,duration,weight shift,class) is provided

        :return: Drift Information DataFrames
        :rtype: (pandas.DataFrame of shape (concept_drifts,5),
                pandas.DataFrame of shape (data_drifts,6))
        """
        concept = self.concept.get_shift_information()
        centres = pd.DataFrame(self.data_shifts, columns=['time_stamp(centre)'])
        info = pd.DataFrame(self.data_shift_info,
                            columns=['affected_feature', 'radius', 'shift of distribution parameters', 'class'])
        return concept, pd.concat([centres, info], axis=1)
