__all__ = ['monitor_pipeline', 'update_microstrategy', 'get_incremental_update_param']

from ._imports import *
from ._pconfig import *
from ._misc import *
from ._train import get_run_months

def monitor_pipeline(df, model_config, env='dev', logger=None):
    """
    Monitor pipeline that detects target and feature drifting.

    Parameters
    ----------
    df : pd.DataFrame
        Scoring dataset.
    model_config : dict
        Dictionary that stores the model's metadata.
    env : str
        The run environment, 'dev' or 'prod'.
    logger : logging.RootLogger
        Logger for information, warnings, and errors.

    Returns
    -------
    dict
        Dictionary of model monitoring results.
    """

    model_version  = get_model_version(env)
    s3_file        = os.path.join(pconfig[f's3_model_assets_folder_{env}'], f'cancel_v{model_version}.csv')
    kwargs         = {'na_values': pd._libs.parsers.STR_NA_VALUES - {'N/A'}, 'keep_default_na': False}
    reference      = s3_download_and_read_csv(s3_file, read_csv_kwargs=kwargs, env=env)

    concept_drift  = target_drift(df, reference)
    data_drift     = feature_drift(df, reference, model_config)
    null_drift     = feature_null_drift(df, reference, model_config)
    unseen_drift   = feature_unseen_category(df, reference, model_config)

    trackers = {
        'records':           len(df),
        'distribution': {
            'concept_drift': concept_drift,
            'data_drift':    data_drift,
        },
        'null_drift':        null_drift,
        'unseen_drift':      unseen_drift,
    }
    alert = alert_drift(trackers, model_config)
    trackers['alert'] = alert
    return trackers

def target_drift(df, reference):
    drift_preds = chi_square(reference['PREDICTION'], df['PREDICTION'])
    drift_probs = population_stability_index(reference['PROBABILITY'], df['PROBABILITY'])
    return {'PREDICTION': drift_preds, 'PROBABILITY': drift_probs}

def feature_drift(df, reference, model_config):
    df, reference = df.dropna(), reference.dropna()
    drift_feats = {}
    for f in model_config['features']:
        if pd.api.types.is_numeric_dtype(reference[f]):
            drift = population_stability_index(reference[f], df[f])
        else:
            drift = chi_square(reference[f], df[f])
        drift_feats[f] = drift
    return drift_feats

def feature_null_drift(df, reference, model_config):
    null_drift = []
    for f in model_config['features']:
        drift = chi_square(reference[f].isna(), df[f].isna())
        if drift: null_drift.append(f)
    return null_drift

def feature_unseen_category(df, reference, model_config):
    unseen_feats = {}
    for f in model_config['features']:
        if not pd.api.types.is_numeric_dtype(reference[f]):
            unseen = set(df[f].unique()) - set(reference[f].unique())
            if unseen:
                unseen_feats[f] = unseen
    return list(unseen_feats)

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 alert_drift(trackers, model_config, k=5):
    alert = False
    if 'error' in trackers:
        if trackers['error'] == 1:
            return False
    if sum(trackers['distribution']['concept_drift'].values()) > 0:
        return True

    model_config['importances'] = {
       k:v for (k,v) in model_config['importances'].items() if k not in ("STATE", "STATE_TERRITORY", "SOURCE_OF_BUSINESS")
    }
    if 'importances' in model_config:
        importances_v = model_config['importances']
        importances_f = sorted(importances_v.items(), key=lambda x: x[1], reverse=True)
        for f in importances_f[:k]:
            if trackers['distribution']['data_drift'][f[0]] > 0:
                alert = True
    return alert

