""" 
This module was created to handle DataFrames and returns 
pandas DataFrame or pandas Series.

"""


from drtools.utils import list_ops
from typing import List, Union, Tuple
from pandas import DataFrame, Series, cut
from scipy import stats
import numpy as np
import math
from copy import deepcopy
import pandas as pd


ByPercentage = 'by-percentage'
ByCount = 'by-count'
Lower = 'lower'
Upper = 'upper'


def combine_variables(
    dataframe: DataFrame,
    varibles: List[str],
    sep: str='|'
) -> DataFrame:
    """Combine categorical variables of DataFrame in order to analyse 
    impact of combined variables on response.

    Parameters
    ----------
    dataframe : DataFrame
        The DataFrame
    varibles : List[str]
        List of categorical variables that will be combined
    sep : str, optional
        Separation character when combine variables, 
        by default '|'

    Returns
    -------
    DataFrame
        The DataFrame with new colum containing all combinations 
        of desired variables.
    """
    
    df = dataframe.copy()
    title = sep.join(varibles)
    def get_all_values(x):
        text = ''
        variables_len = len(varibles)
        for idx, var in enumerate(varibles):
            text += str(x[var])
            if idx != variables_len - 1:
                text += str(sep)
        return text
    df[title] = df.apply(get_all_values, axis=1)
    return df


def remove_boundary_values(
    dataframe: DataFrame,
    by: Union[ByPercentage, ByCount]=ByPercentage,
    value: float=0.05,
    boundary: Union[Lower, Upper]=Upper
) -> DataFrame:
    """Remove values of top part of the DataFrame

    Parameters
    ----------
    df : DataFrame
        The DataFrame
    by : Union[ByPercentage, ByCount], optional
        Rule that will be applied to remove top
        values, by default ByPercentage
    value : float, optional
        The value of the rule, by default 0.05
    boundary : Union[Lower, Upper], optional
        Extremity from where values 
        will be removed, by default 0.05

    Returns
    -------
    DataFrame
        The DataFrame after removing selected values of the 
        top part of DataFrame

    Raises
    ------
    ValueError
        If **by** option is invalid.
    """
    
    methods = {
        Lower: {
            ByPercentage: math.floor(value * len(dataframe)),
            ByCount: value,
            'df_filter': lambda x: dataframe.iloc[x:]
        },
        Upper: {
            ByPercentage: math.ceil((1-value) * len(dataframe)),
            ByCount: value,
            'df_filter': lambda x: dataframe.iloc[:x]
        }
    }    
    limit = methods[boundary][by]
    df_copy = methods[boundary]['df_filter'](limit)
    return df_copy


def remove_outliers(
    df: DataFrame,
    col_name: str,
    zscore_limit: int=3,
    fillna: any=None
) -> DataFrame:
    """Remove outliers using zscore rule
    
    Only will be keeped values with zscore lower than **zscore_limit**

    Parameters
    ----------
    df : DataFrame
        The DataFrame
    col_name : str
        The column name where outliers will be removed
    zscore_limit : int, optional
        Z-Score limit, by default 3
    fillna : any, optional
        Value to fill na before apply the zscore rule, by default None

    Returns
    -------
    DataFrame
        The DataFrame after removing the outliers
    """
    df_copy = df.copy()
    if fillna is not None:
        df_copy[col_name] = df_copy[col_name].fillna(fillna)
    df_copy = df[
        np.abs(stats.zscore(df_copy[col_name])) < zscore_limit
    ]
    return df_copy


ColumnInfo = dict({
    str: float
})


def keep_categories(
    df: DataFrame,
    col_name: str,
    keep_categories: List[str]=[]
) -> DataFrame:
    """Keep only selected categories from column of DataFrame.

    Parameters
    ----------
    df : DataFrame
        The DataFrame
    col_name : str
        The column name where categories will be keeped
    keep_categories : List[str], optional
        List of categories that will be keeped, by default []

    Returns
    -------
    DataFrame
        The DataFrame after remove unselected categories and keeped
        selected categories
    """
    df_copy = df[col_name].isin(
        keep_categories
    )
    return df[df_copy]


