__all__ = ['admission', 'train_pipeline']

from ._imports import *
from ._pconfig import *
from ._features import *
from ._misc import *

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

def train_pipeline(train_start, train_end, env='dev', logger=None, automl=True, update_model=False):
    error_status = 0
    if logger:
        logger.info(f'Train from {train_start} to {train_end}.')
    try:
        runtime_train = create_train_runtime(train_start, train_end, env=env, logger=logger)
        run_months = get_run_months(train_start, train_end)

        df = pd.DataFrame()
        for run_month in run_months:
            runtime = create_runtime_per_runmonth(run_month, env, logger)
            df_base = dataset_initialize(runtime)
            df_feat = dataset_features(df_base, runtime)
            df_targ = dataset_target(df_feat, runtime)
            df_filt = dataset_filter(df_targ, runtime)
            df = pd.concat([df, df_filt], ignore_index=True)
            if logger:
                logger.info(f'Run date {run_month} with {len(df_targ)} records.')

        if automl:
            evals = dataset_retrain(df, runtime_train, automl=automl_flaml, automl_kwargs={'time_budget': 60*60})
        else:
            evals = dataset_retrain(df, runtime_train)
        save_output_mstr(evals, runtime_train)

        if update_model:
            update_model_version(env, promote=True)
    except Exception as e:
        if logger:
            logger.error('Error', exc_info=e)
        error_status = 1
    finally:
        return error_status

def create_runtime_per_runmonth(run_month, env, logger, run_day=pconfig['run_day']):
    features = Features()
    # features.remove_features([c for c in features.get_feature_dict() if c.startswith('PC') or c.startswith('CC')]) # debugging for now
    features.remove_features(['CC_CLAIM_NOTES'])
    #remove demo
    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)

    runtime = {}
    runtime['timestamp']      = dt.datetime.now().strftime("%Y%m%d%H%M%S")
    runtime['env']            = env
    runtime['logger']         = logger

    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')
    runtime['run_dt']          = DtFind.first_day_last_month(run_month).replace(day=run_day).strftime('%Y-%m-%d')

    return runtime

def create_train_runtime(run_start, run_end, env, logger):
    runtime                   = {}
    runtime['timestamp']      = dt.datetime.now().strftime("%Y%m%d%H%M%S")
    runtime['env']            = env
    runtime['logger']         = logger
    runtime['teradata_table'] = f'{config["teradata"]["user"]}.temp_table'
    runtime['run_start']      = run_start
    runtime['run_end']        = run_end
    return runtime 

def get_run_months(run_start, run_end, dt_format='%Y-%m'):
    a = dt.datetime.strptime(run_start, dt_format)
    b = dt.datetime.strptime(run_end, dt_format)
    run_months = [a.strftime(dt_format)]
    while a.strftime(dt_format) != b.strftime(dt_format):
        a = DtFind.first_day_next_month(a)
        run_months.append(a.strftime(dt_format))
    return run_months

def pc_initialize(df):
    pc_temp(df['POLICY_NO'].rename('policy_number'), '#tmp_policynumber')
    pc_temp(df['ACCOUNT_NO'].rename('Amica_number'), '#tmp_amicanumber')

def dataset_initialize(runtime):
    table_name  = runtime['teradata_table']
    table_query = f'''
        select 
            trim(a.policy_key) as POLICY_KEY
            , trim(a.policy_no) as POLICY_NO
            , trim(a.account_no) as ACCOUNT_NO
            , trim(a.perpetual_pol_no) as PERPETUAL_POL_NO
            , trim(a.expired_policy_no) as EXPIRED_POLICY_NO
            , a.effective_date as EFFECTIVE_DATE
            , a.expiration_date as EXPIRATION_DATE
            , a.process_ts as PROCESS_TS
        from atuprv.dw_agreement a
        left join atuprv.dw_state b on a.iso_state_number = b.iso_state_number
        left join atuprv.dw_form_desc c on a.form_cd = c.code
        where 
            a.{runtime["type_dt"]}_date between date \'{runtime["run_month_start"]}\' and date \'{runtime["run_month_end"]}\'
            and cast(a.process_ts as Date) <= \'{runtime["run_dt"]}\'
            and a.policy_type = 1 -- auto line of business
            and a.policy_activity = 'Renewal'
        qualify row_number() over (partition by a.policy_key order by a.process_ts desc)=1
    '''

    # TERADATA TEMP TABLE
    td_drop_table(table_name)
    td_volatile(table_query, table_name)
    df = td_query(query_creator(table_columns='*', table_name=table_name, table_ops=''))

    # POLICY CENTER TEMP TABLE
    #pc_initialize(df)
    return df