def update_microstrategy(env='prod', logger=None, **incremental_kwargs):
    '''
    currently, we run this function after score pipeline evaluation pipeline
    all dataset will be updated although evaluatin pipeline only impact eval and impact
    when score month does not exist in eval and impact in incremental mode, it will create empty df. merge result will be the data from s3
    if want to update all record
    update_microstrategy(env, logger)
    if just want to increment load what have been changed
    incremental_kwargs = get_incremental_update_param(eval_month, eval_month) #begin and end to update a few month
    update_microstrategy(env=env, logger=logger, **incremental_kwargs)
    '''
    error_status = 0
    try:        
        df = pd.DataFrame()
        eds_file_list = [i for i in s3_list_folder_contents(pconfig[f's3_model_output_eds_folder_{env}'], env=env) if ".csv" in i]
        if incremental_kwargs:
            eds_file_list = [i for i in eds_file_list if (i.split('/')[-1].split('_')[1] in incremental_kwargs['target_month_first_day']) ]  #
        for f in eds_file_list:        
            t = s3_download_and_read_csv(f, env=env)
            t['filename'] = f
            t['runstart'] = f.split('_')[-3]
            t['runend'] = pd.to_datetime(t['runstart']).apply(lambda x:DtFind.last_day_of_month(x)).astype(str)
            t['rundate'] = f.split('_')[-2]
            t['timestamp'] = dt.datetime.strptime(f.split('_')[-1].replace('.csv',''),'%Y%m%d%H%M%S').strftime("%Y-%m-%d %H:%M:%S")
            if df.shape[0] == 0:
                df = t
            else:
                df = pd.concat([df,t])        
        df = df.sort_values(by=['rundate','timestamp'], ascending=False )
        df = df.drop_duplicates(subset=['POLICY_NO', 'rundate'], keep='first')
        df_sum = df.groupby(['filename','rundate','runstart','runend','timestamp']).agg(
            policy_count=('POLICY_NO','count'),
            prediction_average=('PREDICTION','mean'),
            prediction_min=('PREDICTION','min'),
            prediction_max=('PREDICTION','max'),
            prediction_std=('PREDICTION','std')
        ).reset_index()
        if logger:
            logger.info(f"Mstr df_sum: {len(df_sum)}.")
        
        monitor_files = s3_list_folder_contents(
                os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'monitor/'),
                pattern=r'^.*cancel_.*\.json$',
                env=env,
            )
        if incremental_kwargs:
            monitor_files  = [i for i in monitor_files if  (i.split('/')[-1].split('_')[1] in incremental_kwargs['run_dt']) ]  #
        concept_drift_df = pd.DataFrame(columns=['Date','Prediction','Probability'])
        for j in monitor_files:
            file_path  = os.path.join(config['local']['local_folder'], os.path.basename(j))
            s3_download_file(j, file_path, env=env)
            with open(file_path,"r") as f:
                dtime = j.split('/')[-1].split("_")[1]
                data = json.load(f)
                if 'distribution' in data.keys():
                    pred_drift = data['distribution']['concept_drift']['PREDICTION']
                    prob_drift = data['distribution']['concept_drift']['PROBABILITY']
                    t = pd.DataFrame.from_dict({'Date':[dtime],'Prediction':[pred_drift],'Probability':[prob_drift]})
                    concept_drift_df = pd.concat([concept_drift_df,t])
            
            os.remove(file_path)
        concept_drift_df = concept_drift_df.drop_duplicates()

        data_drift_df = pd.DataFrame(columns=['Date','Feature','Drift'])
        for j in monitor_files:
            file_path  = os.path.join(config['local']['local_folder'], os.path.basename(j))
            s3_download_file(j, file_path, env=env)
            with open(file_path,"rb") as f:
                dtime = j.split('/')[-1].split("_")[1]
                data = json.load(f)
                if 'distribution' in data.keys():                
                    for k, v in data['distribution']['data_drift'].items():                
                        t = pd.DataFrame.from_dict({'Date':[dtime],'Feature':[k],'Drift':[v]})
                        data_drift_df = pd.concat([data_drift_df,t])            
            os.remove(file_path)
        data_drift_df = data_drift_df.drop_duplicates()

        evaluate_files = s3_list_folder_contents(
                os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'evaluate/'), 
                pattern=r'^.*cancel_.*\.json$',
                env=env,
            )
        if incremental_kwargs:
            evaluate_files  = [i for i in evaluate_files if  (i.split('/')[-1].split('_')[2] in incremental_kwargs['target_month_first_day']) ]  #        
        eval_df = pd.DataFrame(columns=['Date','Metric','Value'])
        for j in evaluate_files:
            file_path = os.path.join(config['local']['local_folder'], os.path.basename(j))
            s3_download_file(j, file_path, env=env)
            with open(file_path,"rb") as f:
                dtime = j.split('/')[-1].split("_")[2]
                data = json.load(f)
                if 'all' in data.keys():
                    for k, v in data['all'].items():
                        t = pd.DataFrame.from_dict({'Date':[dtime],'Metric':[k],'Value':[v]})
                        eval_df = pd.concat([eval_df,t])
            os.remove(file_path)
        eval_df = eval_df.drop_duplicates()
        eval_df['Value'] = eval_df['Value'].astype(float)

        impact_files = s3_list_folder_contents(
                os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'impact/'),
                pattern=r'^.*cancel_.*\.json$',
                env=env,
            )
        if incremental_kwargs:
            impact_files  = [i for i in impact_files if  (i.split('/')[-1].split('_')[2] in incremental_kwargs['target_month_first_day']) ]  #
        
        impact_df = pd.DataFrame(columns=['Date','Metrics','Values','Datetime'])
        
        if impact_files:   #when there are files
            for j in impact_files:
                file_path  = os.path.join(config['local']['local_folder'], os.path.basename(j))
                s3_download_file(j, file_path, env=env)
                with open(file_path,"rb") as f:
                    date = j.split('/')[-1].split("_")[2]
                    datetime = j.split('/')[-1].split("_")[-1][:8]        
                    data = json.load(f)
                    for k0 in data.keys():                
                        for k, v in data[k0].items():                
                            t = pd.DataFrame.from_dict({'Date':[date],'Metrics':[k0+': '+k],'Values':[v], 'Datetime':[datetime]})
                            impact_df = pd.concat([impact_df,t])            
                os.remove(file_path)
            impact_df = impact_df.sort_values(by='Datetime').drop_duplicates(subset = ['Date','Metrics'], keep='last')
            impact_df[['Category', 'Metrics']] = impact_df['Metrics'].str.split(': ', n=1, expand=True)
            impact_df = impact_df[['Date', 'Category', 'Metrics', 'Values', 'Datetime']]
            impact_df['Values'] = impact_df['Values'].astype(float)
            impact_df = impact_df.pivot(index=['Category', 'Date',], columns=['Metrics'], values='Values').reset_index()
        else: # empty impact_files, this create empty df to merge with 
            impact_df = pd.DataFrame(columns = ['Category', 'Date', 'accounts predicted cancel',
            'accounts predicted cancel & called', 'accounts retained',
            'average call duration (sec)', 'calls inbound', 'calls outbound',
            'calls total', 'premium gained'])

        if incremental_kwargs:
            df_sum, concept_drift_df, data_drift_df, eval_df, impact_df = upsert_all_evaluation(df_sum, concept_drift_df, data_drift_df, eval_df, impact_df, env)

        
        #sort df
        df_sum.sort_values(by ='rundate', inplace=True)
        concept_drift_df.sort_values(by ='Date', inplace=True)
        data_drift_df.sort_values(by ='Date', inplace=True)
        eval_df.sort_values(by ='Date', inplace=True)
        impact_df.sort_values(by =['Category','Date'], inplace=True)

        save_output_mstr(df_sum, concept_drift_df, data_drift_df, eval_df, impact_df, env, athena=True)

    except Exception as e:
        if logger:
            logger.error('Error', exc_info=e)
        error_status = 1
    finally:
        return error_status
        
    
