__all__ = ['admission', 'score_pipeline']

from ._imports import *
from ._pconfig import *
from ._features import *
from ._monitor import *
from ._train import train_pipeline, dataset_initialize, dataset_filter, load_model_config, load_model
from ._misc import *

def admission():
    print('testing...')

def score_pipeline(run_month, run_day=pconfig['run_day'], env='dev', logger=None):
    error_status = 0
    if logger:
        logger.info(f'Score for {run_month} performed on date {run_day}.')
    try:
        runtime = create_runtime(run_month, run_day, env, logger)
        df_base = dataset_initialize(runtime)
        df_feat = dataset_features(df_base, runtime)
        df_filt = dataset_filter(df_feat, runtime)
        df_pred = dataset_score(df_filt, runtime)

        save_output_eds(df_pred, runtime)
        save_output_care(df_pred, runtime)

        model_config = load_model_config(runtime)
        trackers     = monitor_pipeline(df_pred, model_config, logger=logger, env=env)
        trackers['error'] = error_status
        
        if trackers['alert'] or run_month.endswith('-01'):
            # if check_previous_retrain(runtime):
            #     update_model_version(env=env, promote=False)
            # else:
            today_dt     = dt.datetime.today()
            train_start  = DtFind.first_day_of_month(today_dt - dt.timedelta(days=8*30), '%Y-%m')
            train_end    = DtFind.first_day_of_month(today_dt - dt.timedelta(days=5*30), '%Y-%m')
            if logger:
                logger.info(f'Train from {train_start} to {train_end}.')
            train_pipeline(train_start, train_end, env=env, logger=logger, update_model=True)


    except Exception as e:
        if logger:
            logger.error('Error', exc_info=e)
        error_status = 1
        trackers     = {'error': error_status}
    finally:
        save_output_mstr(trackers, runtime)
        return error_status

def create_runtime(run_month, run_day, env, logger):
    timestamp = dt.datetime.now().strftime("%Y%m%d%H%M%S")
    if logger:
        logger.info(f'Run timestamp: {timestamp}')

    features = Features()
    features.remove_features(['CANCEL_RENEWAL', 'CANCEL_EXPIRING'])
    features.remove_features(['CC_CLAIM_NOTES'])
    #remove demographic 
    demo_feature = ['GOOD_STUDENT','NUM_CONVICTIONS','DRIVER_AGE_DISC','DRIVER_TRAINING','YOUTHFUL_OPER','STUDENT_AWAY','MILITARY_IND','RIDER_TRAINING','HOUSEHOLD_INCOME']
    features.remove_features(demo_feature)
    # features.remove_features([c for c in features.get_feature_dict() if c.startswith('PC') or c.startswith('CC')]) # debugging purpose for now

    runtime = {}
    runtime['timestamp']      = timestamp
    runtime['env']            = env
    runtime['features']       = features.get_feature_dict()
    runtime['teradata_table'] = f'{config["teradata"]["user"]}.temp_table'
    runtime['type_dt']        = pconfig['type_dt']

    run_month = dt.datetime.strptime(run_month, '%Y-%m')
    runtime['run_month']       = run_month
    runtime['run_month_start'] = DtFind.first_day_of_month(run_month).strftime('%Y-%m-%d')
    runtime['run_month_end']   = DtFind.last_day_of_month(run_month).strftime('%Y-%m-%d')
    if str(run_day).isnumeric():
        runtime['run_dt']      = DtFind.first_day_last_month(run_month).replace(day=run_day).strftime('%Y-%m-%d')
    else: # allow passing in '2023-04-27' as run day for example
        runtime['run_dt']      = dt.datetime.strptime(run_day,'%Y-%m-%d').strftime('%Y-%m-%d')
    return runtime

def dataset_score(df, runtime):
    model_config = load_model_config(runtime)
    model        = load_model(model_config, runtime)

    df_feat  = df[model_config['features']]
    probs    = model.predict_proba(df_feat)[:,1]
    preds    = (probs>float(model_config['threshold'])).astype(int)
    df_score = df.assign(
        MODEL_VERSION = model_config['model_version'],
        THRESHOLD     = model_config['threshold'],
        PROBABILITY   = probs,
        PREDICTION    = preds,
    )
    return df_score

