# Source code for ml4sec.scale

    from sklearn.decomposition import PCA

    import matplotlib.pyplot as plt
    import numpy             as np
    import pandas            as pd

    ################################################################################
    #                                   Scaling                                    #
    ################################################################################

    def scale(self, matrix, minimum=None, maximum=None):
        """Scale features in matrix.

            Parameters
            ----------
            matrix : np.array of shape=(n_samples, n_features)
                Matrix to scale.

            minimum : array-like of shape=(n_features,), optional
                If given, should contain the minimum values for each feature to
                scale with.

            maximum : array-like of shape=(n_features,), optional
                If given, should contain the maximum values for each feature to
                scale with.

            Returns
            -------
            scaled_matrix : array-like of shape=(n_samples, n_features)
                Scaled matrix.

            minimum : array-like of shape=(n_features,)
                The minimum values for each feature used for scaling.

            maximum : array-like of shape=(n_features,)
                The maximum values for each feature used for scaling.
            """
        raise NotImplementedError(
            "You will need to implement the 'scale()' function yourself."
        )

    ################################################################################
    #                                Scale as numpy                                #
    ################################################################################

    def scale_numpy(self, matrix, minimum=None, maximum=None):
        """Scale features in matrix to numpy arrays.

            Parameters
            ----------
            matrix : np.array of shape=(n_samples, n_features)
                Matrix to scale.

            minimum : array-like of shape=(n_features,), optional
                If given, should contain the minimum values for each feature to
                scale with.

            maximum : array-like of shape=(n_features,), optional
                If given, should contain the maximum values for each feature to
                scale with.

            Returns
            -------
            scaled_matrix : np.array of shape=(n_samples, n_features)
                Scaled matrix.

            minimum : np.array of shape=(n_features,)
                The minimum values for each feature used for scaling.

            maximum : np.array of shape=(n_features,)
                The maximum values for each feature used for scaling.
            """
        # Perform scaling
        scaled_matrix, minimum, maximum = self.scale(matrix, minimum, maximum)

        # Transform to numpy arrays
        scaled_matrix = np.asarray(scaled_matrix)
        minimum       = np.asarray(minimum)
        maximum       = np.asarray(maximum)

        # Return result
        return scaled_matrix, minimum, maximum

    ################################################################################
    #                          Perform checks on scaling                           #
    ################################################################################

    def test_scale(self, verbose=True):
        """Method that runs some tests on data scaling.

            Parameters
            ----------
            verbose : boolean, default=True
                If True, print test cases.
            """

        ########################################################################
        #                             Test case 1                              #
        ########################################################################

        # Input
        matrix = np.asarray([
            [0, 1],
            [2, 3],
            [4, 5],
            [6, 7],
            [8, 9],
        ])
        minimum = None
        maximum = None

        # Target output
        target_scaled = np.asarray([
            [0   , 0],
            [0.25, 0.25],
            [0.5 , 0.5],
            [0.75, 0.75],
            [1   , 1],
        ])
        target_minimum = np.asarray([0, 1])
        target_maximum = np.asarray([8, 9])

        # Run single test case
        self.test_scale_single(
            scaled_input   = matrix,
            minimum_input  = minimum,
            maximum_input  = maximum,
            scaled_target  = target_scaled,
            minimum_target = target_minimum,
            maximum_target = target_maximum,
            verbose        = verbose,
        )

        # ########################################################################
        # #                             Test case 2                              #
        # ########################################################################

        # Input
        matrix = np.asarray([
            [0, 1, 10],
            [2, 3, 11],
            [4, 5, 12],
            [6, 7, 13],
            [8, 9, 14],
        ])
        minimum = np.asarray([4, -20, 0])
        maximum = np.asarray([6,  20, 1])

        # Target output
        target_scaled = np.asarray([
            [-2, 0.525, 10],
            [-1, 0.575, 11],
            [0 , 0.625, 12],
            [1 , 0.675, 13],
            [2 , 0.725, 14],
        ])
        target_minimum = np.asarray([4, -20, 0])
        target_maximum = np.asarray([6,  20, 1])

        # Run single test case
        self.test_scale_single(
            scaled_input   = matrix,
            minimum_input  = minimum,
            maximum_input  = maximum,
            scaled_target  = target_scaled,
            minimum_target = target_minimum,
            maximum_target = target_maximum,
            verbose        = verbose,
        )

        ########################################################################
        #                        Print success message                         #
        ########################################################################

        # Print success message
        print("All tests passed.")

    def test_scale_single(self,
            scaled_input, minimum_input, maximum_input,
            scaled_target, minimum_target, maximum_target,
            verbose = False,
        ):
        """Perform a single test case of input and target.

            Parameters
            ----------
            scaled_input : np.array of shape=(n_samples, n_features)
                matrix input given to scale function.

            minimum_input : np.array of shape=(n_features), optional
                minimum input given to scale function.

            maximum_input : np.array of shape=(n_features), optional
                maximum input given to scale function.

            scaled_target : np.array of shape=(n_samples, n_features)
                Target scaled matrix if scale is implemented correctly.

            minimum_target : np.array of shape=(n_features)
                Target minimum of features if scale is implemented correctly.

            maximum_target : np.array of shape=(n_features)
                Target maximum of features if scale is implemented correctly.

            verbose : boolean, default=False
                If True, print information on test case.
            """

        # If verbose, print test case
        if verbose:
            print("Test case:")
            print("matrix  = \n{}".format(scaled_input))
            print("minimum = {}"  .format(minimum_input))
            print("maximum = {}\n".format(maximum_input))

        # Call scale function
        output = self.scale_numpy(
            matrix  = scaled_input,
            minimum = minimum_input,
            maximum = maximum_input,
        )

        # Perform checks
        self.assert_scale(
            scaled_target, minimum_target, maximum_target,
            *output,
        )

    def assert_scale(self,
            scaled_target     , minimum_target     , maximum_target     ,
            scaled_output=None, minimum_output=None, maximum_output=None,
        ):
        """Performs checks comparing a computed output and target output.

            Parameters
            ----------
            scaled_target : array-like of shape=(n_samples, n_features)
                Target scaled matrix if scale is implemented correctly.

            minimum_target : array-like of shape=(n_features)
                Target minimum of features if scale is implemented correctly.

            maximum_target : array-like of shape=(n_features)
                Target maximum of features if scale is implemented correctly.

            scaled_output : array-like of shape=(n_samples, n_features)
                Scaled output by running it through the scale implementation.

            minimum_output : array-like of shape=(n_features)
                Minimum of features computed by scale implementation.

            maximum_output : array-like of shape=(n_features)
                Maximum of features computed by scale implementation.
            """

        ########################################################################
        #                       Check if output is given                       #
        ########################################################################

        """Perform check on scale implementation to check for common bugs."""
        # Check if scaled matrix is returned
        if scaled_output is None:
            raise ValueError(
                "Your method did not return a scaled matrix. Please return:\n"
                " - the scaled matrix\n"
                " - minimum value for each column\n"
                " - maximum value for each column\n"
            )

        # Check if minimum is returned
        if minimum_output is None:
            raise ValueError(
                "Your method did not return a minimum value. Please return:\n"
                " - the scaled matrix\n"
                " - minimum value for each column\n"
                " - maximum value for each column\n"
            )

        # Check if maximum is returned
        if minimum_output is None:
            raise ValueError(
                "Your method did not return a maximum value. Please return:\n"
                " - the scaled matrix\n"
                " - minimum value for each column\n"
                " - maximum value for each column\n"
            )

        ########################################################################
        #                 Check if output is given as an array                 #
        ########################################################################

        # Transform arrays to numpy
        try:
            scaled_output  = np.asarray(scaled_output )
            minimum_output = np.asarray(minimum_output)
            maximum_output = np.asarray(maximum_output)
        except Exception as e:
            raise ValueError("Your could not be cast to array: '{}'".format(e))

        # Transform arrays to numpy
        try:
            scaled_target  = np.asarray(scaled_target )
            minimum_target = np.asarray(minimum_target)
            maximum_target = np.asarray(maximum_target)
        except Exception as e:
            raise ValueError(
                "Target array could not be cast to array: '{}'. If you called the "
                "'assert_scale' method yourself, please check if your input is "
                "correct. If you get this message when running test_scale, please "
                "contact t.s.vanede@utwente.nl".format(e)
            )

        ########################################################################
        #                  Check if output shapes are correct                  #
        ########################################################################

        # Check if minimum is of correct shape
        if minimum_output.shape[0] != scaled_target.shape[1]:
            raise ValueError(
                "The size of your computed minimum ({}) did not match the expected "
                "size ({}). Minimum should be computed based on the features "
                "(columns). Please check if you did not accidentally compute the "
                "minimum over the samples (rows) instead.".format(
                minimum_output.shape[0], scaled_target.shape[1]
            ))

        # Check if minimum is of correct shape
        if minimum_output.shape != minimum_target.shape:
            raise ValueError(
                "The size of your computed minimum ({}) did not match the expected "
                "size ({})".format(minimum_output.shape, minimum_target.shape)
            )

        # Check if maximum is of correct shape
        if maximum_output.shape[0] != scaled_target.shape[1]:
            raise ValueError(
                "The size of your computed maximum ({}) did not match the expected "
                "size ({}). Maximum should be computed based on the features "
                "(columns). Please check if you did not accidentally compute the "
                "maximum over the samples (rows) instead.".format(
                maximum_output.shape[0], scaled_target.shape[1]
            ))

        # Check if maximum is of correct shape
        if maximum_output.shape != maximum_target.shape:
            raise ValueError(
                "The size of your computed maximum ({}) did not match the expected "
                "size ({})".format(maximum_output.shape, maximum_target.shape)
            )

        # Check if scaled matrix is of correct shape
        if scaled_output.shape != scaled_target.shape:
            raise ValueError(
                "The size of your computed scaled matrix ({}) did not match the "
                "expected size ({})".format(scaled_output.shape,
                scaled_target.shape
            ))

        ########################################################################
        #                  Check if output values are correct                  #
        ########################################################################

        # Check if computed minimum is correct
        if not np.all(minimum_output == minimum_target):
            raise ValueError(
                "Minimum was incorrectly computed. Expected minimum to be:\n{}\n, "
                "but minimum was:\n{}\n".format(minimum_target, minimum_output)
            )

        # Check if computed maximum is correct
        if not np.all(maximum_output == maximum_target):
            raise ValueError(
                "Maximum was incorrectly computed. Expected maximum to be:\n{}\n, "
                "but maximum was:\n{}\n".format(maximum_target, maximum_output)
            )

        # Check if computed scaled matrix is correct
        if not np.all(scaled_output == scaled_target):
            raise ValueError(
                "Scaled matrix was incorrectly computed. Expected scaled matrix to "
                "be:\n{}\n, but scaled matrix was:\n{}\n".format(scaled_output,
                scaled_target
            ))

    ################################################################################
    #                                Visualisation                                 #
    ################################################################################

    def show_scaled(self):
        """Shows the scaled feature matrix for benign flows.

            1. Extracts the feature matrix.
            2. Scales the data.
            3. Turns it into a pd.DataFrame.

            """
        return pd.DataFrame(self.scale_numpy(self.feature_matrix(self.flows_benign))[0])

    def plot_scaled(self):
        """Plot scaled vs unscaled.

            1. Computes feature matrix for both benign.csv and unknown.csv.
            2. Applies scaling to both benign and unknown feature matrix.
            3. Plots the differences.

            """
        # Get unscaled matrices
        X_unscaled_benign  = self.feature_matrix(self.flows_benign )
        X_unscaled_unknown = self.feature_matrix(self.flows_unknown)

        # Get scaled matrices
        X_scaled_benign , min_, max_ = self.scale_numpy(X_unscaled_benign)
        X_scaled_unknown, _, _ = self.scale_numpy(X_unscaled_unknown, min_, max_)

        # Create 2x2 grid
        fig, axs = plt.subplots(2, 2, figsize=(12, 8))

        # Create plots
        axs[0, 0].set_title("Unscaled data", fontsize=20)
        axs[0, 0].set_ylabel("Benign data", fontsize=20)
        self.plot_data(axs[0, 0], X_unscaled_benign, self.labels_benign)
        axs[0, 1].set_title("Scaled data", fontsize=20)
        self.plot_data(axs[0, 1], X_scaled_benign  , self.labels_benign)
        axs[0, 1].legend(loc='center left', bbox_to_anchor=(1, 0.5))
        axs[1, 0].set_ylabel("Unknown data", fontsize=20)
        self.plot_data(axs[1, 0], X_unscaled_unknown, self.labels_unknown)
        self.plot_data(axs[1, 1], X_scaled_unknown  , self.labels_unknown)
        axs[1, 1].legend(loc='center left', bbox_to_anchor=(1, 0.5))

        # Show plot
        plt.show()

    def plot_data(self, ax, X, y):
        """Plots the feature vectors given by X and labelled by y.

            Parameters
            ----------
            ax : matplotlib.axis
                Axis on which to plot data.

            X : array-like of shape=(n_samples, n_features)
                Feature vectors to plot.

            y : array-like of shape=(n_features,)
                Labels corresponding to the feature vectors.
            """
        # Transform to numpy arrays
        X = np.asarray(X)
        y = np.asarray(y)

        # The PCA(n) takes all features from X and compresses them
        # into n dimensions, in our case to 2 dimensions for plotting.
        X = PCA(2).fit_transform(X)

        # Next for each label, we scatter the corresponding points
        for label in np.unique(y):
            ax.scatter(X[:, 0][y == label], X[:, 1][y == label], label=label)