from ._imports import *
from ._pconfig import *
from ._score import *
from ._misc import get_model_version
from ._train import get_train_data
from statsmodels.tsa.seasonal import STL

def monitor_pipeline (runtime):
    '''
    Example
    -------
    cf_type = 'life_inflow_premiums'
    #train_begin = None   
    train_end = None 
    env = 'dev'  
    runtime = create_runtime (cf_type, env, train_end=None)
    tracker = monitor_pipeline (runtime)
    print(tracker)

    Parameters
    ----------
    runtime : dict
        dictionary contains parameters which will be used monitor pipeline.   

    Returns
    ------- 
    dict
        monitor results.

    '''
    df_prev = get_prev_data(runtime)
    df_curr = get_curr_data(runtime) 

    df_prev_compare, df_curr_compare = adjust_data_per_features_used(df_prev, df_curr)    
    data_drift_flag, data_drift_detail = get_data_drift(df_prev_compare, df_curr_compare)    
    outlier_flag, outlier_detail = get_outlier(df_curr, runtime)
    null_flag, missing_months = check_ts_missing_month(df_curr, runtime['train_end'])
    tracker ={
        'Data Drift' :data_drift_flag,
        'Data Drift Detail': data_drift_detail,
        'Outlier Drift': outlier_flag,
        'Outlier Detail': outlier_detail, 
        'Null Drift': null_flag,
        'Missing Months':missing_months,
        'Prev Dataset': convert_ts_to_dict(df_prev),
        'Curr Dataset': convert_ts_to_dict(df_curr)
    }
    return tracker

def convert_series_to_df (df):
    '''change to df if it is series type'''
    df = df.copy()     
    if isinstance(df , pd.Series):
                df  = df.to_frame()
    return df

def get_data_with_ts(df, runtime):
    '''
    add timeseries to the index of existing df    
    '''
    df_ind = df.copy()
    df_ind.index =df_ind.iloc[:, :2].apply(lambda x: str(x[0]) + ' ' + str(x[1]), axis =1).apply(pd.to_datetime).values  #first two column are year month
    df_ind = df_ind.loc[df_ind.index <= runtime['train_end']].iloc[:, 2:] #what if there are mulitple feature iloc[:, 2:]?
    df_ind = convert_series_to_df(df_ind)
    return df_ind

def pre_process_mutual_data(df, runtime):
    df = df.copy( )
    if 'Date' in df.columns:
        df['Date'] = pd.to_datetime(df['Date'])
    else:
        df['Date'] = pd.to_datetime(df['trans_year'].astype(str) + '-' + df['trans_month'] +'-01')
    extra_columns = [ i for i in df.columns if i in ['trans_year', 'trans_month']]
    for i in extra_columns:
        df.drop(i, axis =1, inplace = True)                
    df = df[df['Date']<=pd.to_datetime(runtime['train_end'])].set_index('Date') 
    return df      

def get_prev_data(runtime):
    '''
    based on version to get reference file, make date as index
    '''
    model_version = get_model_version(runtime)
    if not model_version:  #first time run
         df_prev = get_curr_data(runtime)
    else:         
        file_list = s3_list_folder_contents(pconfig[f's3_model_assets_folder_{runtime['env']}'], None, f'^.*cash_flow_model_{runtime['cf_type']}_{model_version}.csv$', runtime['env'])
        latest_file = sorted(file_list)[-1]
        df_prev = s3_download_and_read_csv(latest_file, env = runtime['env'])
        #mutal side
        if runtime['cf_type'] in ['claim_loss', 'claim_subrogation', 'inflow_premiums', 'outflow_premiums', 'inflow_receipts', 'vendor', 'salary'] :            
            df_prev = pre_process_mutual_data(df_prev, runtime)
        else:
            df_prev = get_data_with_ts(df_prev, runtime)
        
    return df_prev

def get_curr_data(runtime):
    type, env, train_begin, train_end= runtime['cf_type'], runtime['env'], runtime['train_begin'], runtime['train_end']
    df_curr, *_ = get_train_data (type, env, train_begin, train_end)
    if runtime['cf_type'] in [ 'inflow_premiums', 'outflow_premiums', 'claim loss', 'claim subrogation','inflow_receipts', 'vendor', 'salary'] :            
            df_curr= pre_process_mutual_data(df_curr, runtime)
    else:
        df_curr = get_data_with_ts(df_curr, runtime)
    return df_curr

def adjust_data_per_features_used(df_prev, df_curr):
    shared_columns =  [ i  for i in df_prev.columns if i in df_curr.columns]
    df_prev_compare = df_prev[shared_columns]
    df_curr_compare = df_curr[shared_columns]
    return df_prev_compare, df_curr_compare 