def join_categories(
    dataframe: DataFrame,
    col_name: str,
    categories: List[str]=[], 
    to_category: str='others'
) -> DataFrame:
    """Join categories and fill these values with desired value.

    Parameters
    ----------
    dataframe : DataFrame
        The DataFrame
    col_name : str
        The column name where categories will be joined.
    categories : List[str], optional
        List of categories will be joined, by default []
    to_category : str, optional
        New value for selected categories, by default 'others'

    Returns
    -------
    DataFrame
        The DataFrame after join the categories
    """
    df = dataframe.copy()
    if len(categories) == 0:
        return df
    condition = df[col_name].isin(categories)
    df.loc[condition, col_name] = to_category
    return df


def keep_categories_and_join_remaining(
    dataframe: DataFrame,
    column: str,
    categories: List[str]=[], 
    to_category: str='others'
) -> DataFrame:
    """Join categories and fill these values with desired value.

    Parameters
    ----------
    dataframe : DataFrame
        The DataFrame
    column : str
        The column name where categories will be joined.
    categories : List[str], optional
        List of categories will be joined, by default []
    to_category : str, optional
        New value for selected categories, by default 'others'

    Returns
    -------
    DataFrame
        The DataFrame after join the categories
    """
    df = dataframe.copy()
    if len(categories) == 0:
        return df
    df[column] = np.where(
        pd.isnull(df[column]), 
        df[column], 
        df[column].astype(str)
    ) 
    join_cat = list(df[column].dropna().unique())
    join_cat = list_ops(join_cat, categories)
    df = join_categories(df, column, join_cat, to_category)
    return df


def drop_categories(
    dataframe: DataFrame,
    col_name: str,
    categories: List[str],
) -> DataFrame:
    """Drop categories of column

    Parameters
    ----------
    dataframe : DataFrame
        The DataFrame
    col_name : str
        The column name where categories will be dropped
    categories : List[str]
        List of categories that will be dropped

    Returns
    -------
    DataFrame
        The DataFrame after drop the categories
    """
    condition = ~dataframe[col_name].isin(
        categories
    )
    return dataframe[condition]


def get_categories_ge(
    df: DataFrame,
    col_name: str,
    by: Union[ByPercentage, ByCount]=ByPercentage,
    value: float=0.01,
    ignore_categories: List[str]=[]
) -> List[str]:
    """Get name of categories that contains percentage or counts 
    of total samples greater or equal than value.

    Parameters
    ----------
    df : DataFrame
        DataFrame where categoreis will be found
    col_name : str
        Column name where categories will be found
    by : Union[ByPercentage, ByCount], optional
        Rule that will be applied to get categories based
        on percentage of total or count, by default ByPercentage
    value : float, optional
        The value of the rule, by default 0.01
    ignore_categories : List[str], optional
        Ignore this categories when apply percentage 
        filter, by default []

    Returns
    -------
    List[str]
        The list of categories with at least desired percentage of total
    """
    resp = []
    v = df[col_name].value_counts()
    col_len = df[col_name].notna().sum()
    
    for category, category_count_value in zip(v.index, v.values):
        if category in ignore_categories:
            continue
        
        if by == ByPercentage:
            compare_value = category_count_value / col_len
        elif by == ByCount:
            compare_value = category_count_value
        else:
            raise ValueError('Invalid "by" option.')       
        
        if compare_value >= value:
            resp.append(category)
    return resp


def get_labels_from_bin_interval(
    bins: List[float],
) -> List[str]:
    """Get bins labels from bins interval of values

    Parameters
    ----------
    bins : List[float]
        List containing bins boundary

    Returns
    -------
    List[str]
        The list with labels for each interval
    """
    return [f'{x}_{y}' for x, y in zip(bins[:-1], bins[1:])]