def dataset_filter(df, runtime):
    if runtime['features']['SOURCE_OF_BUSINESS']:
        df = df.loc[df['SOURCE_OF_BUSINESS']!='MAIP POLICY']
    if runtime['features']['STATE']:
        df = df.loc[
            (df['STATE_CODE'].isin(
                ('AK', 'AL', 'AR', 'AZ', 'CA', 'CO', 'CT', 'DC', 'DE', 'FL', 'GA', 'HI', 'IA', 'ID', 'IL', 'IN', 'KS', 'KY', 'LA', 'MA', 
                'MD', 'ME', 'MI', 'MN', 'MO', 'MS', 'MT', 'NC', 'ND', 'NE', 'NH', 'NJ', 'NM', 'NV', 'NY', 'OH', 'OK', 'OR', 'PA', 'RI', 'SC', 'SD', 'TN', 
                'TX', 'UT', 'VA', 'VT', 'WA', 'WI', 'WV', 'WY')
            ))
        ]
    if runtime['features']['POLICY_FORM']:
        df = df.loc[
            (df['POLICY_FORM'].isin(
                ('MOTORCYCLE', 'TRAILER', 'SNOWMOBILE', 'MANUAL AUTO', 'ANTIQUE AUTO', 'PRIVATE PASSENGER', 'MOBILE HOME', 'MOTOR HOME', 'GOLF CART')
            ))
        ]
    return df

def dataset_target(df, runtime):
    if not runtime['features']['CANCEL_RENEWAL'] or not runtime['features']['CANCEL_EXPIRING']:
        return df

    def target_cancel(row):
        cancel_reasons = [
            'RatesAnotherCarrier', 
            'nonpayment', 
            'EffortsDeterReasonFailed', 
            'RatesAnotherCarrierOutJur', 
            'PlacedLocalAgentSpouseCo',
        ]
        target = (
            (
                row['CANCEL_POST_RENEWAL_DAYS'] < np.inf
            ) and (
                (
                    row['CANCEL_RENEWAL_REASON'] in cancel_reasons 
                    and 
                    row['CANCEL_POST_RENEWAL_DAYS'] <= pconfig['target_post_window']
                ) or (
                    row['CANCEL_EXPIRING_REASON'] in cancel_reasons 
                    and 
                    row['CANCEL_PRE_RENEWAL_DAYS'] <= pconfig['target_pre_window']
                )
            )
        )
        return target

    df[pconfig['target']] = (
        df
        .assign(
            CANCEL_PRE_RENEWAL_DAYS = (
                pd.to_datetime(df['EFFECTIVE_DATE'], errors='coerce')
                - 
                pd.to_datetime(df['CANCEL_EXPIRING_DATE'], errors='coerce')
            ).dt.days.fillna(np.inf), 
            CANCEL_POST_RENEWAL_DAYS = (
                pd.to_datetime(df['CANCEL_RENEWAL_DATE'], errors='coerce')
                - 
                pd.to_datetime(df['EFFECTIVE_DATE'], errors='coerce')
            ).dt.days.fillna(np.inf),
        )
        .apply(target_cancel, axis=1).astype(int)
    )
    return df

