import uuid
from datetime import datetime
import sys
from xmlrpc.client import Boolean
import ipdb
import itertools
import pandas as pd
from sklearn.base import clone
# from sklearn.experimental import enable_halving_search_cv # noqa
# from sklearn.model_selection import HalvingGridSearchCV
from sklearn.model_selection import GridSearchCV, KFold, train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import (roc_auc_score, accuracy_score,
                             average_precision_score)
# from ml.pmc.utils import MC_loss
import warnings
import time
from tempfile import mkdtemp
from shutil import rmtree
from read_file import read_file
import pdb
import numpy as np
import json
import os
import inspect
from util import jsonify, hasattranywhere
from pmc.auditor import Auditor
from pmc.metrics import (differential_calibration, 
                            multicalibration_score,
                            proportional_multicalibration_score,
                            multicalibration_loss,
                            proportional_multicalibration_loss
                           )
DEFAULT_GROUPS=['ethnicity','gender']

def evaluate_model(
    dataset, 
    results_path,
    emb_path, 
    random_state, 
    ml, 
    est, 
    alpha,
    n_bins,
    gamma,
    rho,
    one_hot_encoded,
    text_features,
    pre_train=None,
    groups=DEFAULT_GROUPS
    
):
    """Main evaluation routine."""


    run_id = str(uuid.uuid4())
    ##################################################
    # setup data
    ##################################################
    tmp = read_file(dataset,
                    one_hot_encode = one_hot_encoded,
                    label='y',
                    text_features=text_features,
                    emb_path = emb_path)
    if len(tmp) == 2:
        features, labels = tmp 
    elif len(tmp) == 3:
        print('Note: changing "groups" from',groups,'to',tmp[2])
        features, labels, groups = tmp 
    print('features:')
    print(features.head())
    print(features.shape)
    # generate train/test split
    X_train, X_test, y_train, y_test = train_test_split(features, labels,
                                                    train_size=0.75,
                                                    test_size=0.25,
                                                    shuffle=True,
                                                    random_state=random_state,
                                                    stratify=labels
                                                    )                                                      

    ########################################
    # configure estimators
    ########################################
    setatts = {
        'random_state':random_state, 
        'alpha':alpha,
        'n_bins':n_bins,
        'gamma':gamma,
        'rho':rho,
    }
    np.random.seed(random_state)
    for k,v in setatts.items():
        if hasattr(est, k):
            setattr(est, k, v)
    # if hasattr(est, 'n_jobs'):
    #     est.n_jobs = 1
    if groups is not None:
        if hasattr(est, 'auditor_type'):
            est.auditor_type = Auditor(groups=groups)
    print(40*'=','Evaluating '+ml+' on ',dataset,40*'=',sep='\n')


    print('X_train:',X_train.shape)
    print('y_train:',y_train.shape)
    

    ################################################## 
    # Fit models
    ################################################## 
    print('training',est)
    t0p = time.process_time()
    t0t = time.time()
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        est.fit(X_train, y_train)
    process_time = time.process_time() - t0p
    time_time = time.time() - t0t
    print('Training time measures:',process_time, time_time)
    
    ##################################################
    # store results
    dataset_name = dataset.split('/')[-1].split('.')[0]
    results = {
        'algorithm':ml,
        'dataset':dataset,
        'params':jsonify(est.get_params()),
        'process_time': process_time, 
        'time_time': time_time, 
        'text_encoding': one_hot_encoded,
        'text_features' : text_features,
        'groups': groups,
        'run_id': run_id
    }
    results.update(setatts)
    if hasattr(est, 'n_updates_'):
        results['n_updates'] = est.n_updates_

    ##############################
    # scores
    ##############################
    for fold, target, X in zip(['train','test'],
                               [y_train, y_test], 
                               [X_train, X_test]
                              ):
            y_pred = est.predict(X).reshape(-1,1)
            y_pred_proba = est.predict_proba(X)[:,1]
            for score, scorer in [('roc_auc',roc_auc_score),
                                  ('auprc',average_precision_score)]:
                # ipdb.set_trace()
                results[score + '_' + fold] = scorer(target, y_pred_proba) 
                print(score + '_' + fold,
                      f'{scorer(target, y_pred_proba):.3f}')
            for score, scorer in [('accuracy',accuracy_score)]:
                results[score + '_' + fold] = scorer(target, y_pred) 
                print(score + '_' + fold,
                      f'{scorer(target, y_pred):.3f}')
            
            y_pred_proba = pd.Series(y_pred_proba, index=target.index)
            X = X.set_index(target.index)
            # MC
            results['MC_loss_' + fold] = multicalibration_loss(
                estimator=est,
                X=X, 
                y_true=target, 
                groups=groups,
                n_bins=n_bins,
                alpha=alpha,
                gamma=gamma,
                rho=rho,
                proportional=False,
            )
            print('MC_loss_' + fold,
                  f"{results['MC_loss_' + fold]:.3f}")
            # PMC
            results['PMC_loss_' + fold] = \
                    proportional_multicalibration_loss(
                estimator=est,
                X=X, 
                y_true=target, 
                groups=groups,
                n_bins=n_bins,
                proportional=True,
                alpha=alpha,
                gamma=gamma,
                rho=rho
            )
            print('PMC_loss_' + fold,
                  f"{results['PMC_loss_' + fold]:.3f}")
            # differential calibration loss
            results['DC_loss_' + fold] = differential_calibration(
                estimator=est,
                y_true=target, 
                X=X, 
                groups=groups,
                n_bins=n_bins,
                alpha=alpha,
                gamma=gamma,
                rho=rho
            )
            print('DC_loss_' + fold,
                  f"{results['DC_loss_' + fold]:.3f}")
    ######################################## 
    # save predictions and sensitive attributes
    test_pred_df = X_test[groups]
    test_pred_df.loc[:,'y_true'] =y_test
    test_pred_df.loc[:,'y_pred'] = est.predict(X).reshape(-1,1)
    test_pred_df.loc[:,'y_pred_proba'] = est.predict_proba(X_test)[:,1]
    
    if hasattr(est, 'feature_importances_'):
        results['feature_importances_'] = \
                {fn:imp 
                 for fn,imp in zip(X_train.columns, est.feature_importances_)}
        print('feature importances:',results['feature_importances_'])

    ##############################
    # write to file
    ##############################
    if not os.path.exists(results_path):
        os.makedirs(results_path, exist_ok=True)

    save_file = os.path.join(
        results_path, 
        '_'.join([ f'{n}' for n in [
            dataset_name,
            ml,
            one_hot_encoded,
            random_state,
            alpha,
            gamma,
            rho,
            datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
            os.environ['LSB_JOBID'] if 'LSB_JOBID' in os.environ.keys() else '',
            run_id
            ]
        ])
    )

    print('save_file:',save_file)
    # scores
    with open(save_file + '.json', 'w') as out:
        json.dump(jsonify(results), out, indent=4)
    # predictions
    test_pred_df.to_csv(save_file + '.csv')
    # stats
    if hasattr(est, 'stats_'):
        est.stats_.to_csv(save_file + '.stats')

    # store CV detailed results
    # turning off for now as I dont think we'll need this for our analysis
    # cv_results = grid_est.cv_results_
    # cv_results['random_state'] = random_state

    # with open(save_file + '_cv_results.json', 'w') as out:
    #     json.dump(jsonify(cv_results), out, indent=4)
    ##################################################

