from ._imports import *
from ._train import *
from ._pconfig import *
from ._misc import *
from ._monitor import monitor_pipeline
pd.set_option('display.float_format', lambda x: '%.2f' % x)

'''
cf_type = 'life_inflow_premiums'
#train_begin = None   
train_end = None 
env = 'dev'  
run_time = create_runtime (cf_type, env, train_end=None)
dict_results = training_pipeline( run_time)
df_score = get_forecast(dict_results, run_time)
df_score

'''


def scoring_pipeline (cf_type = 'claim loss', env = 'dev', logger = None, train_end = None):  
    '''
    Example
    -------
    This function will run monitor and training pipeline and use optmized model to forecast next three month
    example:
    scoring_pipeline (cf_type = 'claim loss', env = 'dev', logger = logger, train_end = '2024-10-01')
    add train_end to recreate previous month prediction
    loggile, logger = create_logger()
    scoring_pipeline (cf_type = 'claim loss', env = 'dev', logger = logger, train_end = None)

    Parameters
    ----------
    cf_type : str
        cash flow type.
    env : str
        environment that function is running.
    logger : logging.RootLogger
        keep job log.
    train_end : str
        training data end, by default previous month.

    Returns
    ------- 
    int
        Binary integer for error status.
    '''
    
    error_status = 0  
    try:    
        runtime = create_runtime(cf_type, env, train_end)
        if logger: 
            logger.info(f'''{cf_type} forecasting based on data between {runtime['train_begin']} and {runtime['train_end']},
                         test length is {runtime['test_len']} and forecast length is {runtime['forecast_len']}''')
        tracker = monitor_pipeline(runtime)        
        assert not tracker['Null Drift'], f'Data might be missing in {','.join(tracker['Missing Months'])} '        
        assert not tracker['Outlier Drift'], f'The latest record residue has 3 standard deviation from mean, please make sure data is right'
        dict_results = training_pipeline(runtime) 
        df_score = get_forecast(dict_results, runtime)
        #send to accounting
        if logger:logger.info(df_score.to_string())
        #send_sns(df_score , 'dl')  #TODO change dl to env when dev credential available        
        
        #save model asset
        save_model_assets(dict_results, runtime)
        
        #save monitor 
        save_monitor_mstr(tracker, runtime)
        #save_monitor_dashboard_value(runtime)
        
        #save score to ds folder
        save_cash_flow_score_to_ds(df_score, dict_results, runtime)        
        #export to silver layer
        save_cash_flow_score_to_silver (df_score, runtime)

    except Exception as e:        
        if logger:
            logger.error('Error', exc_info=e)
        error_status = 1
        raise  #pass error to outer try catch
    return error_status



#module variable, in case param cf_type different from folder name 
#life flow is not added, I will use standard name
dict_s3_prediction_folder_mapping = {
    'claim_loss': 'outflow_claims',
    'claim_subrogation': 'inflow_claims',
    'inflow_premiums':'inflow_premiums',
    'outflow_premiums':'outflow_premiums',
    'inflow_receipts': 'inflow_receipts', 
    'vendor':'outflow_vendor' ,
    'salary':'outflow_payroll'
    }

def create_runtime (cf_type, env, train_end=None):
    runtime = {}
    runtime['cf_type'] = cf_type.replace(' ', '_')
    
    if train_end:

        runtime['train_end'] = train_end
    else:         
         runtime['train_end'] = DtFind.first_day_last_month( dt.datetime.today(), dt_format=None).strftime("%Y-%m-%d")
        #runtime['train_end'] = '2024-10-01'
    month_offset = pconfig.get(f'{runtime['cf_type']}_month_offset', None)  #avoid create clutter record in pconfig
    if month_offset:
        runtime['train_begin'] = (dt.datetime.strptime(runtime['train_end'], "%Y-%m-%d") +pd.DateOffset(months=month_offset)).strftime("%Y-%m-%d")   #'2020-07-01'  from last month count back 52 month    
    else:
        runtime['train_begin'] = None  #if none, need to get that from query result    
    runtime['test_len'] = 8
    runtime['forecast_len'] = 3
    runtime['env'] = env    
    runtime['timestamp']      = dt.datetime.now().strftime("%Y%m%d_%H%M%S")
    return runtime


