import numpy as np
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import query_method
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import OrdinalEncoder
from sklearn import preprocessing
from dataPrep import preparing

randomSeed = 3567
np.random.seed(randomSeed)

# Retrieving and preparing data
koguAndmestik = query_method.alusAndmed()
X = koguAndmestik.drop(columns = 'defaultedInYear')
y = koguAndmestik['defaultedInYear']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=randomSeed, stratify=y)
X_train.reset_index(inplace=True)
X_test.reset_index(inplace=True)
y_train.reset_index(drop=True, inplace=True)
y_test.reset_index(drop=True, inplace=True)
enc = OrdinalEncoder(categories=ordinal_categories, handle_unknown='use_encoded_value', unknown_value=-1)
enc.fit(X_train[ordinals])
X_train[ordinals] = enc.transform(X_train[ordinals]).astype(np.int32)
X_test[ordinals] = enc.transform(X_test[ordinals])
one_hot = OneHotEncoder(sparse=False)
one_hot.fit(X[oneHotEnc]) 
df_train = pd.DataFrame(one_hot.transform(X_train[oneHotEnc]),  columns=one_hot.get_feature_names_out(X[oneHotEnc].columns))
X_train_pre = df_train.merge(X_train, left_index=True, right_index=True).drop(columns=oneHotEnc)
df_test = pd.DataFrame(one_hot.transform(X_test[oneHotEnc]), columns=one_hot.get_feature_names_out(X[oneHotEnc].columns))
X_test_pre = df_test.merge(X_test, left_index=True, right_index=True).drop(columns=oneHotEnc)
scaler = preprocessing.StandardScaler().fit(X_train_pre[regular_columns])
X_train_temp = pd.DataFrame(scaler.transform(X_train_pre[regular_columns]), columns=regular_columns)
X_test_temp = pd.DataFrame(scaler.transform(X_test_pre[regular_columns]), columns=regular_columns)
X_train_preped = pd.merge(X_train_temp, X_train_pre[one_hot_encoded_columns], left_index=True, right_index=True)
X_test_preped = pd.merge(X_test_temp, X_test_pre[one_hot_encoded_columns], left_index=True, right_index=True)

# Visualization of prepared data
sns.set_style('white')
with plt.rc_context({'figure.figsize': (5.7, 4.27)}):
    ax = np.sum(X_train_viz[one_hot_encoded_columns_x]).plot(kind='bar')
    ax.set_ylabel('Näidete arv', fontsize=16)
    ax.set_xlabel('Tunnus', fontsize=16)
    plt.tight_layout()
with plt.rc_context({'figure.figsize': (5.7, 4.27)}):
    ax = sns.histplot(X_train_viz[regular_numerical_columns_x], element="poly", stat='count', log_scale=True)
    ax.set_ylabel('Näidete arv', fontsize=16)
    ax.set_xlabel('Skaleeritud väärtused', fontsize=16)
    plt.show()
    
# Baseline (logistic regression)
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_validate, cross_val_predict
from sklearn.metrics import make_scorer, precision_recall_curve, auc, roc_auc_score, confusion_matrix
from validateResults import validation

logBaseline = LogisticRegression(random_state=randomSeed)
logBaseline.fit(X_train_preped, y_train)
baselineResults = validation(logBaseline, X_test_preped, y_test, estimator_name='LogisticRegression baseline')

# model size
import pickle
import sys
p = pickle.dumps(logBaseline)
print('Size in bytes: {0}'.format(sys.getsizeof(p)))

# Dealing with imbalance
from imblearn.over_sampling import SMOTE

sm = SMOTE(random_state=randomSeed)  
X_over_sampled_train, y_over_sampled_train = sm.fit_resample(X_train_preped, y_train)
X_train_sampled, X_val_sampled, y_train_sampled, y_val_sampled = train_test_split(X_over_sampled_train, y_over_sampled_train, test_size=0.2, random_state=randomSeed)

# Other model baselines
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.neural_network import MLPClassifier
import xgboost as xgb

