from typing import Any
import itertools

import holidays
import numpy as np
import pandas as pd


def datepart(df: pd.Series) -> pd.DataFrame:
    """
    Extracts various date/time parts from a pandas Series of datetime values,
    and returns them as a new DataFrame.

    Parameters
    ----------
    df : pd.Series
        A pandas Series of datetime values to extract date/time parts from.

    Returns
    -------
    pd.DataFrame
        A new pandas DataFrame containing the extracted date/time parts,
        with the same index as the input Series.

    Examples
    --------
    >>> dates = pd.Series(['2022-01-01 12:00:00', '2022-01-02 18:00:00'])
    >>> datepart(dates)
                        dayofmonth  dayofweek  dayofyear  hour  year  month  quarter  holiday  weekend
    2022-01-01 12:00:00           1          5          1    12  2022      1        1     True     True
    2022-01-02 18:00:00           2          6          2    18  2022      1        1    False     True
    """
    it_holidays = holidays.country_holidays('IT')
    df_ds = pd.to_datetime(df.values)
    new_df = pd.DataFrame()
    new_df['dayofmonth'] = df_ds.day
    new_df['dayofweek'] = df_ds.dayofweek
    new_df['dayofyear'] = df_ds.dayofyear
    # new_df['hour'] = df_ds.hour
    new_df['year'] = df_ds.year
    new_df['month'] = df_ds.month
    new_df['quarter'] = df_ds.quarter
    new_df['holiday'] = df_ds.map(lambda x: x in it_holidays)
    new_df['weekend'] = df_ds.dayofweek.map(lambda x: x in [5, 6])
    
    new_df.index = df.values

    return new_df

def get_hierarchy_column(df: pd.DataFrame, data: str, values: str, hier: str, n_hier: int = None, fill_value: int = 0,
                         sons: bool = False, father: bool = False,type_: str=""):
        
    pivot = pd.pivot_table(df.fillna(0),values=values,index=data,columns=hier,fill_value=fill_value)
    pivot.columns = [type_ + "_" + i for i in pivot.columns]
    n_hiers = np.array([n_hier-1,n_hier,n_hier+1])[[father,True,sons]]
    
    if n_hier:
        col = [i for i in pivot.columns if len(i.split("/")) in n_hiers]
        pivot = pivot.loc[:,col]
        
    cols = [list(filter(lambda x : i in x,pivot.columns)) 
            for i in np.unique(np.array([i.rsplit("/",1)[0] for i in pivot.columns if len(i.split("/")) == n_hiers[-1]]))]
    
    return [pivot.loc[:,i] for i in cols]

def divide_by_father(df: pd.DataFrame):
    """
    Divides all columns in a pandas DataFrame by a "father" column,
    and returns the result as a new DataFrame. The father column is
    the one with the least number of "/"

    Parameters
    ----------
    df : pd.DataFrame
        A pandas DataFrame containing the columns to be divided,
        as well as a "father" column to divide by.

    Returns
    -------
    pd.DataFrame
        A new pandas DataFrame containing the resulting values of the division.
        The "father" column is included in the new DataFrame.

    Examples
    --------
    >>> data = {'A/B': [1, 2, 3], 'A/C': [4, 5, 6], 'A': [7, 8, 9]}
    >>> df = pd.DataFrame(data)
    >>> divide_by_father(df)
              A/B       A/C    A
    0  0.142857  0.571429  7
    1  0.250000  0.625000  8
    2  0.333333  0.666667  9
    """
    father_index = np.argmin([i.count("/") for i in df.columns])
    father = df.columns[father_index]
    
    df_div = df.div(df[father].values,axis=0)
    df_div[father] = df[father]
    
    df_div.fillna(1/(len(df.columns)-1),inplace=True)
    
    return df_div

def get_lagged_column(df: pd.DataFrame, lags: list[int], columns: list[str] = None):
    """
    Computes lagged versions of one or more columns in a pandas DataFrame,
    and returns them as a new DataFrame.

    Parameters
    ----------
    df : pd.DataFrame
        A pandas DataFrame containing the original columns to be lagged.
    lags : list[int]
        A list of integer values indicating the number of lags to compute.
    columns : list[str], optional
        The name of the columns to be lagged. If not specified, all columns in
        the input DataFrame will be lagged.

    Returns
    -------
    pd.DataFrame
        A new pandas DataFrame containing the lagged versions of the specified column(s).

    Examples
    --------
    >>> data = {'A': [1, 2, 3, 4], 'B': [5, 6, 7, 8]}
    >>> df = pd.DataFrame(data)
    >>> get_lagged_column(df, [1, 2], 'A')
       Lag1_A  Lag2_A
    0     NaN     NaN
    1     1.0     NaN
    2     2.0     1.0
    3     3.0     2.0
    """
    if columns:
        cols = [columns]
    else:
        cols = df.columns
        
    col_lagged = [df[cols].shift(lag) for lag in lags]
    col_names = ["_".join(i) for i in itertools.product(["Lag"+str(i) for i in lags],cols)]
    
    df_lag = pd.concat(col_lagged,axis=1)
    df_lag.columns = col_names
    
    return df_lag

