#!/usr/bin/env python3

## Preprocessing functions
# Copyright (C) <2018-2022>  <Agence Data Services, DSI Pôle Emploi>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.


import os
import re
import json
import logging
import functools
import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator
from typing import List, Dict, Tuple, Any, Union, Optional
from sklearn.utils.validation import check_is_fitted
from sklearn.preprocessing._function_transformer import FunctionTransformer

logger = logging.getLogger(__name__)


class Estimator(BaseEstimator):
    '''Base class for the classes defined below. Implements _validate_input and fit_transform.'''

    def __init__(self, input_length: Union[int, None]) -> None:
        '''Initialization of the class

        Args:
            input_length (int): The number of columns of the input (used in _validate_input())
        '''
        self.input_length = input_length

    def _validate_input(self, X: Union[np.ndarray, pd.DataFrame]) -> Union[np.ndarray, pd.DataFrame]:
        '''Validates input format

        Args:
            X (np.ndarray or pd.DataFrame): Input to validate
        Raises:
            ValueError: If the shape of the input does not correspond to self.input_length
        Returns:
            np.ndarray or pd.DataFrame: A copy of X
        '''
        if self.input_length is not None and X.shape[1] != self.input_length:
            raise ValueError(f"Bad shape: ({X.shape[1]} != {self.input_length})")

        # Mandatory copy in order not to modify the original !
        if isinstance(X, pd.DataFrame):
            return X.copy(deep=True)
        else:
            return X.copy()

    def fit(self, X: Union[np.ndarray, pd.DataFrame], y: Union[np.ndarray, pd.Series, pd.DataFrame]) -> Any:
        '''Fit transformer

        Args:
            X (np.ndarray or pd.DataFrame): Array-like, shape = [n_samples, n_features]
            y (np.ndarray or pd.Series or pd.DataFrame): Array-like, shape = [n_samples, n_targets]
        Returns:
            self
        '''
        raise NotImplementedError("'fit' needs to be overridden")

    def transform(self, X: Union[np.ndarray, pd.DataFrame]) -> np.ndarray:
        '''Transforms X

        Args:
            X (np.ndarray or pd.DataFrame): Array-like, shape = [n_samples, n_features]
        Returns:
            Transformed X
        '''
        raise NotImplementedError("'transform' needs to be overridden")

    def fit_transform(self, X: Union[np.ndarray, pd.DataFrame], y: Union[np.ndarray, pd.Series, pd.DataFrame, None] = None) -> np.ndarray:
        '''Applies both fit & transform.

        Args:
            X (np.ndarray or pd.DataFrame): Shape = [n_samples, n_features]
        Kwargs:
            y: Array-like, shape = [n_samples]
        Returns:
            X_out: Array-like, shape [n_samples, n_features]
                        Transformed input.
        '''
        self.fit(X, y)
        return self.transform(X)