def save_output_mstr(df_sum, concept_drift_df, data_drift_df, eval_df, impact_df, env, athena=False):
    s3_save_and_upload_file(
        df_sum, 
        s3_file_path=os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'cancel_record_counts.csv'),
        overwrite=True,
        env=env,
    )
    s3_save_and_upload_file(
        concept_drift_df, 
        s3_file_path=os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'concept_drift.csv'),
        overwrite=True,
        env=env,
    )
    s3_save_and_upload_file(
        data_drift_df, 
        s3_file_path=os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'data_drift.csv'),
        overwrite=True,
        env=env,
    )
    s3_save_and_upload_file(
        eval_df, 
        s3_file_path=os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'eval.csv'),
        overwrite=True,
        env=env,
    )
    s3_save_and_upload_file(
        impact_df, 
        s3_file_path=os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'impact.csv'),
        overwrite=True,
        env=env,
    )

    if athena:
        s3_save_and_upload_file(
            df_sum, 
            s3_file_path=os.path.join(
                pconfig[f's3_model_output_mstr_folder_{env}'], 
                'Athena/cancel_record_counts/', 
                'cancel_record_counts.parquet'
            ),
            overwrite=True,
            env=env,
        )
        s3_save_and_upload_file(
            concept_drift_df, 
            s3_file_path=os.path.join(
                pconfig[f's3_model_output_mstr_folder_{env}'],
                'Athena/concept_drift/', 
                'concept_drift.parquet'
            ),
            overwrite=True,
            env=env,
        )
        s3_save_and_upload_file(
            data_drift_df, 
            s3_file_path=os.path.join(
                pconfig[f's3_model_output_mstr_folder_{env}'],
                'Athena/data_drift/', 
                'data_drift.parquet'
            ),
            overwrite=True,
            env=env,
        )
        s3_save_and_upload_file(
            eval_df, 
            s3_file_path=os.path.join(
                pconfig[f's3_model_output_mstr_folder_{env}'],
                'Athena/eval/',
                'eval.parquet'
            ),
            overwrite=True,
            env=env,
        )
        s3_save_and_upload_file(
            impact_df, 
            s3_file_path=os.path.join(
                pconfig[f's3_model_output_mstr_folder_{env}'],
                'Athena/impact/',
                'impact.parquet'
            ),
            overwrite=True,
            env=env,
        )