################################################################################
# main entry point
################################################################################
import argparse
import importlib
import logging
import sys

if __name__ == '__main__':

    logging.basicConfig(stream=sys.stdout, level=logging.INFO) 

    # parse command line arguments
    parser = argparse.ArgumentParser(
        description="Evaluate a method on a dataset.", add_help=False)
    parser.add_argument('-file', action='store', type=str,
                        default='data/mimic4_admissions.csv',
                        help='Data file to analyze; ensure that the '
                             'target/label column is labeled as "y". '
                             'If you use the preprocessing file, '
                             'you do not need to do anything')    
    parser.add_argument('-h', '--help', action='help',
                        help='Show this help message and exit.')
    parser.add_argument('-ml', action='store', default='xgb',type=str, 
            help='Name of estimator (with matching file in ml/)')
    parser.add_argument('-results_path', action='store', dest='RDIR',
                        default='results/', type=str, 
                        help='Name of save file')
    parser.add_argument('-emb_path', action='store', dest='EMB',
                        default='preprocessing/BCH_embedding.npy', type=str, 
                        help='Path of pre_trained embedding saved file')
    parser.add_argument('-seed', action='store', dest='RANDOM_STATE',
                        default=42, type=int, help='Seed / trial')
    parser.add_argument('-alpha', action='store', default=0.01, type=float, 
                        help='Calibration tolerance (for metrics)')
    parser.add_argument('-n_bins', action='store', default=10, type=int, 
                        help='Number of bins to consider for calibration')
    parser.add_argument('-gamma', action='store', default=0.05, type=float, 
                        help='Min subpop prevalence (for metrics)')
    parser.add_argument('-rho', action='store', default=0.001, type=float, 
                        help='Min subpop prevalence (for metrics)')
    parser.add_argument('-ohc', action='store', default=None, type = str,
                        choices=['ohc','label_encoding','embedding'],
                        help='Specificy how text should be one-hot-encoded.')
    parser.add_argument('-groups', action='store', default='ethnicity,gender', type = str,
                        help='groups to protect')
    parser.add_argument("-text", action='store',type=str, default='chiefcomplaint', 
        help = 'Specify text features with comma seperated'
    )
    args = parser.parse_args()
    groups = args.groups.split(',')
    # import algorithm 
    print('import from','ml.'+args.ml)
    algorithm = importlib.__import__('ml.'+args.ml,
                                     globals(),
                                     locals(),
                                     ['*']
                                    )

    print('algorithm:',algorithm.est)

    if args.ohc is None:
        ohc = 'ohc' if 'lr' in args.ml else 'label_encoding'
    else:
        ohc = args.ohc 

    # optional keyword arguments passed to evaluate
    eval_kwargs = {}
    if 'eval_kwargs' in dir(algorithm):
        eval_kwargs = algorithm.eval_kwargs

    # check for conflicts btw cmd line args and eval_kwargs
    # if args.SKIP_TUNE:
    #     eval_kwargs['skip_tuning'] = True
    evaluate_model(
        dataset=args.file, 
        results_path=args.RDIR,
        emb_path = args.EMB,
        random_state=args.RANDOM_STATE,
        ml=args.ml,
        est=algorithm.est, 
        alpha=args.alpha,
        n_bins=args.n_bins,
        gamma=args.gamma,
        rho=args.rho,
        text_features=args.text.split(','),
        one_hot_encoded=ohc,
        groups=groups,
        **eval_kwargs
    )