model_names = ['svc', 'dtc', 'rfc', 'mlp', 'xgb']
models = {}
models['svc'] = SVC(probability=True, random_state=randomSeed}).fit(X_train_preped, y_train)
models['dtc'] = DecisionTreeClassifier(random_state=randomSeed).fit(X_train_preped, y_train)
models['rfc'] = RandomForestClassifier(random_state=randomSeed).fit(X_train_preped, y_train)
models['mlp'] = MLPClassifier(random_state=randomSeed).fit(X_train_preped, y_train)
models['xgb'] = xgb.XGBClassifier(random_state=randomSeed).fit(X_train_preped, y_train)

results = {}
for i, model in enumerate(models):
    results[model] = validation(models[model], X_test_preped, y_test, estimator_name=model, visualize=False)
tulemus = pd.DataFrame.from_dict(results).T
tulemus.reset_index(inplace=True)
hindamiseks = ['roc_auc', 'precision_recall_auc']
fig, ax = plt.subplots(1, 2, figsize= (20, 5))
for variable, subplot in zip(hindamiseks, ax.flatten()):
    pilt = sns.barplot(x=tulemus[variable], y='index', data = tulemus, ax=subplot, palette="Blues_d")
    pilt.axvline(baselineResults[variable], color='r')
    pilt.set(xlim = (min(tulemus[variable])-0.05, max(tulemus[variable])+0.05))
plt.show()

# Hyperparameter search
from sklearn.ensemble import RandomForestClassifier
from sklearn.experimental import enable_halving_search_cv  
from sklearn.model_selection import HalvingGridSearchCV

model_names = ['log', 'svc', 'dtc', 'rfc', 'mlp', 'xgb']
models_rs = {}
models_rs['log'] = LogisticRegression(random_state=randomSeed)
models_rs['svc'] = SVC(probability=True, random_state=randomSeed)
models_rs['dtc'] = DecisionTreeClassifier(random_state=randomSeed)
models_rs['rfc'] = RandomForestClassifier(random_state=randomSeed)
models_rs['mlp'] = MLPClassifier(random_state=randomSeed)
models_rs['xgb'] = xgb.XGBClassifier(random_state=randomSeed)

log_params_dist = {"penalty": ['none', 'l2'],
                   "class_weight": [None, {0: 10, 1: 1}, {0: 0.1, 1: 0.9}, {0: 1, 1: 100}, {0: 1, 1: 1000}]}
svc_params_dist = {"kernel": ['rbf'],  
                   "class_weight": [None, {0: 0.1, 1: 0.9}, {0: 1, 1: 100}, {0: 1, 1: 1000}]}
dt_params_dist = {"max_depth": [None, 3, 5, 7],
                  'min_samples_split': [2, 4, 6, 8, 11],
                  "splitter": ['best', 'random']}
rf_params_dist = {"max_depth": [3, 5, 7, None],
                  'min_samples_split': [2, 4, 6, 8, 11]}
mlp_params_dist = {"hidden_layer_sizes": [(256, 128), (256, 256), (256, 256, 128), (128, 64), (64, 64)],
                   "learning_rate": ['constant', 'invscaling', 'adaptive'],
                   "activation": ['logistic', 'relu'],
                   "batch_size": ['auto', 500, 1000, 2000]}
xgb_params_dist = {"max_depth": [3, 5, 7, None]}
search_models = {}
best_params = {}
params = [log_params_dist, svc_params_dist, dt_params_dist, rf_params_dist, mlp_params_dist, xgb_params_dist]
resources = ['n_samples',  # log
             'C',  #svc
             'n_samples',  #dtc
             'n_estimators',  #rfc
             'n_samples',  #mlp
             'n_estimators'  #xgb
             ]
for i, model in enumerate(models_rs):
    search_models[model] = HalvingGridSearchCV(estimator=models_rs[model],
                                               param_grid=params[i],
                                               resource=resources[i],
                                               max_resources=100,
                                               factor=10,  
                                               random_state=0,
                                               scoring='roc_auc',
                                               cv=5
                                               ).fit(X_train_preped, y_train)
    best_params[model] = search_models[model].best_params_