def upsert_all_evaluation(df_sum, concept_drift_df, data_drift_df, eval_df, impact_df, env):
    df_sum = upsert_single_evaluation(df_sum, 
                            'runstart', 
                            os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'cancel_record_counts.csv'),
                            env)

    concept_drift_df = upsert_single_evaluation(concept_drift_df, 
                            'Date', 
                            os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'concept_drift.csv'),
                            env)
    
    data_drift_df = upsert_single_evaluation(data_drift_df, 
                            'Date', 
                            os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'data_drift.csv'),
                            env)

    eval_df = upsert_single_evaluation(eval_df, 
                            'Date', 
                            os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'eval.csv'),
                            env)

    impact_df = upsert_single_evaluation(impact_df, 
                            'Date', 
                            os.path.join(pconfig[f's3_model_output_mstr_folder_{env}'], 'impact.csv'),
                            env)
    return df_sum, concept_drift_df, data_drift_df, eval_df, impact_df

def upsert_single_evaluation(df, join_key, s3_file_path, env):
    df = df.copy()    
    df_aws = s3_download_and_read_file(s3_file_path=s3_file_path, env=env)
    df = pd.concat([df_aws[~df_aws[join_key].isin(df[join_key])], df], ignore_index=True)
    return df

def get_incremental_update_param(month_start, month_end):
    months = get_run_months(month_start, month_end)
    kwargs ={}
    kwargs['target_month'] = months
    run_day=pconfig['run_day']
    if str(run_day).isnumeric():
        kwargs['run_dt']      = [DtFind.first_day_last_month(dt.datetime.strptime(i, '%Y-%m')).replace(day=run_day).strftime('%Y-%m-%d') for i in months]
        kwargs['target_month_first_day'] = [dt.datetime.strptime(i, '%Y-%m').strftime('%Y-%m-%d') for i in months]
    return kwargs