def dataset_retrain(df, runtime, holdout_size=0.2, valid_size=0.2, random_state=42, resample=False, automl=None, automl_kwargs={}):
    model_config = load_model_config(runtime)
    model        = load_model(model_config, runtime)

    train, holdout = sklearn.model_selection.train_test_split(
        df, test_size=holdout_size, random_state=random_state, stratify=df[pconfig['target']]
    )
    train, valid   = sklearn.model_selection.train_test_split(
        train, test_size=valid_size, random_state=random_state,
    )

    train   = train[model_config['features'] + [pconfig['target']]]
    holdout = holdout[model_config['features'] + [pconfig['target']]]
    valid   = valid[model_config['features'] + [pconfig['target']]]

    if automl:
        estimator = automl(
            train[model_config['features']], train[pconfig['target']],
            valid[model_config['features']], valid[pconfig['target']],
            **automl_kwargs
        )
        model_retrain = sklearn.pipeline.make_pipeline(model[-2], estimator)
    else:
        model_retrain = sklearn.base.clone(model)
    

    if resample:
        train = undersample_target(train)

    model_retrain.fit(train[model_config['features']], train[pconfig['target']])
    valid_probs   = model_retrain.predict_proba(valid[model_config['features']])[:,1]
    holdout_probs = model_retrain.predict_proba(holdout[model_config['features']])[:,1]

    thr           = get_threshold(valid[pconfig['target']], valid_probs, type='f1')
    holdout_preds = (holdout_probs>thr).astype(int)
    evals         = evaluate_metrics(holdout[pconfig['target']], holdout_probs, holdout_preds)
    importances   = feature_importance(train, model_retrain, model_config)

    model_version_retrain = f'{model_config["model_version"]}.retrain'
    model_version_update  = str(int(model_config["model_version"]) + 1)
    model_config_retrain  = {
        'model_name':   'propensity_to_cancel',
        'model_version': model_version_update,
        'model_type':    model_retrain[-1].__class__.__name__,
        'model_file':    f'cancel_v{model_version_retrain}.pkl',
        'features':      model_config['features'],
        'threshold':     str(np.round(thr, 3)),
        'evaluation':    evals,
        'importances':   importances,
    }
    model_ref_retrain = holdout.assign(
        PROBABILITY = holdout_probs,
        PREDICTION  = holdout_preds,
    )

    #save_model_assets('dev', model_retrain, model_config_retrain, model_ref_retrain, model_version_retrain, model_data=df)
    save_model_assets(runtime['env'], model_retrain, model_config_retrain, model_ref_retrain, model_version_retrain, model_data=df)
    
    return evals

def undersample_target(df):
    df_balanced = df.groupby(df[pconfig['target']])
    df_balanced = df_balanced.apply(lambda x: 
        x.sample(df_balanced.size().min())
    ).reset_index(drop=True)
    return df_balanced

def preprocess_features(model_config, train, holdout=None, valid=None):
    txt_columns = ['CLAIM_NOTES']
    cat_columns = train.select_dtypes(include=['object', 'category']).columns.to_list()
    num_columns = train.select_dtypes(include='number').columns.to_list()
    
    cat_columns = [c for c in cat_columns if c in model_config['features'] and c not in txt_columns]
    num_columns = [c for c in num_columns if c in model_config['features'] and c not in txt_columns]

    pipe_num = sklearn.pipeline.make_pipeline(
        sklearn.impute.SimpleImputer(strategy='median'), 
        sklearn.preprocessing.MinMaxScaler(),
     )
    pipe_cat = sklearn.pipeline.make_pipeline(
        sklearn.impute.SimpleImputer(strategy='constant', fill_value='Missing'), 
        sklearn.preprocessing.OneHotEncoder(handle_unknown='ignore'),
    )

    transformer = sklearn.compose.make_column_transformer(
        (pipe_num, num_columns),
        (pipe_cat, cat_columns),
        remainder='passthrough',
    )

    X_train = transformer.fit_transform(train[model_config['features']])
    X_train = pd.DataFrame(X_train, columns=transformer.get_feature_names_out())
    if valid is not None: 
        X_valid = transformer.transform(valid[model_config['features']])
        X_valid = pd.DataFrame(X_valid, columns=transformer.get_feature_names_out())
    if holdout is not None: 
        X_holdout = transformer.transform(holdout[model_config['features']])
        X_holdout = pd.DataFrame(X_holdout, columns=transformer.get_feature_names_out())
    return X_train, X_valid, X_holdout