class AutoLogTransform(Estimator):
    '''Automatically applies a log transformation on numerical data if the distribution of the variables
       is skewed (abs(skew) > min_skewness) and if there is an amplitude superior to min_amplitude between
       the 10th and 90th percentiles

    WARNING: YOUR DATA MUST BE POSITIVE IN ORDER FOR EVERYTHING TO WORK CORRECTLY
    '''

    def __init__(self, min_skewness: float = 2, min_amplitude: float = 10E3) -> None:
        '''Initialization of the class

        Kwargs:
            min_skewness (float): Absolute value of the required skewness to apply a log transformation
            min_amplitude (float): Minimal value of the amplitude between the 10th and 90th percentiles
                                   required to apply a log transformation
        '''

        super().__init__(None)
        # Set attributes
        self.min_skewness = min_skewness
        self.min_amplitude = min_amplitude

        # Columns on which to apply the transformation
        # Set on fit
        # Warning: sklearn does not support columns name, so we can only use indexes
        # Hence, X input must expose same columns order (this won't be checked)
        self.applicable_columns_index: Optional[List[Any]] = None

    def fit(self, X: Union[np.ndarray, pd.DataFrame], y: Any = None) -> Any:
        '''Fit transformer

        Args:
            X (np.ndarray or pd.DataFrame): Array-like, shape = [n_samples, n_features]
        Kwargs:
            y (None): Not used here
        Returns:
            self
        '''
        X = self._validate_input(X)
        # If x is a numpy array, casts it in pd.DataFrame
        if isinstance(X, np.ndarray):
            X = pd.DataFrame(X)
        # Otherwise, we reset the columns name because sklearn can't manage them
        else:
            X = X.rename(columns={col: i for i, col in enumerate(X.columns)})
        self.input_length = X.shape[1]

        # Get applicable columns
        skew = X.skew()
        candidates = list(skew[abs(skew) > self.min_skewness].index)
        if len(candidates) > 0:
            q10 = X.iloc[:, candidates].quantile(q=0.1)
            q90 = X.iloc[:, candidates].quantile(q=0.9)
            amp = q90 - q10
            # Update applicable_columns_index
            self.applicable_columns_index = list(amp[amp > self.min_amplitude].index)

        self.fitted_ = True
        return self

    def transform(self, X: Union[np.ndarray, pd.DataFrame]) -> np.ndarray:
        '''Transforms X - apply log on applicable columns

        Args:
            X (np.ndarray or pd.DataFrame): Array-like, shape = [n_samples, n_features]
        Returns:
            X_out: Array-like, shape [n_samples, n_features]
                        Transformed input.
        '''
        # Validate input
        check_is_fitted(self, 'fitted_')
        X = self._validate_input(X)

        # If x is a numpy array, casts it in pd.DataFrame
        if isinstance(X, np.ndarray):
            X = pd.DataFrame(X)

        # Log transformation on the applicable columns
        if len(self.applicable_columns_index) > 0:
            X.iloc[:, self.applicable_columns_index] = np.log(X.iloc[:, self.applicable_columns_index])

        # Compatibility -> returns numpy array
        return X.to_numpy()


class ThresholdingTransform(Estimator):
    '''Applies a threshold on columns.
    If min and max values are given, the threshold is manual; otherwise it is statistical.
    '''

    def __init__(self, thresholds: List[Tuple], quantiles: tuple = (0.05, 0.95)) -> None:
        '''Initialization of the class

        Args:
            tresholds (list of tuple): Each tuple contains (min_val, max_val).
        Kwargs:
            quantiles (tuple): Tuple containing (quantile_min, quantile_max)
        Raises:
            ValueError: If quantiles values are not between 0 and 1 and if quantiles[0] >= quantiles[1]
        '''
        if not 0 < quantiles[0] < 1 or not 0 < quantiles[1] < 1 or not quantiles[0] < quantiles[1]:
            raise ValueError(f"The values contained in quantiles should verify quantile_min < quantile_max and both > 0 and < 1. quantiles = {quantiles} is not supported.")

        super().__init__(len(thresholds))

        # Set attributes
        self.thresholds = thresholds
        self.fitted_thresholds: List[tuple] = []
        self.quantiles = quantiles

    def fit(self, X: Union[np.ndarray, pd.DataFrame], y: Any = None) -> Any:
        '''Fits the ThresholdingTransform on X.

        Args:
            X (np.ndarray or pd.DataFrame): Shape [n_samples, n_features]
        Kwargs:
            y (None): Not used here.
        Returns:
            self: ThresholdingTransform
        '''
        X = self._validate_input(X)
        # If X is a numpy array, casts it as a pd.DataFrame
        if isinstance(X, np.ndarray):
            X = pd.DataFrame(X)

        # Fits column one by one
        for col_index, item in enumerate(self.thresholds):
            val_min, val_max = item
            if val_min is None:
                val_min = X.iloc[:, col_index].quantile(q=self.quantiles[0])
            if val_max is None:
                val_max = X.iloc[:, col_index].quantile(q=self.quantiles[1])
            self.fitted_thresholds.append((col_index, val_min, val_max))

        self.fitted_ = True
        return self

    def transform(self, X: Union[np.ndarray, pd.DataFrame]) -> np.ndarray:
        '''Impute all missing values in X.

        Args:
            X (np.ndarray or pd.DataFrame): Shape (n_samples, n_features)
                The input data to complete.
        '''
        check_is_fitted(self, 'fitted_')
        X = self._validate_input(X)
        # If X is a numpy array, casts it to pd.DataFrame
        if isinstance(X, np.ndarray):
            X = pd.DataFrame(X)

        for col_index, val_min, val_max in self.fitted_thresholds:
            X.iloc[:, col_index][X.iloc[:, col_index] < val_min] = val_min
            X.iloc[:, col_index][X.iloc[:, col_index] > val_max] = val_max

        return X.to_numpy() # Compatibility -> return a numpy array


