__all__ = ['get_best_prophet_param_both2']

from abc import ABC, abstractmethod
from cffds.tools._imports import *
from scipy.stats import uniform


class PerformanceTuneBase(ABC):


    def __init__(self, train_df, test_df, test_size, cf_type):
        self.train_df = train_df
        self.test_df = test_df        
        self.test_size = test_size
        self.cf_type = cf_type        
        self.column_mask_comb = get_column_boolean_combination(train_df, cf_type)    #at instance level
      

    @abstractmethod
    def objective(self):
        """Abstract method to create objective function."""
        pass
    

    def adjust_model_input(self, params): 
        '''
        this is to adjust prophet model input featuer per parameter value 
        '''           
        params_without_flag = {x:y for x, y in params.items() if x!='column_masks'}
        model = Prophet(**params_without_flag)
        column_masks = self.column_mask_comb[params['column_masks']]
        selected_columns = [ x  for x, y in zip(self.train_df.columns, column_masks) if y] 
        train_df = self.train_df.loc[:, selected_columns ]
        test_df = self.test_df.loc[:, selected_columns ]
        #add regressor
        additional_columns = selected_columns[2:]
        if additional_columns:
            for i in additional_columns:
                model.add_regressor(i)
        return model, train_df, test_df
    
    def get_additional_features(self, future, params, _train_df, _test_df ):
        column_masks =  self.column_mask_comb[params['column_masks']]               
        selected_columns = [ x for x, y in zip(_train_df.columns, column_masks) if y][2:] #skip ds and y
        for c in selected_columns:                 
            future[c] = pd.concat([_train_df[c], _test_df[c]])
        return future

class MangoTuning(PerformanceTuneBase):  
    
    conf_Dict = dict()
    conf_Dict['initial_random'] = 10
    conf_Dict['num_iteration'] = 100

    def __init__ (self, train_df, test_df, test_size, cf_type):
        super().__init__(train_df, test_df, test_size, cf_type)
        self.param_space = self.get_param_space()                   #using polymorphism for each child class

    def objective(self, args_list): 
        
        params_evaluated = []
        res = []
        
        for params in args_list:
            try:
                
                model, _train_df, _test_df = self.adjust_model_input(params)                
                model.fit(_train_df)                
                future = model.make_future_dataframe(periods=self.test_size, freq='MS') 
                # column_masks =  self.column_mask_comb[params['column_masks']]               
                # selected_columns = [ x for x, y in zip(_train_df.columns, column_masks) if y][2:] #skip ds and y
                # for c in selected_columns:                 
                #     future[c] = pd.concat([_train_df[c], _test_df[c]])
                future = self.get_additional_features(future, params, self.train_df, self.test_df)
                forecast = model.predict(future)
                predictions_tuned = forecast.tail(self.test_size)
                error = mape(_test_df['y'].values, predictions_tuned['yhat'].values)   
                
                params_evaluated.append(params)
                res.append(error)
            except Exception as e:            
                params_evaluated.append(params)
                print(e)
                res.append(100.0)# Giving high loss for exceptions regions of spaces
            
        return params_evaluated, res 
    
    def get_param_space(self):
         
         param_space = dict(growth = ['linear',  'flat'],  #remove 'logistic', need cap in data set
                   n_changepoints  = range(0, 55, 5),
                   changepoint_range  = uniform(0.5, 0.5),
                   yearly_seasonality = [True, False],
                   weekly_seasonality = [True, False],
                   daily_seasonality = [True, False],
                   seasonality_mode = ['additive', 'multiplicative'],
                   seasonality_prior_scale=uniform(5.0, 15.0),
                   changepoint_prior_scale=uniform(0.0, 0.1),
                   interval_width = uniform(0.2, 0.8),
                   uncertainty_samples = [500, 1000, 1500, 2000],
                   column_masks = range(0, len(self.column_mask_comb))   #direct use 
                   )        
         return param_space 
    
class OptunaTuning(PerformanceTuneBase):

    def __init__ (self, train_df, test_df, test_size, cf_type):
        super().__init__(train_df, test_df, test_size, cf_type)  

    def objective(self, trial): 
        params = {

            "growth":trial.suggest_categorical("growth", ['linear',  'flat']),  #remove 'logistic',  it will need cap
            "n_changepoints":trial.suggest_int("n_changepoints", 0, 55, 5),
            "changepoint_range":trial.suggest_float("changepoint_range", 0.5, 0.5),
            'yearly_seasonality': trial.suggest_categorical('yearly_seasonality',[True, False]),
            'weekly_seasonality': trial.suggest_categorical('weekly_seasonality',[True, False]),  
            'daily_seasonality': trial.suggest_categorical('daily_seasonality',[True, False]), 
            "seasonality_mode":trial.suggest_categorical("seasonality_mode", ['additive', 'multiplicative']),
            "seasonality_prior_scale":trial.suggest_float("seasonality_prior_scale", 5, 15),
            "changepoint_prior_scale":trial.suggest_float("changepoint_prior_scale", 0, 0.1),
            "interval_width":trial.suggest_float("interval_width", 0.2, 0.8),
            'uncertainty_samples': trial.suggest_categorical('uncertainty_samples',[500, 1000, 1500, 2000]), 
            'column_masks': trial.suggest_int("column_masks", 0, len(self.column_mask_comb)-1),  #optuna is inclusive for upper boundry, use -1
            
        }  
        
        try:
            model, _train_df, _test_df = self.adjust_model_input(params)        
            model.fit(_train_df)
            future = model.make_future_dataframe(periods=self.test_size, freq='MS')                 
            future = self.get_additional_features(future, params, self.train_df, self.test_df)
            forecast = model.predict(future)            
            predictions_tuned = forecast.tail(self.test_size)
            error = mape(_test_df['y'].values, predictions_tuned['yhat'].values)
            return error 
        except Exception as e:  
             print(e) 



