"""Module to run train test splits

Author(s):
    Mike Skarlinski (michael.skarlinski@weightwatchers.com)
    Carl Anderson (carl.anderson@weightwatchers.com)

"""
import logging
import pandas as pd
from sklearn.model_selection import train_test_split

from primrose.base.pipeline import AbstractPipeline, PipelineModeType
from primrose.data_object import DataObject, DataObjectResponseType


class TrainTestSplit(AbstractPipeline):
    """Parent pipeline to split into training / testing data, and run a transform

    Note:
        This class will split data into testing and training portions, then write the objects to the data_object
        if this is the only the pipeline operation needed, then you can use this class directly as a pipeline.
        Otherwise, make a child class and write an init_pipeline method to perform operations on your data.

        This class also handles writing your transformer_sequence into the data_object, so there's no need
        to write in child classes.

    """

    def __init__(self, configuration, instance_name):
        super(TrainTestSplit, self).__init__(configuration, instance_name)
        self.training_fraction = self.node_config.get("training_fraction", 0)
        self.seed = self.node_config.get("seed", 0)

    @staticmethod
    def necessary_config(node_config):
        """Return the necessary configuration keys for the DataFrameJoiner object

        Args:
            node_config (dict): set of parameters / attributes for the node

        Note:
            target_variable (string): column name which holds the target variable \
            training_fraction (float): 0->1 float for the fraction of data rows to be used for training \
            seed (int): random number to control the stochastic row selection \

        Returns:
            set of keys

        """
        return set(["training_fraction", "seed"]).union(
            AbstractPipeline.necessary_config(node_config)
        )

    def features(self, data):
        """Use user-specified features if available, otherwise use all non-target columns

        Args:
            data (DataFrame): pandas dataframe

        Returns:
            lsit of feature names

        """
        if "features" in self.node_config:
            return self.node_config["features"]

        else:
            if "target_variable" in self.node_config:
                return [
                    f for f in data.columns if f != self.node_config["target_variable"]
                ]
            else:
                return data.columns

    def _train_test_split(self, data):
        """Split data into test/train sets

        Returns:
            train_data_to_transform (DataFrame)

            test_data_to_transform (DataFrame)

        """
        logging.info("Splitting data into testing and training sets.")

        test_size = 1.0 - float(self.training_fraction)

        if test_size == 0:
            train_data_to_transform = data
            test_data_to_transform = pd.DataFrame()

        else:
            if "target_variable" in self.node_config:
                data_train, data_test, target_train, target_test = train_test_split(
                    data[
                        sorted(
                            list(
                                set(data.columns)
                                - set([self.node_config.get("target_variable")])
                            )
                        )
                    ],
                    data[self.node_config["target_variable"]],
                    test_size=(1.0 - float(self.training_fraction)),
                    random_state=self.seed,
                )

                # re-merge training and target data into a single dataframe for transforming
                train_data_to_transform = pd.concat([data_train, target_train], axis=1)
                test_data_to_transform = pd.concat([data_test, target_test], axis=1)

            else:
                data_train, data_test = train_test_split(
                    data[
                        sorted(
                            list(
                                set(data.columns)
                                - set([self.node_config.get("target_variable")])
                            )
                        )
                    ],
                    test_size=(1.0 - float(self.training_fraction)),
                    random_state=self.seed,
                )

                train_data_to_transform = data_train
                test_data_to_transform = data_test

        logging.info(
            "Training data rows: {}, Testing data rows: {}".format(
                len(train_data_to_transform), len(test_data_to_transform)
            )
        )

        return train_data_to_transform, test_data_to_transform

    def final_data_object_additions(self, data_object):
        """DataObject: Template method to be overloaded in child classes

        Note:
            Allows for child class TransformerSequence objects to be added into the data_object

        Args:
            data_object (DataObject): instance of DataObject

        Returns:
            data_object (DataObject): instance of DataObject

        """
        return data_object

    @staticmethod
    def _concatenate_upstream_dataframes(data):
        """Concatenate multiple upstream data sources into a single dataframe

        Note:

            Upstream objects must add dataframes with the same schema, otherwise mismatched columns will have
            all NULL values in the mutual exclusive indexes

        Args:
            data (list): list of dicts keyed to instances

        Returns:
            dataframe (DataFrame): concatenated dataframes from the data list

        """

        dataframes_to_join = []

        for source in data:

            for key in data[source]:

                if isinstance(data[source][key], pd.DataFrame):

                    if dataframes_to_join:
                        if set(data[source][key].columns) != set(
                            dataframes_to_join[0].columns
                        ):
                            logging.warning(
                                "Concatenated dataframe schemas do not match, subbed with NULL values."
                            )

                    dataframes_to_join.append(data[source][key])

        return pd.concat(dataframes_to_join).reset_index(drop=True)

    def fit_transform(self, data_object):
        """Split data into testing and training sets, then applies the categorical transform to each

        Args:
            data_object (DataObject): instance of DataObject

        Returns:
            data_object (DataObject): instance of DataObject

        """

        # we're can expect multiple objects from a reader, so we need to concatenate
        data_list = data_object.get_upstream_data(
            self.instance_name,
            pop_data=False,
            rtype=DataObjectResponseType.INSTANCE_KEY_VALUE.value,
        )

        data = self._concatenate_upstream_dataframes(data_list)

        train_data, test_data = self._train_test_split(data)

        if not train_data.empty:
            train_data = self.execute_pipeline(
                train_data, PipelineModeType.FIT_TRANSFORM
            )
            data_object.add(
                self,
                train_data[self.features(train_data)],
                key="data_train",
                overwrite=False,
            )

            if "target_variable" in self.node_config:
                data_object.add(
                    self,
                    train_data[self.node_config["target_variable"]],
                    key="target_train",
                    overwrite=False,
                )

        # run the pre-trained pipeline on the testing data
        if not test_data.empty:
            # run the pipeline in Transform mode since we've already fit the pipeline with training data
            test_data = self.execute_pipeline(test_data, PipelineModeType.TRANSFORM)
            self.data = test_data  # assign the data to the testing data if available
            data_object.add(
                self,
                test_data[self.features(train_data)],
                key="data_test",
                overwrite=False,
            )

            if "target_variable" in self.node_config:
                data_object.add(
                    self,
                    test_data[self.node_config["target_variable"]],
                    key="target_test",
                    overwrite=False,
                )

        # save pipeline for writing later
        data_object.add(
            self, self.transformer_sequence, key="transformer_sequence", overwrite=False
        )
        data_object = self.final_data_object_additions(data_object)

        return data_object

    def transform(self, data_object):
        """Transform the data into label encoded data using the pre-trained transformer sequence

        Args:
            data_object (DataObject): instance of DataObject

        Returns:
            data_object (DataObject): instance of DataObject

        """
        # we're can expect multiple objects from a reader, so we need to concatenate
        data_list = data_object.get_upstream_data(
            self.instance_name,
            pop_data=False,
            rtype=DataObjectResponseType.INSTANCE_KEY_VALUE.value,
        )

        data = self._concatenate_upstream_dataframes(data_list)

        data = self.execute_pipeline(data, PipelineModeType.TRANSFORM)

        self.data = (
            data  # keep the data for use in the final_data_object_additions method
        )

        data_object.add(
            self, data[self.features(data)], key="data_test", overwrite=False
        )

        if "target_variable" in self.node_config:
            data_object.add(
                self,
                data[self.node_config["target_variable"]],
                key="target_test",
                overwrite=False,
            )

        data_object = self.final_data_object_additions(data_object)

        return data_object