def select_data_range (df_ts, begin, end=None):  
    '''
    select_data_range(df_claim_ts, begin = '2020-04-01')
    '''
    if not begin:
        begin=df_ts.index[0]

    if not end:
        end = df_ts.index[-1]

    return df_ts.loc[begin:end]




############################ test new function 
def score_adjust_model_input(params, train_df, test_df): 
        '''
        this is to adjust prophet model input featuer per parameter value 
        '''           
        params_without_flag = {x:y for x, y in params.items() if x not in ['column_masks', 'columns_included']}
        model = Prophet(**params_without_flag)
        # column_masks = column_mask_comb[params['column_masks']]
        selected_columns = params['columns_included']
        train_df = train_df.loc[:, selected_columns ]
        test_df = 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 score_get_additional_features( future, params, _train_df, _test_df ):                 
#     additional_columns = params['columns_included'][2:] #skip ds and y
#     for c in additional_columns:                 
#         future[c] = pd.concat([_train_df[c], _test_df[c]])
#     return future

def get_forecast_month_features(future, runtime):
    if runtime['cf_type'] in ['outflow_premiums', 'inflow_premiums']:
        df_features = get_additional_features(runtime['env'])
        df_features.iloc[:, 0 ] = pd.to_datetime(df_features.iloc[:, 0 ])
        ts_column_name = df_features.columns.to_list()[0] 
        df_features.set_index(ts_column_name, drop=True, inplace=True)
        # inflow premimium has addition column from accounting 
        if runtime['cf_type'] == 'inflow_premiums':
            df_features0 = prep_premium_inflow_invoice(runtime['env'])   
            df_features0.index = (pd.to_datetime(df_features0[['expected_clear_year', 'expected_clear_month']]
                                                        .rename(columns={'expected_clear_year':'YEAR', 'expected_clear_month':'MONTH'})
                                                        .assign(DAY=1)))
            df_features = df_features0.join(df_features)
    elif runtime['cf_type'] in ['salary', 'life_outflow_payroll' ]:
        df_features = get_payday_datataframe()
        df_features['Date'] = pd.to_datetime(df_features['trans_year'].astype(str) + '-' + df_features['trans_month'] + '-01')
        df_features = df_features[['Date','payday_count','success_share']]
        df_features = df_features.set_index('Date')
    forecast_begin = future['ds'][-runtime['forecast_len']:].min()
    forecast_end = future['ds'][-runtime['forecast_len']:].max()
    filtered_features= select_data_range(df_features, forecast_begin, forecast_end)
    return filtered_features  

def get_forecast(dict_results, runtime):
    '''
    Use trained model to forecast for next three months 
    Parameters
    ----------
    dict_results : dict 
                a dictionary return from training step containing best hyperparm, best score as well data used for training.
    None : dict 
                a parameter set used in this process.
    Returns
    ----------
    None : DataFrame 
                prediction saved in  the dataframe .
    Examples
    ----------
    df_score = get_forecast(dict_results, runtime)
    
    '''
    df = dict_results['dataset inference']
    best_params = dict_results['best params']
    forecast_size = runtime['forecast_len']
    #modify model and df based on the best param, could 
    model, df, _ = score_adjust_model_input(best_params, df, df)  #use 2 df to meeting the orginal function requirement    
    model.fit(df)
    future = model.make_future_dataframe(periods=forecast_size, freq='MS')
    if len(best_params['columns_included'])>2:  #need to include other feature in future to forecast
            
            additional_features = get_forecast_month_features(future, runtime)
            column_name_list = df.iloc[:, 2:].columns.to_list()
            selected_columns =  [i for i in column_name_list if i in dict_results['best params'][ 'columns_included']] 
            for c in selected_columns:                 
                future[c] = pd.concat([df[c], additional_features[c]], ignore_index=True)
            dict_results['additional_feature_used'] = future
    forecast = model.predict(future)
    predictions_tuned = forecast.tail(forecast_size)
    dict_additional_columns ={
    "Type": dict_s3_prediction_folder_mapping.get(runtime['cf_type'], runtime['cf_type']).replace('_', ' '),
    "Timestamp" : str(runtime['timestamp']).replace('_', ' '),
    "Forecast Month": predictions_tuned['ds'].astype(str),
    "Values": predictions_tuned['yhat'],
    "Base Month":str(runtime['train_end']) }   
    predictions_tuned = predictions_tuned.assign(**dict_additional_columns)
    return predictions_tuned[['Timestamp', 'Type', 'Forecast Month', 'Values', 'Base Month']] 

