"""Module to run a basic decision tree model

Author(s):
    Mike Skarlinski (michael.skarlinski@weightwatchers.com)

"""
import pandas as pd
import numpy as np
import logging
from sklearn import preprocessing

from primrose.base.transformer import AbstractTransformer


class ExplicitCategoricalTransform(AbstractTransformer):

    DEFAULT_NUMERIC = -9999

    def __init__(self, categoricals):
        """initialize the ExplicitCategoricalTransform
        Args:
            categoricals: dictionary containing for each column to be transformed:
            - transformations: list of strings to be executed on the data ('x' represents the current categorical variable)
            - rename: if present, rename the current categorical variable to that name
            - to_numeric: if true, attempt to apply to_numeric after previous transformations

        """
        self.categoricals = categoricals

    def fit(self, data):
        pass

    @staticmethod
    def _process_transformations(data, input_data, categorical, x):
        """transform a column

        Args:
            data (dataframe): dataframe

            input configuration (JSON): JSON categorical config for this variable

            categorical (str): varible name

            x (str): transformation string

        Returns:
            data (dataframe)

        """
        if 'transformations' in input_data.keys():
            logging.info("Applying key {} to variable {}".format('transformations', categorical))
            for transformation in input_data['transformations']:
                exec(transformation.format(x=x))

    @staticmethod
    def _process_rename(data, input_data, categorical):
        """rename a field

        Args:
            data (dataframe): dataframe

            input configuration (JSON): JSON categorical config for this variable

            categorical (str): varible name

        Returns:
            (tuple): tuple containing:

                data (dataframe): dataframe

                name (str): original name (if not "to_numeric": True), new_name otherwise
        """
        if 'rename' in input_data.keys():
            logging.info("Applying key {} to variable {}".format('rename', categorical))
            data = data.rename({categorical: input_data['rename']}, axis='columns')
            return data, input_data['rename']
        return data, categorical

    @staticmethod
    def _process_numeric(data, input_data, name):
        """convert column to numeric

        Args:
            data (dataframe): dataframe

            input configuration (JSON): JSON categorical config for this variable

            name (str): field name

        Returns:
            data with the colun converted to numeric

        """
        if input_data.get('to_numeric', False):

            logging.info("Applying key {} to variable {}".format('to_numeric', name))

            # if there are errors converting to numerical values, we need to sub in a reasonable value
            if sum(pd.to_numeric(data[name], errors='coerce').isnull()) > 0:
                logging.info("Can't convert these entries in {}. Replacing with {}: {}".format(
                    name, ExplicitCategoricalTransform.DEFAULT_NUMERIC,
                    np.unique(data[name][pd.to_numeric(data[name], errors='coerce').isnull()].astype(str))))

                data[name][pd.to_numeric(data[name], errors='coerce').isnull()] = ExplicitCategoricalTransform.DEFAULT_NUMERIC
            try:
                data[name] = pd.to_numeric(data[name])
                return data

            except:
                raise TypeError('Failed to convert feature {} to numeric'.format(name))

        else:
            return data

    def transform(self, data):
        """Transform categorical variables into one or more numeric ones, no need to separate testing & training data

        Args:
            data: dictionary containing dataframe with all categorical columns present

        Returns:
            data with all categorical columns recoded and/or deleted

        """
        for categorical in self.categoricals.keys():

            x = "data['{}']".format(categorical)

            input_data = self.categoricals[categorical]

            ExplicitCategoricalTransform._process_transformations(data, input_data, categorical, x)

            data, new_name = ExplicitCategoricalTransform._process_rename(data, input_data, categorical)

            data = ExplicitCategoricalTransform._process_numeric(data, input_data, new_name)

        return data


class ImplicitCategoricalTransform(AbstractTransformer):
    """Class which implicitly transforms all string columns of a dataframe with sklearn LabelEncoder"""

    def __init__(self, target_variable):
        """initialize this ImplicitCategoricalTransform

        Args:
            target_variable (str): target variable name

        """
        self.target_variable = target_variable
        self._encoder = {}
        self.target_encoder = None

    def fit(self, data):
        """encode the data as categorical labels

        Args:
            data (dataframe)

        Returns:
            dataframe (dataframe)

        """

        logging.info('Fitting LabelEncoders on all string-based dataframe columns...')

        data.is_copy = False

        for column_name in data.columns:

            if data[column_name].dtype == object:

                logging.info('Fitting LabelEncoder for column {}'.format(column_name))

                self._encoder[column_name] = preprocessing.LabelEncoder()
                self._encoder[column_name].fit(data[column_name])

                if column_name == self.target_variable:
                    self.target_encoder = self._encoder[column_name]
            else:
                pass

        return data

    def transform(self, data):
        """Transform data into categorical variables using pre-trained label encoder

        Args:
            data (dataframe)

        Returns:
            dataframe (dataframe)

        """

        data.is_copy = False

        for column_name in data.columns:

            if column_name in self._encoder:

                logging.info('LabelEncoding column {}'.format(column_name))

                data[column_name] = self._encoder[column_name].transform(data[column_name])

        return data