def transform_data(
    real: pd.DataFrame, pro: pd.DataFrame, params: dict[str, Any]
) -> list[pd.DataFrame]:
    PP = PreprocessBoosting(**params)
    train_data = PP.transform(
        real, pro, "ds", "Timestamp", "y", "Pred_Mean", "unique_id", "Id_Pred"
    )
    return train_data

class PreprocessBoosting():
    def __init__(self, datepart: bool, prophet: bool, target: bool, conf_inter: bool, lags: list = None, n_hier: int = None,
                 father_pro: bool = False, sons_pro: bool = False, father_target: bool = False, 
                 sons_target: bool = False, divide_target: bool = False):
        
        """preprocessing class to obtain dataset input for boosting models
        
        Parameters
        ----------
        datepart : bool,
            if True adds date information to dataset
        prophet : bool
            if True adds Prophet predictions to dataset
        target : bool,
            if True adds Target value to dataset
        conf_inter : bool,
            if True adds prophet confidence intervals to dataset
        lags : list
            if True lags the prophet columns and adds them to the dataset
        n_hier : int
            level of hierarchy considered, if None considers all levels
        father : bool
            if True adds father of the hierarchy level (selected by n_hier) to dataset
        sons : bool
            if True adds all sons of the hierarchy level (selected by n_hier) to dataset
        divide_target : bool
            if True divide all target sons by target father
        """
        self.datepart = datepart
        self.prophet = prophet
        self.target = target
        self.conf_inter = conf_inter
        self.lags = lags
        self.n_hier = n_hier
        self.father_pro = father_pro
        self.sons_pro = sons_pro
        self.father_target = father_target
        self.sons_target = sons_target
        self.divide_target = divide_target
        
    def transform(self, df_real: pd.DataFrame, df_pro: pd.DataFrame, date_real: str, date_pro: str, values_real: str, 
                  values_pro: str, hier_real: str, hier_pro: str) -> pd.DataFrame:
        """Transforms a DataFrame
        
        Parameters
        ----------
        df_real : pd.DataFrame
            dataframe of target values
        df_pro : pd.DataFrame
            dataframe of Prophet predictions
        date_real : str
            name of date column (target dataframe)
        date_pro : str
            name of date column (Prophet dataframe)
        values_real : str
            name of value (y_true) column (target dataframe)
        values_pro : str
            name of value prediction (y_hat) column (Prophet dataframe)
        hier_real : str
            name of hierarchy id column (target dataframe)
        hier_pro : str
            name of hierarchy id column (Prophet dataframe)
        
        Returns
        -------
        list[pd.DataFrame]
            list of dataframes with extracted features for boosting models
        """
        
        dfs = []
        
        if self.target:
            target_list = get_hierarchy_column(df_real,date_real,values_real,hier_real,self.n_hier,sons=self.sons_target,
                                               father=self.father_target,type_="Target")
            
            if self.divide_target:
                target_list = [divide_by_father(target) for target in target_list]
                
            dfs.append(target_list)
            
        if self.prophet:
            pro_list = get_hierarchy_column(df_pro,date_pro,values_pro,hier_pro,self.n_hier,sons=self.sons_pro,
                                            father=self.father_pro,type_="Prophet")
            
            dfs.append(pro_list)
            
            if self.lags:
                dfs.append([get_lagged_column(pro,self.lags) for pro in pro_list])
                
        if self.conf_inter:
            dfs.append(get_hierarchy_column(df_pro,date_pro,"Pi_Lower_95",hier_pro,self.n_hier,sons=self.sons_pro,
                                            father=self.father_pro,type_="Pi_Lower"))
            dfs.append(get_hierarchy_column(df_pro,date_pro,"Pi_Upper_95",hier_pro,self.n_hier,sons=self.sons_pro,
                                            father=self.father_pro,type_="Pi_Upper"))
            
        dfs = [pd.concat([i[j] for i in dfs],axis=1) for j in range(0,len(dfs[0]))]
        
            
        if self.datepart:
            timestamp = pd.Series(df_real[date_real].unique()).sort_values()
            date_col = datepart(timestamp)
            dfs = [pd.concat([i,date_col],axis=1) for i in dfs]

        return dfs
    
    @staticmethod
    def get_X_y(df: pd.DataFrame, n_hier: int) -> tuple[pd.DataFrame, pd.DataFrame]:
        target_cols = [c for c in df.columns if c.count("/") == n_hier - 1 and "Target" in c]
        feature_cols = [c for c in df.columns if "Target" not in c]

        X_train = df[feature_cols]
        y_train = df[target_cols]
        return X_train, y_train
    