class AutoBinner(Estimator):
    '''Automatically creates a "other" category when the categories are heavily unbalanced
    /!\ Replaces the values of some categories /!\
    '''

    def __init__(self, strategy: str = "auto", min_cat_count: int = 3, threshold: float = 0.05) -> None:
        '''Initialization of the class

        Kwargs:
            strategy (str): 'auto' or 'threshold'
                - 'auto': Aggregates all categories as long as their cumulated frequence is less than threshold
                - 'threshold': Aggregates all category whose frequence is less than threshold
            min_cat_count (int): Minimal number of category to keep
            threshold (float): The threshold to consider
        Raises:
            ValueError: The object strategy must be in the list of allowed strategies
            ValueError: The object min_cat_count must be non negative
            ValueError: The object threshold must be in ]0,1[
        '''
        super().__init__(None)

        allowed_strategies = ["threshold", "auto"]
        self.strategy = strategy
        if self.strategy not in allowed_strategies:
            raise ValueError(f"Can only use these strategies: {allowed_strategies}. " +\
                             f"Got strategy={strategy}")
        if min_cat_count < 0:
            raise ValueError("min_cat_count must be non negative")
        if not (0 < threshold < 1):
            raise ValueError(f"threshold must be in ]0,1[, not {threshold}")

        # Set attributes
        self.min_cat_count = min_cat_count
        self.threshold = threshold
        self.kept_cat_by_index: Dict[int, list] = {}

    def fit(self, X: Union[np.ndarray, pd.DataFrame], y: Any = None) -> Any:
        '''Fit the AutoBinner on X.

        Args:
            X (np.ndarray or pd.DataFrame): Shape (n_samples, n_features)
        Kwargs:
            y: Not used here
        Returns:
            self: AutoBinner
        '''
        X = self._validate_input(X)
        # If x is a numpy array, casts it in pd.DataFrame
        if isinstance(X, np.ndarray):
            X = pd.DataFrame(X)

        self.input_length = X.shape[1]
        # Fits column one by one
        for col_index in range(self.input_length):
            # Get col serie
            X_tmp_ser = X.iloc[:, col_index]
            # Get unique vals
            unique_cat = list(X_tmp_ser.unique())
            # If less vals than min threshold, set this column allowed values with all uniques values
            if len(unique_cat) <= self.min_cat_count:
                self.kept_cat_by_index[col_index] = unique_cat
                continue

            # If more vals than min threshold, keep values based on strategy
            table = X_tmp_ser.value_counts() / X_tmp_ser.count()
            table = table.sort_values()
            if self.strategy == 'auto':
                table = np.cumsum(table)
            # If only one category is less than the threshold, we do not need to transform it
            if table[1] > self.threshold:
                self.kept_cat_by_index[col_index] = unique_cat
                continue
            # Otherwise, we get rid of the superfluous categories
            else:
                to_remove = list(table[table<self.threshold].index)
                for item in to_remove:
                    unique_cat.remove(item)
                self.kept_cat_by_index[col_index] = unique_cat

        self.fitted_ = True
        return self

    def transform(self, X: Union[np.ndarray, pd.DataFrame]) -> np.ndarray:
        '''Imputes all missing values in X.

        Args:
            X (np.ndarray or pd.DataFrame): Shape (n_samples, n_features)
                The input data to complete.
        '''
        check_is_fitted(self, 'fitted_')
        X = self._validate_input(X)
        # If x is a numpy array, casts it in pd.DataFrame
        if isinstance(X, np.ndarray):
            X = pd.DataFrame(X)

        for col_index in range(self.input_length):
            X.iloc[:, col_index] = X.iloc[:, col_index].apply(lambda x: x if x in self.kept_cat_by_index[col_index] else 'other_')

        return X.to_numpy() # Compatibility, returns a numpy array