best_params = {}
for i, model in enumerate(models_rs):
    best_params[model] = search_models[model].best_params_
    
#Parimad mudelid
models_calib = {}
models_calib['rf'] = RandomForestClassifier().set_params(**best_params['rfc'])
models_calib['log'] = LogisticRegression()
confidences = {}
scores_calib = {}
for model in models_calib:
    models_calib[model].fit(X_train_preped, y_train)
    scores_calib[model] = models_calib[model].predict_proba(X_test_preped)[:, 1]
    confidences[model] = models_calib[model].predict_proba(X_test_preped)
# netcal kalibreerimiseks
from netcal.scaling import TemperatureScaling
from netcal.presentation import ReliabilityDiagram

X_train_cal, X_test_cal, y_train_cal, y_test_cal = train_test_split(X_train_preped, y_train, test_size=0.125, stratify=y_train)
models_rfc = RandomForestClassifier().set_params(**best_params['rfc'])
models_rfc.fit(X_train_cal, y_train_cal)
confidences_netcal = models_rfc.predict_proba(X_test_cal)
temperature = TemperatureScaling()
temperature.fit(confidences_netcal, np.array(y_test_cal))
calibrated = temperature.transform(confidences_netcal)
n_bins = 10
diagram = ReliabilityDiagram(bins=n_bins)
diagram.plot(confidences_netcal[:, 1], np.array(y_test_cal))

# Kalibreerimise joonised ja ece
from sklearn.calibration import CalibrationDisplay

fig, axes = plt.subplots(1, 2, squeeze=True, figsize=(11.7, 5.87))
n_bins = 20
n_bins2 = 5
disp = CalibrationDisplay.from_estimator(models_calib['rf'], X_test_preped, y_test, ax=axes[0], n_bins=n_bins, c='r')
disp2 = CalibrationDisplay.from_estimator(models_calib['log'], X_test_preped, y_test, ax=axes[0], n_bins=n_bins)
disp3 = CalibrationDisplay.from_estimator(models_calib['rf'], X_test_preped, y_test, ax=axes[1], strategy='quantile', n_bins=n_bins2, c='r')
disp4 = CalibrationDisplay.from_estimator(models_calib['log'], X_test_preped, y_test, ax=axes[1], strategy='quantile', n_bins=n_bins2)
axes[1].legend(['Ideaalne Kalibreeritus', 'Otsustusmets', 'Baas mudel'], fontsize=16)
axes[0].legend().remove()
axes[0].set_xlabel('Enesekindlus (confidence)', fontsize=16)
axes[1].set_xlabel('Enesekindlus (confidence)', fontsize=16)
axes[0].set_ylabel('Õigsus (accuracy)', fontsize=16)
axes[1].set_ylabel(' ')
axes[1].set_xlim((0.0, 0.05))
axes[1].set_ylim((0.0, 0.05))
axes[0].set_title('Ühtlase jaotusega 20 binni', fontsize=16)
axes[1].set_title('Kvantiili jaotusega 5 binni', fontsize=16)
plt.tight_layout()

def ece_meetod(y, proba, bins='fd'):
    import numpy as np
    bin_count, bin_edges = np.histogram(proba, bins=bins)
    n_bins = len(bin_count)
    bin_edges[0] -= 1e-8  
    bin_id = np.digitize(proba, bin_edges, right=True) - 1
    bin_ysum = np.bincount(bin_id, weights=y, minlength=n_bins)
    bin_probasum = np.bincount(bin_id, weights=proba, minlength=n_bins)
    bin_ymean = np.divide(bin_ysum, bin_count, out=np.zeros(n_bins), where=bin_count > 0)
    bin_probamean = np.divide(bin_probasum, bin_count, out=np.zeros(n_bins), where=bin_count > 0)
    ece = np.abs((bin_probamean - bin_ymean) * bin_count).sum() / len(proba)
    return ece