#### save model results
def save_model_assets (dict_results,runtime):
    '''
    this function saves model asset to AWS S3 
    Parameters
    ----------
    dict_results : dict 
                dictionary contains all training meta data.
    runtime : dict 
                dictionary contains all parameter used for training  and scoring.
    Examples
    ----------
    save_model_assets (dict_results,runtime)
    '''
    #first time run, if no version file, create one    
    if not (version:= get_model_version(runtime)):
           archive_model_version (runtime)

    #only update version when end date == last month, to avoid testing update
    if runtime['train_end'] ==DtFind.first_day_last_month( dt.datetime.today(), dt_format=None).strftime("%Y-%m-%d"):
             archive_model_version (runtime)
    
    #handle model with more than 1 feature
    #dict_results_without_flag = {k:v for k, v in dict_results['best params'].items() if k != 'additional_feature_flag'}

    params_without_flag = {x:y for x, y in dict_results['best params'].items() if x not in ['column_masks', 'columns_included']}
        

    #save model performance 
    s3_save_and_upload_file(
            #get best params and best loss
            {k: dict_results[k] for k in ('best params', 'best loss', 'tuning results')},
            s3_file_path = os.path.join(pconfig[f's3_model_assets_folder_{runtime['env']}'], f'cash_flow_model_{runtime["cf_type"]}_{runtime["timestamp"]}.json'),
            overwrite='overwrite',
            env=runtime['env'],
        )

    #save model itself
    s3_save_and_upload_file(
                
        Prophet(**params_without_flag),
        s3_file_path = os.path.join(pconfig[f's3_model_assets_folder_{runtime['env']}'], f'cash_flow_model_{runtime["cf_type"]}_{runtime["timestamp"]}.pkl'),
        overwrite='overwrite',
        env=runtime['env'],
    )

    #save model ref dataset
    s3_save_and_upload_file(
        dict_results['original dataset'],
        s3_file_path = os.path.join(pconfig[f's3_model_assets_folder_{runtime['env']}'], f'cash_flow_model_{runtime["cf_type"]}_{runtime["timestamp"]}.csv'),
        overwrite='overwrite',
        env=runtime['env'],
    )

#save score data
def save_cash_flow_score_to_ds(df_score, dict_results, runtime):
    '''
    Save inference dataset, additional features and prediction result in s3 for future audition 
    Parameters
    ----------
    df_score : DataFrame 
                dataFrame containing prediction.
    dict_results : dict 
                dictionary cotaining all training metadata.
    runtime : dict 
                dictionary contain  training params.
    Examples
    ----------
    save_cash_flow_score_to_ds(df_score, dict_results, runtime)
    '''
    s3_save_and_upload_file(
        dict_results['dataset inference'],
        s3_file_path = os.path.join(pconfig[f's3_model_output_eds_folder_{runtime['env']}'], f'cash_flow_model_{runtime["cf_type"]}_inference_data_{runtime["timestamp"]}.csv'),
        overwrite='overwrite',
        env=runtime['env'],
    )

    if 'additional_feature_used' in dict_results:
        s3_save_and_upload_file(
        dict_results['additional_feature_used'],
        s3_file_path = os.path.join(pconfig[f's3_model_output_eds_folder_{runtime['env']}'], f'cash_flow_model_{runtime["cf_type"]}_additional_features_{runtime["timestamp"]}.csv'),
        overwrite='overwrite',
        env=runtime['env'],
    )

    s3_save_and_upload_file(
        df_score,
        s3_file_path = os.path.join(pconfig[f's3_model_output_mstr_folder_{runtime['env']}'], f'score/cash_flow_model_{runtime["cf_type"]}_scoring_{runtime["timestamp"]}.csv'),
        overwrite='overwrite',
        env=runtime['env'],
    )