def check_previous_retrain(runtime, window=30, dt_format='%Y-%m-%d'):
    fnames = s3_list_folder_contents(os.path.join(
        pconfig[f's3_model_output_mstr_folder_{runtime["env"]}'],
        'retrain/'
    ), pattern=r'^.*cancel_model_update.*\.json$')
    fnames = pd.DataFrame({'file': fnames})
    if len(fnames)==0:
        return False
    pat    = r'cancel_model_update_([\d]*)\.json'
    fnames['timestamp'] = fnames['file'].str.extract(pat)
    fnames['run_date']  = pd.to_datetime(fnames['timestamp']).dt.strftime('%Y-%m-%d')
    
    for _, row in fnames.iterrows():
        trackers = s3_download_and_read_json(row['file'])
        trackers_df = pd.Series(trackers)
        fnames.loc[fnames['file']==row['file'], trackers_df.index] = trackers_df.values
    fnames = fnames.loc[fnames['action']=='promote']
    fnames = fnames.sort_values(by=['datetime'])
    if len(fnames)==0:
        return False
    promotion_date = fnames.iloc[-1].to_dict()['datetime']

    run_date       = dt.datetime.strptime(runtime['run_dt'], dt_format)
    promotion_date = dt.datetime.strptime(promotion_date, dt_format)
    return promotion_date > run_date - dt.timedelta(window)

def save_output_eds(df, runtime):
    s3_save_and_upload_csv(
        df, 
        s3_file_path = os.path.join(
            pconfig[f's3_model_output_eds_folder_{runtime["env"]}'],
            f'cancel_{runtime["run_month_start"]}_{runtime["run_dt"]}_{runtime["timestamp"]}.csv'
        ),
        env=runtime['env'],
    )

def save_output_care(df, runtime):
    df = df[['POLICY_NO', 'ACCOUNT_NO', 'PREDICTION', 'PROBABILITY', 'THRESHOLD', 'EFFECTIVE_DATE', 'EXPIRATION_DATE']]
    df['ID']             = 'Retention_' + df['POLICY_NO']
    df['Source']         = 'Retention'
    df['RecCreatedDate'] = dt.datetime.strptime(runtime['timestamp'], "%Y%m%d%H%M%S")
    df['PREDICTION']     = df['PREDICTION'].apply(lambda r: 'ActionableCancel' if r>0 else 'NotActionableCancel')
    df = df.rename(columns={
        'POLICY_NO': 'PolicyNumber', 'ACCOUNT_NO': 'AccountNumber',
        'EFFECTIVE_DATE': 'PolicyEffectiveDate', 'EXPIRATION_DATE': 'PolicyExpiredDate',
        'PREDICTION': 'Prediction', 'PROBABILITY': 'Probability', 'THRESHOLD': 'Threshold',
    })
    df = df[['ID', 'Source', 'PolicyNumber', 'AccountNumber', 'Prediction', 'Probability', 'Threshold', 'PolicyEffectiveDate', 'PolicyExpiredDate', 'RecCreatedDate']]

    s3_save_and_upload_csv(
        df, 
        s3_file_path = os.path.join(
            pconfig[f's3_model_output_care_folder_{runtime["env"]}'],
            f'cancel_{runtime["run_month_start"]}_{runtime["timestamp"]}.csv'
        ),
        env=runtime['env'],
    )

    s3_save_and_upload_csv(
        df, 
        s3_file_path = os.path.join(
            pconfig[f's3_model_output_care_folder_{runtime["env"]}'],
            'prediction_current.csv'
        ),
        overwrite=True,
        env=runtime['env'],
    )

    if any([(os.path.basename(f)=='prediction_history.csv') for f in s3_list_folder_contents(pconfig[f's3_model_output_care_folder_{runtime["env"]}'])]):
        df_history = s3_download_and_read_csv(os.path.join(
            pconfig[f's3_model_output_care_folder_{runtime["env"]}'],
            'prediction_history.csv'
        ), env=runtime['env'])
        df_history['RecCreatedDate'] = pd.to_datetime(df_history['RecCreatedDate'], errors='coerce')
        df_history = pd.concat([df_history, df], ignore_index=True)
        df_history = (
            df_history.sort_values(by='RecCreatedDate', ascending=False)
            .drop_duplicates(subset=['PolicyNumber', 'AccountNumber', 'PolicyEffectiveDate'], keep='first')
        )
        s3_save_and_upload_csv(
            df_history, 
            s3_file_path = os.path.join(
                pconfig[f's3_model_output_care_folder_{runtime["env"]}'],
                'prediction_history.csv'
            ),
            overwrite=True,
            env=runtime['env'],
        )
        
    else:
        s3_save_and_upload_csv(
            df, 
            s3_file_path = os.path.join(
                pconfig[f's3_model_output_care_folder_{runtime["env"]}'],
                'prediction_history.csv'
            ),
            overwrite=False,
            env=runtime['env'],
        )

def save_output_mstr(data, runtime):
    s3_save_and_upload_json(
        data,
        s3_file_path = os.path.join(
            pconfig[f's3_model_output_mstr_folder_{runtime["env"]}'],
            'monitor/',
            f'cancel_{runtime["run_dt"]}_{runtime["timestamp"]}.json'
        ),
        env=runtime['env'],
    )