def mce_meetod(y, proba, bins='fd'):
    import numpy as np
    bin_count, bin_edges = np.histogram(proba, bins=bins)
    n_bins = len(bin_count)
    bin_edges[0] -= 1e-8  
    bin_id = np.digitize(proba, bin_edges, right=True) - 1
    bin_ysum = np.bincount(bin_id, weights=y, minlength=n_bins)
    bin_probasum = np.bincount(bin_id, weights=proba, minlength=n_bins)
    bin_ymean = np.divide(bin_ysum, bin_count, out=np.zeros(n_bins), where=bin_count > 0)
    bin_probamean = np.divide(bin_probasum, bin_count, out=np.zeros(n_bins), where=bin_count > 0)
    test = (bin_probamean - bin_ymean) * bin_count
    mce = np.max(np.abs((bin_probamean - bin_ymean) * bin_count))
    return mce, test

ece = ece_meetod(y_test, scores_calib['rf'], bins=10)
print(ece)
mce, tmce = mce_meetod(y_test, scores_calib['rf'], bins=10)
print(mce)
print(tmce)
ece = ece_meetod(y_test, scores_calib['log'], bins=10)
print(ece)
mce, tmce = mce_meetod(y_test, scores_calib['log'], bins=10)
print(mce)
print(tmce)

# Partial dependence plot
from sklearn.inspection import PartialDependenceDisplay

fig, axes = plt.subplots(int(len(idx) / 2), 2, squeeze=True, figsize=(11.7, 50.87))
for i in range(int(len(idx) / 2)):
    PartialDependenceDisplay.from_estimator(models_calib['rf'], X=X_train_preped, features=[i], feature_names=idx,
                                            ax=axes[i][0], line_kw={"color": "red"})
    PartialDependenceDisplay.from_estimator(models_calib['log'], X=X_train_preped, features=[i], feature_names=idx, ax=axes[i][1])
plt.tight_layout()
fig, axes = plt.subplots(int(len(idx) / 2), 2, squeeze=True, figsize=(11.7, 50.87))
j = 0
for i in range(int(len(idx) / 2), len(idx)):
    PartialDependenceDisplay.from_estimator(models_calib['rf'], X=X_train_preped, features=[i], feature_names=idx,
                                            ax=axes[j][0], line_kw={"color": "red"})
    PartialDependenceDisplay.from_estimator(models_calib['log'], X=X_train_preped, features=[i], feature_names=idx, ax=axes[j][1])
    j += 1
plt.tight_layout()

# LIME 
from lime import lime_tabular

explainer = lime_tabular.LimeTabularExplainer(training_data = X_train_preped.to_numpy(),
                                               feature_names = idx,
                                               class_names=['Ei', 'Jah'],
                                               categorical_features=categorical_columns,
                                               categorical_names=['True', 'False']
                                               )
contractToLookAt = 450 
explain1 = explainer.explain_instance(data_row=X_test_preped.to_numpy()[contractToLookAt],
                                       predict_fn=models_calib['log'].predict_proba)
explain2 = explainer.explain_instance(data_row=X_test_preped.to_numpy()[contractToLookAt],
                                       predict_fn=models_calib['rf'].predict_proba)
explain_df1 = pd.DataFrame(explain1.as_list())
explain_df2 = pd.DataFrame(explain2.as_list())
with plt.rc_context({'figure.figsize': (5.2, 6.2)}):
    fig, (ax1, ax2) = plt.subplots(2, 1)
    pilt1 = sns.barplot(y=explain_df1[0], x=explain_df1[1], palette="Blues_d", ax=ax1)
    pilt1.set(title='Otsustusmetsa tunnuste mõju')
    pilt1.set_ylabel('Tunnused')
    pilt1.set_xlabel(' ')
    pilt2 = sns.barplot(y=explain_df2[0], x=explain_df2[1], palette="Blues_d", ax=ax2)
    pilt2.set(title='Logistilise regressiooni tunnuste mõju ')
    pilt2.set_ylabel('Tunnused')
    pilt2.set_xlabel('Mõju')
    plt.show()

# Feature importances
from bootstrapTools import *

