import pandas as pd
from inference import InferenceWrapper, pass_through
from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.models import Model
from tensorflow.keras import regularizers
from tensorflow.keras.optimizers import Adam
from sklearn.metrics import roc_auc_score
from IPython.display import display

import os
import torch
from torch.cuda import empty_cache
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="2"

combined_sig_df = pd.read_hdf('combined_sig_df.hdf')
feat_df = pd.read_csv('feat_df.csv')

train_mask = combined_sig_df.part != 6
holdout_mask = ~train_mask

# count the labels in the holdout set
from collections import Counter
c = Counter()

for row in combined_sig_df[combined_sig_df.part == 6].labels:
    c.update(row)

X = feat_df[train_mask].values
X.shape

label_columns = [x for x in combined_sig_df.columns if 'sig/' in x]
y = combined_sig_df[label_columns][train_mask].values
y.shape

X_holdout = feat_df[holdout_mask].values
y_holdout = combined_sig_df[label_columns][holdout_mask].values


def calculate_auc(predictions):
    auc_scores = []
    counts = []

    for i, l in enumerate(label_columns):
        y_hat = predictions[:, i]
        y = y_holdout[:, i]
        auc = roc_auc_score(y_true=y, y_score=y_hat)
        auc_scores.append(auc)
        counts.append(c[l])
    
    df = pd.DataFrame({'label': label_columns, 'auc': auc_scores, 'count': counts})    
    display(df)
    weightedavg_auc = df.apply(lambda x: x.auc * x['count'], axis=1).sum() / df['count'].sum()
    print(f'Weighted Average AUC: {weightedavg_auc}')
    return df, weightedavg_auc

"""
Explanation: Data For this notebook:
combined_sig_df.hdf
https://storage.googleapis.com/issue_label_bot/notebook_files/combined_sig_df.hdf
feat_df.csv https://storage.googleapis.com/issue_label_bot/notebook_files/feat_df.csv
End of explanation
"""


def shallow_model(l1=.01, l2=.01):

    inp = Input(shape=(1600,))
    x = Dense(units=30)(inp)
    out = Dense(units=28, activation='sigmoid')(x)
    model = Model(inputs=inp, outputs=out)
    model.compile(optimizer=Adam(lr=.001), loss='categorical_crossentropy')
    return model

shallow_model = shallow_model()

shallow_model.fit(x=X, y=y, batch_size=64, epochs=50, validation_split=.15)

shallow_model.fit(x=X, y=y, batch_size=64, epochs=1, validation_split=0)

y_hat_holdout = shallow_model.predict(X_holdout)

auc_scores = []

for i, l in enumerate(label_columns):
    y_hat = y_hat_holdout[:, i]
    y = y_holdout[:, i]
    auc = roc_auc_score(y_true=y, y_score=y_hat)
    auc_scores.append(auc)    

pd.DataFrame({'label': label_columns, 'auc': auc_scores})

"""
Explanation: Part 1: Feed Embeddings From Language Model To Downstream Algorithims and Do Greedy Training On Top of That
Keras
End of explanation
"""


from sklearn.neural_network import MLPClassifier
mlp = MLPClassifier(early_stopping=True, n_iter_no_change=5, max_iter=500, solver='adam', 
                   random_state=1234)

mlp.fit(X, y)

mlp_predictions = mlp.predict_proba(X_holdout)

mlp_df, mlp_auc = calculate_auc(mlp_predictions)

"""
Explanation: wow that sucks, lets try something else
Sklearn Nueral Network
b/c multi-label is supported naturally
End of explanation
"""


params = {'hidden_layer_sizes': [(100,), (200,), (400, ), (50, 50), (100, 100), (200, 200)],
              'alpha': [.001, .01, .1, 1, 10],
              'learning_rate': ['constant', 'adaptive'],
              'learning_rate_init': [.001, .01, .1]}
              
mlp_clf = MLPClassifier(early_stopping=True, validation_fraction=.2, n_iter_no_change=4, max_iter=500)

gscvmlp = GridSearchCV(mlp_clf, params, cv=5, n_jobs=-1)

gscvmlp.fit(X, y)

print(f'The best model from grid search is:\n=====================================\n{gscvmlp.best_estimator_}')

mlp_tuned_predictions = gscvmlp.predict_proba(X_holdout)

mlp_tuned_df, mlp_tuned_auc = calculate_auc(mlp_tuned_predictions)

"""
Explanation: Try Tuning the MLP
End of explanation
"""


from sklearn.ensemble import RandomForestClassifier
from sklearn.multiclass import OneVsRestClassifier
from sklearn.svm import NuSVC

rf = RandomForestClassifier(n_estimators=300,
                             random_state=1234, min_samples_leaf=3, class_weight='balanced')


clf = OneVsRestClassifier(rf, n_jobs=-1)

clf.fit(X, y)

rf_predictions = clf.predict_proba(X_holdout)

rf_df, rf_auc = calculate_auc(rf_predictions)

"""
Explanation: Sklearn Random Forest
End of explanation
"""


from sklearn.neighbors import KNeighborsClassifier
knn = KNeighborsClassifier(weights='distance', n_neighbors=10, n_jobs=-1)

knn.fit(X, y)