# save prediction to silver
def save_cash_flow_score_to_silver (df_score, runtime): 
    '''
    Save score to cashflow_forecast/forecast_stage/ 
    Parameters
    ----------
    df_score : DataFrame 
                dataframe contains score results.
    runtime : dict 
                dictionary contains training params.
    Examples
    ----------
    save_cash_flow_score_to_silver (df_score, runtime)
    '''   
    accounting_type = dict_s3_prediction_folder_mapping.get(runtime['cf_type'], runtime['cf_type']) #if cannot find, use type value
    run_date = runtime['timestamp'].split('_')[0]
    type = runtime['cf_type']
    base_month = runtime['train_end'][:7]
    file_name = f'{accounting_type}_{base_month}_{run_date}.parquet'
    local_file_path = config['local']['local_folder']
    local_file_path = os.path.join(local_file_path, file_name)
    s3_file_path = f'cashflow_forecast/forecast_stage/{accounting_type}/{file_name}'
    df_score.to_parquet(local_file_path)
    s3_upload_file_to_silver(local_file_path, s3_file_path, overwrite=True, env=runtime['env'])    
    os.remove(local_file_path)
    return  

def save_monitor_mstr(data, runtime):
    '''
    Save model data monitoring result to s3 folder 
    Parameters
    ----------
    data : dict 
                dictionary that conatins satistics comparing current data vs previous run.
    runtime : dict 
                dictionary contains training params info.
    Examples
    ----------
    save_monitor_mstr(tracker, runtime)
    '''
    if not runtime["train_begin"]:        #some model don't have fixed start date 
         keyname = next(iter(data['Curr Dataset']))
         ts = list(data['Curr Dataset'][keyname].keys())
         runtime["train_begin"] = min(ts)
         
    s3_save_and_upload_json(
        data,
        s3_file_path = os.path.join(
            pconfig[f's3_model_output_mstr_folder_{runtime["env"]}'],
            'monitor/',
            f'{runtime['cf_type']}_monitor_{runtime["train_begin"]}_{runtime["train_end"]}_{runtime["timestamp"]}.json'
        ),
        env=runtime['env'],
    )    
    return



#### these function are  for model data monitor by wcao 
def get_latest_monitor_file(env = 'dev'):
    monitor_files = s3_list_folder_contents(
                os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'monitor/'),
                pattern=r'^.*\.json$',
                env=env,
            )
    df_monitor_files = pd.DataFrame(columns=['file_path'])
    df_monitor_files['file_path'] = monitor_files
    df_monitor_files['file_name'] = df_monitor_files['file_path'].str.split('/').str[-1]
    df_monitor_files['cash_flow_type'] = df_monitor_files['file_name'].str.split('_').str[:-5].apply(lambda x: "_".join(x))
    df_monitor_files['training_end_month'] = df_monitor_files['file_name'].str.split('_').str[-3]
    df_monitor_files['training_date_time'] = df_monitor_files['file_name'].str.split('_').str[-2:].apply(lambda x: "_".join(x)).str.strip('.json')
    #drop duplicate
    df_monitor_files2 = df_monitor_files.sort_values(by=['cash_flow_type','training_end_month', 'training_date_time'])
    df_monitor_files2 = df_monitor_files2.drop_duplicates(subset=['cash_flow_type','training_end_month'], keep='last', ignore_index=True ) 
    return df_monitor_files2

def extract_monitor_value(df_monitor, env):
    df_monitor_extract = pd.DataFrame(columns = ['type', 'train_end',  'run_timestamp', 'data_drift',
                                              'outlier_drift', 'null_drift', 'residuals', 'low_limit', 'high_limit', 'mean_limit'])
    for i, row in df_monitor.iterrows():
        dict_row = s3_download_and_read_json(row['file_path'],None, None, env)
        if type(dict_row['Outlier Detail']) == list:  #this is old format
            Outlier_Detail= dict_row['Outlier Detail']
        elif type(dict_row['Outlier Detail']) == dict:
            Outlier_Detail=  next(iter(dict_row['Outlier Detail'].values()), None)[1]  #get first value
        residuals, low_limit, high_limit,  = Outlier_Detail
        mean_limit = (low_limit + high_limit)/2
        df_monitor_extract.loc[len(df_monitor_extract)] =[row['cash_flow_type'],row['training_end_month'], row['training_date_time'], dict_row['Data Drift'],
                                                        dict_row['Outlier Drift'], dict_row['Null Drift'],residuals, low_limit, high_limit, mean_limit]
    dict_type_map = {
        'claim_loss': 'outflow claims', 'claim_subrogation':'inflow claims',  
        'salary':'outflow payroll',  'vendor':'outflow vendor' 
    }
    df_monitor_extract['type'] =  df_monitor_extract['type'].apply(lambda x:dict_type_map.get(x, x.replace('_', " ")))
    return df_monitor_extract 