class EmbeddingTransformer(Estimator):
    '''Constructs a transformer that apply an embedding mapping to Categorical columns'''

    def __init__(self, embedding: Union[str, dict], none_strategy: str = 'zeros') -> None:
        '''Initialization of the class

        Args:
            embedding (str or dict): Embedding to use
                - If dict -> ok, ready to go
                - If str -> path to the file to load (json)
        Kwargs:
            none_strategy (str): Strategy to fill elements not in embedding
                - Zeros: only 0s
        Raises:
            ValueError: If strategy is not in the allowed strategies
            ValueError: If the embedding is of type str but does not end in .json
            FileNotFoundError: If the path to the embedding does not exist
        '''
        super().__init__(None)

        # Check none strategy
        allowed_strategies = ["zeros"]
        self.none_strategy = none_strategy
        if self.none_strategy not in allowed_strategies:
           raise ValueError(f"Can only use these strategies: {allowed_strategies}, got strategy={self.none_strategy}")

        # If str, loads the embedding
        if isinstance(embedding, str):
            if not embedding.endswith('.json'):
                raise ValueError(f"The file {embedding} must be a .json file")
            if not os.path.exists(embedding):
                raise FileNotFoundError(f"The file {embedding} does not exist")
            with open(embedding, 'r', encoding='{{default_encoding}}') as f:
                self.embedding = json.load(f)
        else:
            self.embedding = embedding

        # Get embedding size
        self.embedding_size = len(self.embedding[list(self.embedding.keys())[0]])
        # Other params
        self.n_missed = 0

    def fit(self, X: Union[np.ndarray, pd.DataFrame], y: Any = None) -> Any:
        '''Fit transformer

        Args:
            X (np.ndarray or pd.DataFrame): Shape (n_samples, n_features)
        Kwargs:
            y: Not used here
        Returns
            self (EmbeddingTransformer)
        '''
        X = self._validate_input(X)
        # If x is a numpy array, casts it in pd.DataFrame
        if isinstance(X, np.ndarray):
            X = pd.DataFrame(X)

        self.input_length = X.shape[1]

        # Nothing to do

        self.fitted_ = True
        return self

    def transform(self, X: Union[np.ndarray, pd.DataFrame]) -> np.ndarray:
        '''Transform X - embedding mapping

        Args:
            X (np.ndarray or pd.DataFrame): Shape (n_samples, n_features)
        Raises:
            ValueError: If there are missing columns
        Returns:
            X_out (np.ndarray): Shape (n_samples, n_features) transformed input.
        '''
        X = self._validate_input(X)
        # If x is a numpy array, casts it in pd.DataFrame
        if isinstance(X, np.ndarray):
            X = pd.DataFrame(X)

        n_rows = X.shape[0]

        # Apply mapping
        new_df = pd.DataFrame()
        for col in X.columns:
            self.n_missed = 0  # Counts the number of missing elements in the embedding
            tmp_serie = X[col].apply(self.apply_embedding)  # Updates self.n_missed
            new_df = pd.concat([new_df, pd.DataFrame(tmp_serie.to_list())], axis=1)
            perc_missed = self.n_missed / n_rows * 100
            if perc_missed != 0:
                logger.warning(f"Warning, {self.n_missed} ({perc_missed} %) missing elements in the embedding for column {col}")

        return new_df.to_numpy() # Compatibility, returns a numpy array

    def apply_embedding(self, content) -> list:
        '''Applies embedding mapping

        Args:
            content: Content on which to apply embedding mapping
        Raises:
            ValueError: If the strategy is not recognized
        Returns:
            list: Applied embedding
        '''
        if content in self.embedding.keys():
            return self.embedding[content]
        else:
            self.n_missed += 1
            if self.none_strategy == 'zeros':
                return [0] * self.embedding_size
            else:
                raise ValueError(f"Strategy {self.none_strategy} not recognized")

    def get_feature_names(self, features_in: list, *args, **kwargs) -> np.ndarray:
        '''Returns feature names for output features.

        Args:
            features_in (list): list of features

        Returns:
            output_feature_names: ndarray of shape (n_output_features,)
                Array of feature names.
        '''
        check_is_fitted(self, 'fitted_')
        new_features = [f"emb_{feat}_{i}" for feat in features_in for i in range(self.embedding_size)]
        return np.array(new_features, dtype=object)


if __name__ == '__main__':
    logger.error("This script is not stand alone but belongs to a package that has to be imported.")
