import numpy as np

################################################################################
#                                    Split                                     #
################################################################################

def split(self, data, labels, apps_train, apps_test, ratio):
    """Extract features from each flow. Must be implemented by the students.

        Parameters
        ----------
        data : np.array of shape=(n_flows, n_features)
            Data for which to produce a split.

        labels : np.array of shape=(n_flows,)
            Labels corresponding to the data.

        apps_train : list()
            List of apps to use for training.

        apps_test : list()
            List of apps to use for testing.

        ratio : float
            Ratio of data for training apps to use for training.
            The rest should be used for testing.

        Returns
        -------
        X_train : array-like of shape=(n_flows_train, n_features)
            Datapoints selected for training.

        y_train : array-like of shape=(n_flows_train,)
            Corresponding labels selected for training.

        X_test : array-like of shape=(n_flows_test, n_features)
            Datapoints selected for testing.

        y_test : array-like of shape=(n_flows_test,)
            Corresponding labels selected for testing.
        """
    raise NotImplementedError(
        "You will need to implement the 'split()' function yourself."
    )


def split_numpy(self, data, labels, apps_train, apps_test, ratio):
    """Perform split function and transform output to numpy arrays.

        Parameters
        ----------
        data : np.array of shape=(n_flows, n_features)
            Data for which to produce a split.

        labels : np.array of shape=(n_flows,)
            Labels corresponding to the data.

        apps_train : list()
            List of apps to use for training.

        apps_test : list()
            List of apps to use for testing.

        ratio : float
            Ratio of data for training apps to use for training.
            The rest should be used for testing.

        Returns
        -------
        X_train : np.array of shape=(n_flows_train, n_features)
            Datapoints selected for training.

        y_train : np.array of shape=(n_flows_train,)
            Corresponding labels selected for training.

        X_test : np.array of shape=(n_flows_test, n_features)
            Datapoints selected for testing.

        y_test : np.array of shape=(n_flows_test,)
            Corresponding labels selected for testing.

        """
    # Get number of features & dtype
    features = data.shape[1]
    dtype    = data.dtype

    # Get regular split
    X_train, y_train, X_test, y_test = self.split(
        data       = data,
        labels     = labels,
        apps_train = apps_train,
        apps_test  = apps_test,
        ratio      = ratio,
    )

    # Transform output to numpy arrays
    X_train = np.asarray(X_train)
    y_train = np.asarray(y_train)
    X_test  = np.asarray(X_test )
    y_test  = np.asarray(y_test )

    # Return result
    return X_train, y_train, X_test, y_test



def get_split(self, apps_train, ratio):
    """Returns the split of 'benign.csv' file for a given input.

        Parameters
        ----------
        apps_train : list()
            Apps to use in the training set, all apps not in apps_train will be
            used only for the testing set.

        ratio : float
            Ratio of datapoints for apps in app_train to use for training.

        Returns
        -------
        X_train : array-like of shape=(n_samples_train, n_features)
            Datapoints used for training.

        y_train : array-like of shape=(n_samples_train,)
            Labels corresponding to training data.
            +1 if the corresponding app is     in the training dataset,
            -1 if the corresponding app is not in the training dataset,
            I.e., all values are +1.

        X_test : array-like of shape=(n_samples_test, n_features)
            Datapoints used for testing.

        y_test : array-like of shape=(n_samples_test,)
            Labels corresponding to testing data.
            +1 if the corresponding app is     in the training dataset,
            -1 if the corresponding app is not in the training dataset.
        """
    # Get relevant data and labels
    data, _, _ = self.scale_numpy(self.feature_matrix(self.flows_benign))
    labels     = self.labels_benign

    # Compute apps in test set
    apps_test = list(set(list(np.unique(labels))) - set(apps_train))

    # Execute split
    X_train, y_train, X_test, y_test = self.split_numpy(
        data       = data,
        labels     = labels,
        apps_train = apps_train,
        apps_test  = apps_test,
        ratio      = ratio,
    )

    # Transform labels to -1/+1 depending on whether they are inside (+1)
    # or outside (-1) of the training set
    y_train = np.ones(y_train.shape[0], dtype=int)
    y_test  = np.isin(y_test, apps_train).astype(int) * 2 - 1

    # Return result
    return X_train, y_train, X_test, y_test


################################################################################
#                                 Test methods                                 #
################################################################################

