from typing import Dict, Tuple, Union, List
import pandas as pd

from .utils import generate_time_series, missing_data_imputation
from ..utility.tele_logger import logger

class ProphetPreprocess():
    def __init__(self, hierarchy: Dict[str, str], time_granularity: str, 
                 date_col: str, missing_data_strategy: Union[str, int, Dict[str, Union[str, int]]]):
        """
        Init of the preprocessing class for the prophet model

        Parameters
        ----------
        hierarchy : Dict[str, str]
            The hierarchy dictionary from the configuration file 
        time_granularity : str
            The time granularity to aggregate <date_col>
        date_col : str
            The date col to aggregate with <time_granularity>
        missing_data_strategy : Union[str, int, Dict[str, Union[str, int]]]
            The missing data strategy specified in the configuration file
        """
        self.hierarchy = hierarchy
        self.time_granularity = time_granularity
        self.date_col = date_col
        self.missing_data_strategy = missing_data_strategy

    def fit(self):
        """
        Fit method of ProphetPreprocess

        Returns
        -------
        Self
        """
        return self
    
    def transform(self, X: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame, Dict[str, List[str]]]:
        """
        Transform method which applies the pipeline of transformations for TC

        Parameters
        ----------
        X : pd.DataFrame
            The dataframe to transform

        Returns
        -------
        Tuple[pd.DataFrame, pd.DataFrame, Dict[str, List[str]]]
            Tuple containing (X, S, tags)
            - X: The processed dataframe after preprocessing
            - S: Dataframe of binary values resulting from the aggregation (incidence matrix of undirected graph)
            - tags: Dictionary containing the information about the hierarchies
        """
        logger.info('Preparing Dataset for Prophet Modelling.')
        X, S, tags = generate_time_series(df = X, date_col = self.date_col, 
                                          time_granularity = self.time_granularity, hierarchy = self.hierarchy)
        X = missing_data_imputation(missing_data_strategy = self.missing_data_strategy, df = X)

        return X, S, tags