def prepare_bins(
    bins: Union[int, List[float]],
    smallest_value: float=None,
    bigger_value: float=None,
    include_upper: bool=False
) -> Tuple[List[float], List[float], List[str]]:
    """Generate values of bins, middle points and labels for the bins
    
    This function will generate points which represents the x axis
    for the bins, the middle points of the intervals of the bins and
    the labels for each bin.

    Parameters
    ----------
    bins : Union[int, List[float]]
        The desired number of bins that will be generated
        or the bins interval
    smallest_value : float, optional
        Smallest value from the interval, by default None.
    bigger_value : float, optional
        Bigger value from the interval, by default None.
    include_upper : bool, optional
        Whether the last interval should 
        be right-inclusive or not, by default False

    Returns
    -------
    Tuple[List[float], List[float], List[str]]
        Returns the values which represents the intervals 
        from the x axis, the middle points of these intervals and
        the labels for the bins interval.
    """
    values = deepcopy(bins) \
        if hasattr(bins, '__iter__') \
        else np.linspace(smallest_value, bigger_value, bins + 1)
    if include_upper:
        upper_boundary = round(values[-1] * (1 + 0.001), 4) \
            if values[-1] != 0 \
            else 1e-4
        values[-1] = upper_boundary
    middle_points = [(x + y) / 2 for x, y in zip(values[:-1], values[1:])]
    labels = get_labels_from_bin_interval(values)
    labels = [
        f'[{label})'.replace('_', ', ')
        for label in labels
    ]
    return values, middle_points, labels
    

def binning_numerical_variable(
    dataframe: DataFrame,
    col_name: str,
    bins: Union[int, List[float]],
    binning_column_prefix: str='binning',
    include_upper: bool=False,
    labels: List[str]=None
) -> DataFrame:
    """Binning numerical variable.

    Parameters
    ----------
    dataframe : DataFrame
        The DataFrame from where numerical variable will be found.
    col_name : str
        The numerical column name.
    bins : Union[int, List[float]]
        Desired number of bins or interval
        of bins
    binning_column_prefix : str, optional
        Prefix of new columns that will be generated
        with the binning of numerical variable, by default 'binning'
    include_upper : bool
        Whether the last interval should 
        be right-inclusive or not, by default False
    labels : List[str]
        Labels for the bins interval, by default **None**

    Returns
    -------
    DataFrame
        The DataFrame with new column containing the binned values
        for the numerical variable
    """
    
    df = dataframe.copy()
        
    sorted_values = df[col_name] \
        .dropna() \
        .sort_values() \
        .values
    
    smallest_value = sorted_values[0]
    bigger_value = sorted_values[-1]
    values, middle_points, _labels = prepare_bins(
        bins,
        smallest_value=smallest_value,
        bigger_value=bigger_value,
        include_upper=include_upper
    )
    df[f'{binning_column_prefix}_{col_name}'] = cut(
        df[col_name], 
        bins=values,
        labels=labels or _labels,
        right=False
    )
    
    return df


def join_categories_lt(
    dataframe: DataFrame,
    col_name: str,
    by: Union[ByPercentage, ByCount]=ByPercentage,
    value: float=0.01,
    to_category_name: str='others',
) -> DataFrame:
    """Fill value with 'to_category_name' of categories 
    which has count lower than 'value' in count or percentage

    Parameters
    ----------
    dataframe : DataFrame
        DataFrame
    col_name : str
        Name of categorical column
    by : Union[ByPercentage, ByCount], optional
        Rule that will be applied to join categories based
        on percentage of total or count, by default ByPercentage
    value : float, optional
        The value of the rule, by default 0.01
    to_category_name : str, optional
        Value to fill categories, by default 'others'

    Returns
    -------
    DataFrame
        DataFrame after applied changes. 
    """
    df = dataframe.copy()
    keep_categories = get_categories_ge(
        df, 
        col_name, 
        by=by,
        value=value,
    )
    all_categories = list(df[col_name].value_counts().index)
    join_categories_list = list_ops(
        all_categories, 
        keep_categories
    )
    df = join_categories(
        df, 
        col_name, 
        categories=join_categories_list,
        to_category=to_category_name
    )
    return df