def evaluate_metrics(actual, prob, pred):
    evals = {
        'support':     len(actual),
        'accuracy':    sklearn.metrics.accuracy_score(actual, pred).round(2),
        'precision':   sklearn.metrics.precision_score(actual, pred).round(2),
        'recall':      sklearn.metrics.recall_score(actual, pred).round(2),
        'specificity': sklearn.metrics.recall_score(actual, pred, pos_label=0).round(2),
        'f1':          sklearn.metrics.f1_score(actual, pred).round(2),
        'rocauc' :     sklearn.metrics.roc_auc_score(actual, prob).round(2),
    }
    return evals

def get_threshold(test_y, test_probs, type='roc_auc'):
    if type == 'roc_auc':
        fpr, tpr, thr = metrics.roc_curve(test_y, test_probs)
        thr = thr[np.argmax(tpr-fpr)]
    else: # type == 'f1'
        precision, recall, thr = metrics.precision_recall_curve(test_y, test_probs)
        precision_filter = precision[(precision>0) & (recall>0)]
        recall_filter = recall[(precision>0) & (recall>0)]
        precision = precision_filter
        recall = recall_filter
        thr = thr[np.argmax(precision*recall/(precision+recall))]
    return thr

def feature_importance(df, model, model_config, n=10000, random_state=100):
    permute_importance = {}
    n_iter = 10

    feats  = model_config['features']
    thr    = float(model_config['threshold'])

    sample = df.sample(n=min(n, len(df)), random_state=random_state)
    sample_probs = model.predict_proba(sample[feats])[:,1]
    sample_preds = (sample_probs>thr).astype(int)
    sample_acc   = metrics.accuracy_score(sample[pconfig['target']], sample_preds)
    for f in feats:
        permute_acc_n = []
        for n in range(n_iter):
            permute_feats = sample[feats].copy(deep=True)
            permute_feats[f] = permute_feats[f].sample(frac=1, random_state=n).values
            permute_probs = model.predict_proba(permute_feats)[:,1]
            permute_preds = (permute_probs>thr).astype(int)
            permute_acc = metrics.accuracy_score(sample[pconfig['target']], permute_preds)
            permute_acc_n.append(permute_acc)
        permute_importance[f] = sample_acc - np.mean(permute_acc_n)
    permute_importance = pd.Series(permute_importance)
    permute_importance = permute_importance / permute_importance.max()
    permute_importance = permute_importance.clip(lower=0).sort_values(ascending=False)
    return permute_importance.to_dict()

def load_model_config(runtime):
    model_version           = get_model_version(runtime["env"])
    local_model_config_file = os.path.join(config['local']['local_folder'], 'model_config.json')
    s3_model_config_file    = os.path.join(pconfig[f's3_model_assets_folder_{runtime["env"]}'], f'cancel_v{model_version}.json')
    s3_download_file(s3_model_config_file, local_model_config_file, env=runtime['env'])

    with open(local_model_config_file, 'r') as f:
        model_config = json.load(f)
    os.remove(local_model_config_file)
    return model_config

def load_model(model_config, runtime):
    local_model_pickle_file = os.path.join(config['local']['local_folder'], 'model.pkl')
    s3_model_pickle_file    = os.path.join(pconfig[f's3_model_assets_folder_{runtime["env"]}'], model_config['model_file'])
    s3_download_file(s3_model_pickle_file, local_model_pickle_file, env=runtime['env'])

    with open(local_model_pickle_file, 'rb') as f:
        model = pickle.load(f)
    os.remove(local_model_pickle_file)
    return model
    