def save_monitor_dashboard_value (env):
    '''
    collect latest monitor file and extract the value for monitor dashboard use 
    Parameters
    ----------
    env : str 
                environment the function is running.
    Examples
    ----------
    save_monitor_dashboard_value('dev')

    '''  
    df_monitor = get_latest_monitor_file(env)
    df_monitor_extract = extract_monitor_value(df_monitor, env)
    s3_save_and_upload_csv(
        df_monitor_extract,
        s3_file_path = os.path.join(
            pconfig[f's3_model_output_mstr_folder_{env}'] + r'dashboard/',           
            f'data_monitor.csv'
        ),
        overwrite=True,
        env=env
    )
    return

####these function is for combine prediction into one file by wcao

def get_latest_prediction_file(env = 'dev'):
    file_list = s3_get_prediction_file_in_silver (s3_folder_path = r"cashflow_forecast/forecast_stage/", s3_bucket = f'amica-aws-eds-{env}-dl-silver', pattern = r'^.*\.parquet$', env=env)
    df_prediction_file = pd.DataFrame()
    df_prediction_file['file_path'] = file_list
    df_prediction_file['timestamp'] =  df_prediction_file['file_path'].str.split('/').str[-1].str.split('_').str[-1]
    df_prediction_file['base_month'] = df_prediction_file['file_path'].str.split('/').str[-1].str.split('_').str[-2]
    df_prediction_file['category'] = df_prediction_file['file_path'].str.split('/').str[-1].str.split('_').str[:-2].apply(lambda x:'_'.join(x))
    #drop duplicate
    df_prediction_file2 =  df_prediction_file.sort_values(by=['category','base_month', 'timestamp'])
    df_prediction_file2 = df_prediction_file2.drop_duplicates(subset=['category','base_month'], keep='last', ignore_index=True ) 
    return df_prediction_file2 

def extract_prediction(df_prediction_file, env = 'dev' ): 
    #some old file contain old name passing by param, this mapping fix, in new env no this issue
    dict_category_mapping  = { x.replace('_', " "):y.replace('_', " ")  for x, y in dict_s3_prediction_folder_mapping.items()  }
       
    df_prediction =  pd.DataFrame(columns=['Timestamp', 'Type', 'Forecast Month', 'Values', 'Base Month'])
    for i, row in df_prediction_file.iterrows():
        df_file = s3_download_and_read_file(row['file_path'], local_file_path=None, s3_bucket=f'amica-aws-eds-{env}-dl-silver', read_file_kwargs={}, env=env)
        df_file['Type']= df_file['Type'].apply(lambda x: dict_category_mapping.get(x, x))  #this fix old name to new name
        df_prediction = pd.concat([df_prediction, df_file], ignore_index= True)
        #remove local file after load
        s3_remove_temp_local_download (row['file_path'])
    df_prediction = df_prediction.sort_values(by=['Type', 'Base Month', 'Timestamp'], ignore_index= True)
    return df_prediction

def save_latest_prediction (env): 
    '''
    collect latest model prediction from different cash flow type to prediction_combined.csv 
    Parameters
    ----------
    env : str 
                environment the function is running.
    Examples
    ----------
    save_latest_prediction('dev')
    ''' 
    df_prediction_file = get_latest_prediction_file(env = env)
    df_prediction_extract = extract_prediction(df_prediction_file, env = env )
    timestamp =  df_prediction_extract['Timestamp'].max().replace(' ', '_')
    local_file_path = os.path.join(config['local']['local_folder'], 'prediction_combined.csv' )    
    df_prediction_extract.to_csv(local_file_path, index = False)
    s3_file_path = r'cashflow_forecast/forecast_stage/forecast_combined/' + 'prediction_combined_' + timestamp + '.csv'
    s3_upload_file_to_silver(local_file_path, s3_file_path, overwrite=True, env=env)
    os.remove(local_file_path)    
    return 