def keep_categories_ge(
    df: DataFrame,
    col_name: str,
    by: Union[ByPercentage, ByCount]=ByPercentage,
    value: float=0.01,
    force_remove: List[str]=[]
) -> DataFrame:
    """Keep categories of 'col_name' which contains 
    percentage or counts greater or equal than 'value'.

    Parameters
    ----------
    df : DataFrame
        Pandas DataFrame
    col_name : str
        Column name to keep categories
    by : Union[ByPercentage, ByCount], optional
        Rule that will be applied to keep categories based
        on percentage of total or count, by default ByPercentage
    value : float, optional
        Keep categories which has at least
        this percentage or count of total 
        samples, by default 0.01
    force_remove : List[str], optional
        Remove columns from this list even if this columns
        doesn't match the percentage condition, by default []

    Returns
    -------
    DataFrame
        DataFrame with only desired categories
    """
    
    keep_categories = get_categories_ge(
        df, 
        col_name, 
        by=by,
        value=value,
    )
    keep_categories = list_ops(
        keep_categories, 
        force_remove
    )
    df_copy = keep_categories(
        df, 
        col_name, 
        keep_categories=keep_categories
    )
    return df_copy


def na_percentage_of_column(
    df: DataFrame,
    col_name: str,
) -> float:
    """Number representing NA percentage from column of DataFrame 

    Parameters
    ----------
    df : DataFrame
        The DataFrame
    col_name : str
        Columns name to verify nullity

    Returns
    -------
    float
        The percentage of na for the columns passed as 'col_name'
        argument
    """
    return df[col_name].isna().sum() / len(df[col_name])


def percentage_of_most_common_value_by_column(
    df: DataFrame,
    ignore_columns: List[str]=[],
    notna: bool=False
) -> ColumnInfo:
    """Percentage of the most common value for each column.

    Parameters
    ----------
    df : DataFrame
        Pandas DataFrame
    ignore_columns : List[str], optional
        Ignore columns when get percentage, by default []
    notna : bool, optional
        If True, percentage will be based on not na values,
        if False, percentage will be calculated considering na
        values, by default False

    Returns
    -------
    ColumnInfo
        Dict containing information about percentage of most common
        value from each column.
    """
    
    resp = {}
    for col in df.columns:
        if col in ignore_columns:
            continue
        
        data_series = df[col].dropna() \
            if notna \
            else df[col].fillna('nan')
        
        v = data_series.value_counts()
        num_of_occ = v.sort_values().values
        num_of_occ = num_of_occ[-1]
        total_len = len(data_series)
        value_percentage = num_of_occ / total_len
        
        resp[col] = round(value_percentage * 100, 2)
    return resp


def na_percentage_info(
    df: DataFrame,
    ignore_columns: List[str]=[]
) -> ColumnInfo:
    """Get percentage of null of all columns.

    Parameters
    ----------
    df : DataFrame
        Pandas DataFrame.
    ignore_columns : List, optional
        Ignore these columns when check DataFrame, by default []

    Returns
    -------
    ColumnInfo
        Information of null percentage about each column
    """
    
    resp = {}
    for col in df.columns:
        if col not in ignore_columns:
            resp[col] = round(
                na_percentage_of_column(df, col) * 100, 
                2
            )
    return resp


def get_column_names(
    df: DataFrame,
    contains: List[str]=[]
) -> List[str]:
    """Get name of columns that contains some of substrings in 'contains'

    Parameters
    ----------
    df : DataFrame
        DataFrame from which columns will be generated
    contains : List[str], optional
        List of substring to filter DataFrame columns, by default []

    Returns
    -------
    List[str]
        The list of column names which has some of the substrings
        in 'contains'
    """
    
    names = []
    if len(contains) == 0:
        contains = ['']
    for col in df.columns:
        for col_name in contains:
            if col_name in col:
                names.append(col)
                break
    return names