values1 = models_calib['log'].coef_.flatten()
df1 = pd.DataFrame(values1, idx)
df1.sort_values(0, inplace=True)
values2 = models_calib['rf'].feature_importances_
df2 = pd.DataFrame(values2, idx)
df2.sort_values(0, inplace=True)
fig, axs = plt.subplots(ncols=2)
pilt1 = sns.barplot(y=df1.index, x=df1[0], palette="Blues_d", ax=axs[0])
pilt1.set(title='Baasmudeli tunnuste koefitsendid')
pilt1.set_ylabel('Tunnused')
pilt1.set_xlabel('Koefitsent')
plt.figure(figsize=(10, 8))
pilt2 = sns.barplot(y=df2.index, x=df2[0], palette="Blues_d", ax=axs[1])
pilt2.set(title='Otsustusmetsa tunnuste olulisus')
pilt2.set_xlabel('Olulisus')
plt.figure(figsize=(10, 8))
plt.show()

# Sampling to reduce randomness, checking base run with all models, to find top performing
models_bootstrap = {}
models_bootstrap['Base (logistic regression)'] = LogisticRegression()
models_bootstrap['Logistic Regression'] = LogisticRegression().set_params(**best_params['log'])
models_bootstrap['Support Vector Machine'] = SVC(probability=True).set_params(**best_params['svc'])
models_bootstrap['Decision Tree'] = DecisionTreeClassifier().set_params(**best_params['dtc'])
models_bootstrap['Random Forest'] = RandomForestClassifier().set_params(**best_params['rfc'])
models_bootstrap['Multi Layer Perceptron'] = MLPClassifier().set_params(**best_params['mlp'])
models_bootstrap['XGboost'] = xgb.XGBClassifier().set_params(**best_params['xgb'])
finRes_AUC, finRes_PR_AUC, finRes_AUC_rl, finRes_PR_AUC_rl, finRes_fpr, finRes_tpr = runComparisonTest(X, y, models_bootstrap, kordused=10, use_SMOTE=True, risk_level=0.05)  
results = [finRes_AUC, finRes_PR_AUC, finRes_AUC_rl, finRes_PR_AUC_rl]r
for result in results:
    prepAndShow(result, baselineName='Base (logistic regression)_auc', use_label=False)

# Final run on top models
models_bootstrap_final = {}
models_bootstrap_final['Base (logistic regression)'] = LogisticRegression()
models_bootstrap_final['Logistic Regression'] = LogisticRegression().set_params(**best_params['log'])
models_bootstrap_final['Random Forest'] = RandomForestClassifier().set_params(**best_params['rfc'])
models_bootstrap_final['XGboost'] = xgb.XGBClassifier().set_params(**best_params['xgb'])
finalRes_AUC, finalRes_PR_AUC, finalRes_fpr, finalRes_tpr, finalRes_precision, finalRes_recall = runComparisonTest2(X, y, models_bootstrap_final, kordused=1000, use_SMOTE=True)