def chi_square(ref, new, alpha=0.05):
    categories     = set(ref.unique()).union(set(new.unique()))
    ref_proportion = [(ref==c).sum()/len(ref) for c in categories]
    new_proportion = [(new==c).sum()/len(new) for c in categories]
    r = scipy.stats.chisquare(new_proportion, f_exp=ref_proportion)
    return int(r.pvalue < alpha)


def population_stability_index(ref, new, buckets=10, psi_cutoff=0.2):
    '''
    Calculate the PSI (population stability index)
    Adapted from https://github.com/mwburke/population-stability-index/blob/master/psi.py
    '''

    def sub_psi(e_pct, a_pct):
        if a_pct == 0: a_pct = 0.0001
        if e_pct == 0: e_pct = 0.0001
        return (e_pct - a_pct) * np.log(e_pct / a_pct)

    breaks  = np.linspace(np.min(ref), np.max(ref), buckets+1)
    ref_pct = np.histogram(ref, breaks)[0] / len(ref)
    new_pct = np.histogram(new, breaks)[0] / len(new)

    psi_value = np.sum(sub_psi(ref_pct[i], new_pct[i]) for i in range(0, len(ref_pct)))
    return int(psi_value > psi_cutoff)

def get_data_drift (df_prev_compare, df_curr_compare):    
    dict_data_drift = dict()

    if isinstance(df_prev_compare , pd.Series):
        df_prev_compare   = df_prev_compare.to_frame()
    
    if isinstance(df_curr_compare , pd.Series):
        df_curr_compare   = df_curr_compare.to_frame()


    for c in df_curr_compare.columns:
        if pd.api.types.is_numeric_dtype(df_curr_compare[c]):                     
                drift = population_stability_index(df_prev_compare[c], df_curr_compare[c])
        else:
            drift = chi_square(df_prev_compare[c], df_curr_compare[c])
        dict_data_drift[c] = drift
    data_drift_flag = int(any(dict_data_drift.values()))    
    return  data_drift_flag, dict_data_drift

def get_column_residuals_outlier(df_curr_column, runtime):
    '''
    https://towardsdatascience.com/the-ultimate-guide-to-finding-outliers-in-your-time-series-data-part-1-1bf81e09ade4__;!!GMUSBZg!nwXhu_f2jwHgSIK7ElxsOicJKl6YEW0MYG1syVpMxOT1TIFdfHn2G7enUnyaJiZh7KxIeiqm9vGb4Ys$ 
    https://www.geeksforgeeks.org/seasonal-decomposition-of-time-series-by-loess-stl/
    compare new data (last entry) with the previous data residual assume they share the trend, sesonality
    '''  
    df_curr_column.index = pd.to_datetime(df_curr_column.index)
    df_curr_ind = df_curr_column[df_curr_column.index <= pd.to_datetime(runtime['train_end'])]
    # Perform STL decomposition
    stl = STL(df_curr_ind, seasonal=13, robust=False) #this exclude the last incomplete month
    result = stl.fit()
    #get the data up to second last
    residuals =  result.resid
    mean_resid = np.mean(residuals[:-1]) #this spare the last month
    std_resid = np.std(residuals[:-1])
    threshold = 3 * std_resid
    upper_limit = mean_resid + threshold
    lower_limit = mean_resid - threshold
    #compare last data with previous data
    outlier_flag = 0 if residuals[-1] <= upper_limit and residuals[-1] >= lower_limit else 1
    return outlier_flag, [residuals[-1], lower_limit, upper_limit] 

def get_outlier(df_curr, runtime):
    dict_outlier_flag = {}
    date_fields = ['trans_year','trans_month'] #Ignore these
    success_share_fields = ['success_share']
    for c in [i for i in df_curr.columns if i not in date_fields + success_share_fields ]:
        outlier_flag, detail = get_column_residuals_outlier(df_curr[c], runtime)
        dict_outlier_flag[c] = [outlier_flag, detail]
    if runtime['env']== 'dev':  #dev data in policy center s3 table is messy, almost each has outlier
        outlier_flag = dict_outlier_flag[next(iter(dict_outlier_flag))][0]   #this get the y outlier flag only, the data from database in the dictionary 
    else:  #if in prod or qa
        outlier_flag = int(any([v[0] for k, v in dict_outlier_flag.items() if k not in ['earned_premium_this_month']]))   #this metric has larger fluctuation, exclude it
    return outlier_flag,  dict_outlier_flag


def check_ts_missing_month(df, month_end):
    """check time series data if there is missing month"""
    month_begin = df.index.min().strftime("%Y-%m-%d")
    if not month_end:
        month_end = DtFind.first_day_last_month( dt.datetime.today(), dt_format=None).strftime("%Y-%m-%d")
    month_range =  pd.date_range(month_begin, month_end, 
                freq='MS').strftime("%Y-%m-%d").tolist()
    missing_month = [i for i in month_range if pd.to_datetime(i) not in df.index]
    return int(any(missing_month)), missing_month 

def convert_ts_to_dict (df):
    df.index = df.index.astype(str)
    return df.to_dict()