def drop_df_columns_gt_constancy(
    dataframe: DataFrame,
    constacy_level: float=0.75,
    verbose: bool=True,
    print_removed_columns: bool=False,
) -> DataFrame:
    """Drop columns based on variable constancy.
    
    Drop columns that has some class or value with count percentage 
    bigger than constancy_level.

    Parameters
    ----------
    dataframe : DataFrame
        DataFrame that columns will be dropped
    constacy_level : float, optional
        Columns that has some value with percentage of counts 
        above constancy level will be dropped, must be
        between 0 and 1, by default 0.8
    verbose : bool, optional
        If True print extra information, by default True
    print_removed_columns : bool, optional
        If True print removed columns, by default False

    Returns
    -------
    DataFrame
        DataFrame after drop columns
    """
    df = dataframe.copy()
    drop_columns = []
    for col in df.columns:
        v = df[col].value_counts()
        num_of_occ = v.sort_values().values[-1]
        total_len = len(df[col])
        variable_constancy_level = num_of_occ / total_len
        if variable_constancy_level >= constacy_level:
            drop_columns.append(col)
    df.drop(columns=drop_columns, axis=1, inplace=True)
    if verbose:
        len_initial = len(dataframe.columns)
        len_final = len(df.columns)
        print(f'From {len_initial} to {len_final}. Remove: {len_initial - len_final}')
    if print_removed_columns:
        removed = [col for col in df.columns if col not in df.columns]
        print(removed)
    return df


def merge_columns(
    dataframe: DataFrame,
    cols: List[str],
) -> Series:
    """Create pandas Series containing the result of 
    combination of some list of columns.
    
    The resulting Series will contain the first filled value 
    from a column passed on variable `cols`.

    Parameters
    ----------
    dataframe : DataFrame
        The DataFrame.
    cols : List[str]
        List of column names.

    Returns
    -------
    Series
        The resulting Series.
    """
    
    df = dataframe[cols].copy()
    series = dataframe[cols[0]]
    
    for col in cols[1:]:
        series = series.fillna(df[col])
    
    return series


def reduce_mem_usage(
    dataframe: DataFrame, 
    verbose=True
):
    """Reduce memory usage from DataFrame
    
    Parameters
    ----------
    dataframe : DataFrame
        The DataFrame.
    verbose : bool, optional
        If True, will print reduction information, 
        by default True.

    Returns
    -------
    DataFrame
        The resulting DataFrame after reduce memory usage.
    """    
    
    df = dataframe.copy()
    numerics = ['int16', 'int32', 'int64', 'float16', 'float32', 'float64']
    start_mem = df.memory_usage().sum() / 1024**2
    
    for col in df.columns:
        col_type = df[col].dtypes
        if col_type in numerics:
            c_min = df[col].min()
            c_max = df[col].max()
            if str(col_type)[:3] == 'int':
                if c_min > np.iinfo(np.int8).min and c_max < np.iinfo(np.int8).max:
                    df[col] = df[col].astype(np.int8)
                elif c_min > np.iinfo(np.int16).min and c_max < np.iinfo(np.int16).max:
                    df[col] = df[col].astype(np.int16)
                elif c_min > np.iinfo(np.int32).min and c_max < np.iinfo(np.int32).max:
                    df[col] = df[col].astype(np.int32)
                elif c_min > np.iinfo(np.int64).min and c_max < np.iinfo(np.int64).max:
                    df[col] = df[col].astype(np.int64)
            else:
                if c_min > np.finfo(np.float16).min and c_max < np.finfo(np.float16).max:
                    df[col] = df[col].astype(np.float16)
                elif c_min > np.finfo(np.float32).min and c_max < np.finfo(np.float32).max:
                    df[col] = df[col].astype(np.float32)
                else:
                    df[col] = df[col].astype(np.float64)

    end_mem = df.memory_usage().sum() / 1024**2
    
    if verbose:
        print('Memory usage BEFORE optimization is: {:.2f} MB'.format(start_mem))
        print('Memory usage AFTER optimization is: {:.2f} MB'.format(end_mem))
        print('Decreased by {:.1f}%'.format(100 * (start_mem - end_mem) / start_mem))

    return df