knn_preds = knn.predict_proba(X_holdout)

import numpy as np
knn_preds_stacked = np.stack([x[:, 1] for x in knn_preds], axis=0).T

knn_preds_stacked

knn_df, knn_auc = calculate_auc(knn_preds_stacked)

"""
Explanation: Sklearn KNN
End of explanation
"""


from sklearn.ensemble import GradientBoostingClassifier
from sklearn.multiclass import OneVsRestClassifier
gbm = GradientBoostingClassifier(max_depth=5, min_samples_leaf=3, max_features='auto', n_iter_no_change=4)
clf = OneVsRestClassifier(gbm, n_jobs=-1)

clf.fit(X, y)

gbm_predictions = clf.predict_proba(X_holdout)

gbm_df, gbm_auc = calculate_auc(gbm_predictions)

"""
Explanation: Sklearn GBM
End of explanation
"""


import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="2"

from inference import InferenceWrapper, pass_through
iw = InferenceWrapper(model_path='/ds/multi_class_model/',
                      model_file_name='/ds/multi_class_model/export.hdf')

#because the pre-trained mdoel predicts more than just sig/ labels
pred_mask = [x in label_columns for x in iw.learn.data.classes]

"""
Explanation: Part 2: Fine Tune Original Language Model
With FastAI & Pytorch
Note: this model was completely trained in another notebook, it is only evaluated here.
Prepare Inference Wrapper
Model is available for download here:
https://storage.googleapis.com/issue_label_bot/model/multi_class_model/export.hdf
End of explanation
"""


parsed_df = iw.process_df(combined_sig_df)

holdout_text = parsed_df[holdout_mask]

lang_model_predict = np.stack(holdout_text.text.apply(lambda x: tcl.predict(x)[2].numpy()[pred_mask]).values)

lang_model_predict.shape

len(iw.learn.data.classes)

np.array(iw.learn.data.classes)[pred_mask]

lang_model_predict_df = pd.DataFrame(lang_model_predict)
lang_model_predict_df.columns = np.array(iw.learn.data.classes)[pred_mask]
lm_df = lang_model_predict_df[[x for x in label_columns if x in lang_model_predict_df.columns]]

"""
Explanation: Do pre-processing (such as markdown parsing) to prepare data for model.
End of explanation
"""


missing_cols = [x for x in label_columns if x not in lm_df.columns]

for col in missing_cols:
    lm_df[col] = 0.0

lm_df, lm_auc = calculate_auc(lm_df.values)

missing_cols

"""
Explanation: There are some columns that the model was not trained on
End of explanation
"""


from sklearn.feature_extraction.text import CountVectorizer
from sklearn.pipeline import make_pipeline, FeatureUnion
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.neural_network import MLPClassifier
from sklearn.model_selection import GridSearchCV


train_mask = combined_sig_df.part != 6

parameters = {'hidden_layer_sizes': [(50,), (100,), (200,), (400, ), (50, 50), (100, 100)],
              'alpha': [0.0001, .001, .01, .1, 1, 10],
              'learning_rate': ['constant', 'adaptive'],
              'learning_rate_init': [.001, .01]}
              
mlp = MLPClassifier(early_stopping=True, validation_fraction=.2, n_iter_no_change=4, max_iter=500)

class ColumnSelector(BaseEstimator, TransformerMixin):
    def __init__(self, column):
        self.column = column

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        assert isinstance(X, pd.DataFrame)

        try:
            return X[self.column].values
        except KeyError:
            col_error = list(set(self.column) - set(X.column))
            raise KeyError("The DataFrame does not include the columns: %s" % col_error)

model_pipeline = make_pipeline(
    FeatureUnion(transformer_list=[
        ("title", make_pipeline(
            ColumnSelector("title"),
            CountVectorizer(ngram_range=(1,3), min_df=3, strip_accents='unicode'),
        )),
        ("body", make_pipeline(
            ColumnSelector("body"),
            CountVectorizer(ngram_range=(1,3), min_df=3, strip_accents='unicode'),
        ))
    ]),
    GridSearchCV(mlp, parameters, cv=5, n_jobs=-1)
)

y_train = combined_sig_df[label_columns][train_mask].values
x_train_df = combined_sig_df[['body', 'title']][train_mask]

"""
Explanation: Part 3:  Using Text & Bag of Words Instead of Pre-Trained Embeddings (Classic ML)
Count Vectorizer w/ Nueral Net
Choose Nueral Net because supports multi-label classification natively
End of explanation
"""


model_pipeline.fit(x_train_df, y_train)

best_model = model_pipeline.named_steps['gridsearchcv'].best_estimator_
print(f'The best model from grid search is:\n=====================================\n{best_model}')


"""
Explanation: Fit the model
End of explanation
"""


holdout_mask = combined_sig_df.part == 6
x_holdout_df = combined_sig_df[['body', 'title']][holdout_mask]
y_holdout = combined_sig_df[label_columns][holdout_mask].values

no_deep_mlp_preds = model_pipeline.predict_proba(x_holdout_df)

no_deep_mlp_df, no_deep_mlp_auc = calculate_auc(no_deep_mlp_preds)

"""
Explanation: Evaluate Model
End of explanation
"""