def generate_boolean_combinations( n):
    """
    Generates a list of all possible boolean combinations for 'n' elements.

    Args:
        n: The number of boolean elements.

    Returns:
        A list of lists, where each inner list represents a unique boolean combination.
    """
    if n < 0:
        raise ValueError("Number of elements cannot be negative.")

    combinations = []
    for i in range(2**n):
        # Convert the integer 'i' to its binary representation
        # and pad with leading zeros to ensure 'n' digits.
        binary_string = bin(i)[2:].zfill(n)
        
        # Convert each binary digit to a boolean value
        boolean_list = [bool(int(bit)) for bit in binary_string]
        combinations.append(boolean_list)
    return combinations

def get_column_boolean_combination(df, cf_type):
    '''
    the prophet data format will be  [ds, y, additional columns]
    assume the first 2 columns are needed to predict, then from 3rd is optional
    '''        
    if df.shape[1] == 2:  #ds, y
        return [(True,True)]
    
    manditary_column_numb = 2
    if cf_type in ['salary', 'life_outflow_payroll',]:   # 'inflow_premiums' remove this for now
        manditary_column_numb = 3

    if df.shape[1] <2:
        raise ValueError("Dataframe must have at least two columns: 'ds' and 'y'.")
    elif df.shape[1]-manditary_column_numb <0:
        raise ValueError("Dataframe has less columns than required mandatory columns.")
    elif df.shape[1]-manditary_column_numb ==0:
        return [ [True]*df.shape[1] ]
    else:
        column_mask = generate_boolean_combinations(df.shape[1]-manditary_column_numb) #create combination except the first two columns
        
        #add first one in, need to convert between tuple and list    
        for i in  column_mask:
            i[0:0] = [True]*manditary_column_numb  #insert first two column mask
        #column_mask2 = [ tuple(i) for i in column_mask]   #remove tuple need 
        return column_mask

def mape(y_true, y_pred):
            return round(np.mean(np.abs((y_true - y_pred) / y_true)) * 100, 2)

def get_best_prophet_param_mango2 (train_df, test_df,  test_size, cf_type):
    mt = MangoTuning(train_df, test_df, test_size,cf_type)
    tuner = mango.Tuner(mt.param_space, mt.objective, mt.conf_Dict)
    results = tuner.minimize()
    best_params = results['best_params']
    best_params['column_masks'] =mt.column_mask_comb[best_params['column_masks']] #convert to boolean mask
    best_loss = results['best_objective']
    return best_params, best_loss  

def get_best_prophet_param_optuna2 (train_df, test_df,  test_size, cf_type):
    ot = OptunaTuning(train_df, test_df, test_size,cf_type)
    study = optuna.create_study(direction="minimize") 
    study.optimize(ot.objective, n_trials=100) # run the objective function 100 times
    optuna_best_params = study.best_params.copy()   #cannot assign value to study.best_params
    optuna_best_params['column_masks'] = ot.column_mask_comb[optuna_best_params['column_masks']] #convert to boolean mask       
    return optuna_best_params,study.best_value

def get_best_prophet_param_both2(train_df, test_df, test_size, cf_type):
    best_params_mango, best_loss_mango = get_best_prophet_param_mango2 (train_df, test_df,test_size, cf_type)
    best_params_optuna, best_loss_optuna = get_best_prophet_param_optuna2(train_df, test_df,test_size, cf_type)
    best_params = best_params_mango if best_loss_mango < best_loss_optuna else best_params_optuna
    best_params['columns_included'] = [ x  for x, y in zip(train_df.columns, best_params['column_masks'] ) if y]       
    best_loss = best_loss_mango if best_loss_mango < best_loss_optuna else best_loss_optuna  
    winner = 'mango' if best_loss_mango < best_loss_optuna else 'optuna' 
    results = {'mango':[best_loss_mango, best_params_mango], 'optuna':[best_loss_optuna,best_params_optuna]} 
    return best_params, best_loss, winner, results