def test_split(self, verbose=False):
    """Executes several test cases for splitting the data.

        Parameters
        ----------
        verbose : boolean, default=False
            If True, print tests that are executed.

        """
    # Select data
    data, _, _ = self.scale_numpy(self.feature_matrix(self.flows_benign))
    labels     = self.labels_benign

    # Test cases
    all_apps_train = [
        ['Anti Virus', 'Chrome', 'DNS Server', 'Excel', 'Firefox', 'Git', 'IntelliJ'],
        ['Anti Virus', 'Outlook', 'DNS Server', 'Excel', 'Firefox', 'Git', 'IntelliJ', 'PowerPoint', 'Thunderbird', 'Web Server', 'Word'],
        ['IntelliJ', 'PowerPoint', 'Thunderbird', 'Web Server', 'Word'],
        ['IntelliJ', 'PowerPoint', 'Thunderbird', 'Web Server', 'Word', 'Mail Server', 'Chrome', 'Anti Virus', 'Outlook', 'DNS Server', 'Excel', 'Firefox', 'Git'],
        [],
    ]

    all_apps_test = [
        ['Mail Server', 'Outlook', 'PowerPoint', 'Thunderbird', 'Web Server', 'Word'],
        ['Mail Server', 'Chrome'],
        ['Mail Server', 'Chrome', 'Anti Virus', 'Outlook', 'DNS Server', 'Excel', 'Firefox', 'Git'],
        [],
        ['IntelliJ', 'PowerPoint', 'Thunderbird', 'Web Server', 'Word', 'Mail Server', 'Chrome', 'Anti Virus', 'Outlook', 'DNS Server', 'Excel', 'Firefox', 'Git'],
    ]

    all_ratios   = [0.5, 0.8, 0.25, 0.9, 0.0]
    all_expected = [4210, 8786, 784, 11330, 0]

    for case, (apps_train, apps_test, ratio, expected) in enumerate(
            zip(all_apps_train, all_apps_test, all_ratios, all_expected)
        ):
        # Print test case
        if verbose:
            print("Test case {} - benign.csv:".format(case + 1))
            print("apps_train = {}".format(apps_train))
            print("apps_test  = {}".format(apps_test ))
            print("ratio      = {}".format(ratio))
            print()

        # Execute split
        data_train, labels_train, data_test, labels_test = self.split_numpy(
            data       = data,
            labels     = labels,
            apps_train = apps_train,
            apps_test  = apps_test,
            ratio      = ratio,
        )

        # Test results
        self.assert_split(
            original_data          = data,
            original_labels        = labels,
            data_train             = data_train,
            labels_train           = labels_train,
            data_test              = data_test,
            labels_test            = labels_test,
            expected_samples_train = expected,
            apps_test              = apps_test,
        )

    # Confirm that all cases passed
    print("All test cases passed.")



def assert_split(self, original_data, original_labels,
                 data_train, labels_train, data_test, labels_test,
                 expected_samples_train, apps_test
    ):
    """Tests whether the split is correct.
        If an incorrect split is found, it raises an error with a suggestion for
        the students to debug their code.

        Parameters
        ----------
        original_data : np.array of shape=(n_samples, n_features)
            Original data before split.

        original_labels : np.array of shape=(n_samples, n_features)
            Original labels before split.

        data_train : np.array of shape=(n_samples_train, n_features)
            Data selected for training.

        labels_train : np.array of shape=(n_samples_train,)
            Corresponding labels selected for training.

        data_test : np.array of shape=(n_samples_test, n_features)
            Data selected for testing.

        labels_test : np.array of shape=(n_samples_test,)
            Corresponding labels selected for testing.

        expected_samples_train : int
            Number of samples expected to be in training.

        apps_test : list
            List of apps that should only appear in testing dataset.
        """
    # Try casting
    try:
        data_train   = np.asarray(data_train, dtype=float)
        labels_train = np.asarray(labels_train)
        data_test    = np.asarray(data_test , dtype=float)
        labels_test  = np.asarray(labels_test )
    except Exception as e:
        raise ValueError("Output could not be cast to array: '{}'".format(e))

    # Check if apps in original are same as in train+test
    if not (np.unique(original_labels) ==
            np.unique(np.concatenate((labels_train, labels_test)))
           ).all():
        raise ValueError(
            "Labels in train and test were different from original."
        )

    # Check if apps in original are same as in train+test
    if np.isin(labels_train, apps_test).any():
        raise ValueError("Training data contains test labels.")

    # Check if number of train datapoints and labels correspond
    if data_train.shape[0] != labels_train.shape[0]:
        raise ValueError(
            "Number of train datapoints is not equal to number of train labels."
        )

    # Check if number of test datapoints and labels correspond
    if data_test .shape[0] != labels_test .shape[0]:
        raise ValueError(
            "Number of test datapoints is not equal to number of test labels."
        )

    # Check if number of train+test datapoints equals original
    if data_train.shape[0] + data_test.shape[0] != original_data.shape[0]:
        raise ValueError("Train data + test data != original data.")

    # Check if correct number of features is returned
    if data_train.ndim != 2:
        raise ValueError(
            "X_train should be of shape=(n_samples, n_features), i.e. 2 "
            "dimensions, Currently only {} dimension(s) were returned. A likely"
            " cause is that you returned an empty list when no samples were "
            "found, instead try to return a numpy array of "
            "shape=(0, n_features). You can do this with "
            "np.zeros((0, data.shape[0]))."
            .format(data_train.ndim)
        )
    if data_test.ndim != 2:
        raise ValueError(
            "X_test should be of shape=(n_samples, n_features), i.e. 2 "
            "dimensions, Currently only {} dimension(s) were returned. A likely"
            " cause is that you returned an empty list when no samples were "
            "found, instead try to return a numpy array of "
            "shape=(0, n_features). You can do this with "
            "np.zeros((0, data.shape[0]))."
            .format(data_test.ndim)
        )

    # Check if number of features is equal to original
    if  data_train.shape[1] != data_test.shape[1] or\
        data_train.shape[1] != original_data.shape[1]:
        raise ValueError("A different number of features was returned.")

    # Check if labels only has single dimension
    if labels_train.ndim != 1 or labels_test.ndim != 1:
        raise ValueError("Labels should have one dimension.")

    # Check if number of train datapoints is expected
    if data_train.shape[0] != expected_samples_train:
        raise ValueError(
            "Number of training datapoints should be '{}', but was '{}'. Hint: "
            "if you are off by 1, try to round the ratio instead of taking the "
            "floor or ceiling value."
            .format(data_train.shape[0], expected_samples_train)
        )

    expected_test = original_data.shape[0] - expected_samples_train
    # Check if number of test datapoints is expected
    if data_test.shape[0] != expected_test:
        raise ValueError(
            "Number of testing datapoints should be '{}', but was '{}'. Hint: "
            "if you are off by 1, try to round the ratio instead of taking the "
            "floor or ceiling value."
            .format(data_test.shape[0], expected_test)
        )