# ROC and PR AUC plots
with plt.rc_context({'figure.figsize': (11.7 / 2, 8.27 / 2)}):
    tprs = ühtlusta(finalRes_tpr['Random Forest'])
    fprs = ühtlusta(finalRes_fpr['Random Forest'])
    plt.plot(np.mean(fprs, axis=1), np.mean(tprs, axis=1), ':', label='Otsustusmets', c='r')
    tprs = ühtlusta(finalRes_tpr['Base (logistic regression)'])
    fprs = ühtlusta(finalRes_fpr['Base (logistic regression)'])
    plt.plot(np.mean(fprs, axis=1), np.mean(tprs, axis=1), ':', label='Baas mudel (logistiline regressioon)', c='b)
    plt.text(0.4, 0.4, 'Keskmine ROC AUC: ', fontsize=16)
    plt.text(0.4, 0.3, 'Otsustusmets {0:2f}'.format(np.mean(finalRes_AUC['Random Forest'])), c='r', fontsize=16)
    plt.text(0.4, 0.2, 'Baas mudel   {0:2f}'.format(np.mean(finalRes_AUC['Base (logistic regression)'])), c='b', fontsize=16)
    plt.xlabel("Vale-positiivsete määr", fontsize=16)
    plt.ylabel("Päris-positiivsete määr", fontsize=16)
    plt.show()

# Results visualization
avgResults_PR = {}
for key in finalRes_PR_AUC.keys():
    avgResults_PR[key] = np.mean(finalRes_PR_AUC[key])
tulemus_PR = pd.DataFrame.from_records(avgResults_PR, index=['values']).T
tulemus_PR.rename(tulbad, inplace=True)
tulemus_PR.sort_index(inplace=True, ascending=False)
pilt = sns.barplot(y=tulemus_PR.index, x=tulemus_PR['values'], linewidth=2.5, facecolor=(1, 1, 1, 0), errcolor=".2", edgecolor=".2")
pilt.set_ylabel('Mudelid')
pilt.set_xlabel('Keskmine PR AUC')
pilt.axvline(avgResults_PR['Base (logistic regression)'], color='r')  
pilt.bar_label(pilt.containers[0], label_type='center', fmt='%.3f', color='k')
plt.figure(figsize=(10, 8))
plt.show()
avgResults = {}  
for key in finalRes_AUC.keys():
    avgResults[key] = np.mean(finalRes_AUC[key])  
tulemus = pd.DataFrame.from_records(avgResults, index=['values']).T
tulemus.rename(tulbad, inplace=True)
tulemus.sort_index(inplace=True, ascending=False)
pilt = sns.barplot(y=tulemus.index, x=tulemus['values'], palette="Blues_d")
pilt.set_ylabel('Mudelid')
pilt.set_xlabel('Keskmine ROC AUC')
pilt.axvline(avgResults['Base (logistic regression)'], color='r')  
pilt.bar_label(pilt.containers[0], label_type='center', fmt='%.3f', color='k')
plt.figure(figsize=(10, 8))
plt.show()
tulemus = pd.DataFrame.from_records(avgResults, index=['values']).T
tulemus.rename(tulbad, inplace=True)
tulemus.sort_index(inplace=True, ascending=False)
tulemus_PR = pd.DataFrame.from_records(avgResults_PR, index=['values']).T
tulemus_PR.rename(tulbad, inplace=True)
tulemus_PR.sort_index(inplace=True, ascending=False)
ax = plt.subplots()
ax = sns.barplot(y=tulemus.index, x=tulemus['values'], palette="Blues_d")
ax = sns.barplot(y=tulemus_PR.index, x=tulemus_PR['values'], linewidth=2.5, facecolor=(1, 1, 1, 0), errcolor=".2", edgecolor=".2")
ax.set_ylabel(' ', fontsize=16)
ax.set_xlabel('Keskmine ROC AUC ja PR AUC', fontsize=16)
ax.axvline(avgResults['Base (logistic regression)'], color='r')  
ax.bar_label(pilt.containers[0], label_type='center', fmt='%.3f', color='w')  
plt.show()

# Statistical test
from scipy import stats

best_model = 'Random Forest' 
baseline_name = 'Base (logistic regression)'
s1 = pd.DataFrame(finalRes_AUC[baseline_name], index=np.arange(1, len(finalRes_AUC[baseline_name]) + 1, 1),
                  columns=[baseline_name])
s2 = pd.DataFrame(finalRes_AUC[best_model], index=np.arange(1, len(finalRes_AUC[best_model]) + 1, 1),
                  columns=[best_model])
stat_test_df = pd.merge(s1, s2, left_index=True, right_index=True)
sns.distplot(stat_test_df[baseline_name], hist=False, rug=True)
sns.distplot(stat_test_df[best_model], hist=False, rug=True)
plt.show()
print('Baseline normality test result: {0}'.format(stats.shapiro(stat_test_df[baseline_name])))
print('Best model normality test result: {0}'.format(stats.shapiro(stat_test_df[best_model])))
print('T-test result: {0}'.format(stats.ttest_rel(stat_test_df[baseline_name], stat_test_df[best_model], alternative='less')))
print('Average prediction ROC AUC scores: {0}'.format(np.mean(stat_test_df)))