def save_model_assets(env, model, model_config, model_ref, model_version, model_data=None, overwrite=True):
    s3_save_and_upload_file(
        model_config,
        s3_file_path = os.path.join(pconfig[f's3_model_assets_folder_{env}'], f'cancel_v{model_version}.json'),
        overwrite=overwrite,
        env=env,
    )
    s3_save_and_upload_file(
        model,
        s3_file_path = os.path.join(pconfig[f's3_model_assets_folder_{env}'], f'cancel_v{model_version}.pkl'),
        overwrite=overwrite,
        env=env,
    )
    s3_save_and_upload_file(
        model_ref,
        s3_file_path = os.path.join(pconfig[f's3_model_assets_folder_{env}'], f'cancel_v{model_version}.csv'),
        overwrite=overwrite,
        env=env,
    )
    if model_data is not None:
        s3_save_and_upload_file(
            model_data,
            s3_file_path = os.path.join(pconfig[f's3_model_assets_folder_{env}'], f'data_cancel_v{model_version}.csv'),
            overwrite=overwrite,
            env=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"]}'],
            'retrain/',
            f'cancel_train_{runtime["run_start"]}_{runtime["run_end"]}_{runtime["timestamp"]}.json'
        ),
        env=runtime['env'],
    )

def automl_flaml(train_x, train_y, val_x, val_y, time_budget=60):
    weight  = train_y.value_counts()[0]/train_y.value_counts()[1]
    
    custom_hp = {
        "xgb_limitdepth": {
            "n_estimators": {
                "domain": flaml.tune.lograndint(lower=10, upper=200),
                "low_cost_init_value": 10,
            },
            "scale_pos_weight":{
                "domain": weight,
                "init_value": weight,
            },
        },
        "xgboost": {
            "n_estimators": {
                "domain": flaml.tune.lograndint(lower=10, upper=200),
                "low_cost_init_value": 10,
            },
            "scale_pos_weight":{
                "domain": weight,
                "init_value": weight,
            },
        },
        "rf": {
            "n_estimators": {
                "domain": flaml.tune.lograndint(lower=10, upper=200),
                "low_cost_init_value": 10,
            },
            "class_weight": {
                "domain":'balanced',
                "init_value":'balanced',
            },
        },
        "extra_tree": {
            "n_estimators": {
                "domain": flaml.tune.lograndint(lower=10, upper=200),
                "low_cost_init_value": 10,
            },
            "class_weight": {
                "domain":'balanced',
                "init_value":'balanced',
            },
        },
        "catboost": {
            "n_estimators": {
                "domain": flaml.tune.lograndint(lower=10, upper=200),
                "low_cost_init_value": 10,
            },
            "scale_pos_weight": {
                "domain": weight,
                "init_value": weight,
            },
        },
        "lgbm": {
            "n_estimators": {
                "domain": flaml.tune.lograndint(lower=10, upper=200),
                "low_cost_init_value": 10,
            },
            "num_leaves": {
                "domain": flaml.tune.lograndint(lower=5, upper=100),
                "low_cost_init_value": 5,
            },
            "scale_pos_weight": {
                "domain": weight,
                "init_value": weight,
            },
        },
    }
    clf_settings = {
        'time_budget': time_budget,
        'metric': 'log_loss',
        'task': 'classification',
        'estimator_list': [
                'rf', 
                'xgboost',
                'xgb_limitdepth', 
                'extra_tree',  
                'lgbm',
            ],
        'verbose': 0,
        'custom_hp': custom_hp,
    }
    clf_dict = {
            'lgbm':             lightgbm.LGBMClassifier, 
            'rf':               sklearn.ensemble.RandomForestClassifier, 
            'xgboost':          xgboost.XGBClassifier, 
            'xgb_limitdepth':   xgboost.XGBClassifier, 
            'extra_tree':       sklearn.ensemble.ExtraTreesClassifier, 
        }

    automl = flaml.AutoML()
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        automl.fit(
                X_train=train_x, y_train=train_y, 
                X_val=val_x, y_val=val_y,
                **clf_settings
            )
    clf = clf_dict[automl.best_estimator]
    clf_params = automl.best_config
    if automl.best_estimator not in ['xgboost', 'xgb_limitdepth']:
        clf_params = {p:clf_params[p] for p in clf_params if p in clf.__init__.__code__.co_varnames}
    return clf(**clf_params)

