# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
DF_PATH       = "../data/processed/2_cleaned_data.pkl"

ROLE_COLS      = ['DevType']
TECH_COLS      = ['LanguageWorkedWith',    'DatabaseWorkedWith',    'WebframeWorkedWith',    'MiscTechWorkedWith']

EXPERIMENT_NAME = "stackoverflow_single_model"
LOG_PATH = "../models/temp/random_forest/"
LOG_DATA_PKL    =  "data.pkl"
LOG_MODEL_PKL   =  "model.pkl"
LOG_METRICS_PKL =  "metrics.pkl"

# + tags=[]
# Load packages
import pandas as pd 
import numpy as np
import logging
import pickle
import random
import plotly 
import os

import mlflow
from mlflow.tracking import MlflowClient

from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.preprocessing import StandardScaler, RobustScaler, MinMaxScaler
from sklearn.pipeline import make_pipeline, FeatureUnion
from sklearn.feature_selection import VarianceThreshold
from sklearn.multioutput import MultiOutputClassifier
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV

from sklearn import metrics
from sklearn.metrics import auc, accuracy_score, confusion_matrix, f1_score, precision_score, recall_score

from sklearn.decomposition import PCA, KernelPCA

from sklearn import tree
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import AdaBoostClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier

from matplotlib import pyplot as plt
# -

# _________

# + [markdown] tags=[]
# ### Read data
# -

# Read Data 
df = pd.read_pickle(DF_PATH)


# + [markdown] tags=[]
# ### Functions
# -

def calculate_quality(ground_truth, prediction, metric_function, sort_values=False):
    quality_scores = {}
    for col in predictions.columns:
        role_pred  = predictions[col].copy()
        role_truth = ground_truth[col].copy()
        quality_scores[col] = round(metric_function(role_truth, role_pred) * 100, 2)
        
    quality_scores = pd.Series(quality_scores.values(), index=quality_scores.keys())
    if sort_values:
        quality_scores = quality_scores.sort_values()
    
    return quality_scores


# + [markdown] tags=[]
# ### Balance classes 
# -

# Check the total samples of roles
roles_df = df["DevType"].copy()
role_sum = df["DevType"].sum(axis=0)
role_sum

# +
# Resample roles
samples_per_class = 500
resampled_roles = []

for role_col in roles_df.columns:
    sub_df = roles_df.loc[roles_df[role_col] == 1].copy()
    
    if len(sub_df) < samples_per_class:
        # Upsample
        sub_df = sub_df.sample(samples_per_class, replace=True, random_state=0)
    else:
        # Downsample
        sub_df = sub_df.sample(samples_per_class, random_state=0) 
    
    resampled_roles.append(sub_df)
# -

# Construct dfs
roles_df  = pd.concat(resampled_roles)
df = df.loc[roles_df.index]

roles_df.sum(axis=0)

# + [markdown] tags=[]
# ### Prep
# -

X_train, X_test, Y_train, Y_test = train_test_split(df.drop("DevType", axis=1), 
                                                    df["DevType"], 
                                                    random_state=0)

X_train.shape

# _________

# + [markdown] tags=[]
# # Train models
# -

# Initialize client and experiment
client = MlflowClient()
mlflow.set_experiment(EXPERIMENT_NAME)
exp = client.get_experiment_by_name(EXPERIMENT_NAME)

# + [markdown] tags=[]
# ### 1. Vanilla Forest 

# +
rf_clf = make_pipeline(StandardScaler(),
                       RandomForestClassifier(random_state=0))

rf_clf.fit(X_train, Y_train)
# -

# Evaluate on train set
predictions =  pd.DataFrame(rf_clf.predict(X_train), columns=Y_train.columns)
train_scores = {score.__name__: calculate_quality(Y_train, predictions, score) 
                for score in [accuracy_score, precision_score, recall_score, f1_score]}
train_scores = pd.concat(train_scores,axis=1)

# Evaluate on test set
predictions =  pd.DataFrame(rf_clf.predict(X_test), columns=Y_test.columns)
test_scores = {score.__name__: calculate_quality(Y_test, predictions, score) 
                for score in [accuracy_score, precision_score, recall_score, f1_score]}
test_scores = pd.concat(test_scores,axis=1)
mean_test_scores = test_scores.mean()

print(test_scores.mean())
test_scores.sort_values("precision_score")

# + [markdown] tags=[]
# #### Log

# + tags=[]
# Data details
data_details = {"data_path": DF_PATH,
                "training_indices": X_train.index.tolist(),
                "test_indices":     X_test.index.tolist(), 
                "features_names":   X_train.columns.droplevel(0).tolist(),
                "targets_names":    Y_train.columns.tolist()}

with open(os.path.join(LOG_PATH, LOG_DATA_PKL), "wb") as output_file:
    pickle.dump(data_details, output_file)

# +
# Model
model = {"model_description": "Random Forest: with non linearity",
         "model_details": str(rf_clf),
         "model_object": rf_clf} 

with open(os.path.join(LOG_PATH, LOG_MODEL_PKL), "wb") as output_file:
    pickle.dump(model, output_file)

# +
# Preformance details
classes_metrics = {"train_scores": train_scores, 
                   "test_scores":  test_scores}

with open(os.path.join(LOG_PATH, LOG_METRICS_PKL), "wb") as output_file:
    pickle.dump(classes_metrics, output_file)
# -

# Start a new run and track 
with mlflow.start_run(experiment_id=exp.experiment_id, 
                      run_name=model["model_description"]):
    # Log pickles
    mlflow.log_artifacts(LOG_PATH)  
    
    # Track metrics 
    for metric, score in mean_test_scores.items():
        mlflow.log_metric(metric, score) 
    

# _________

# + [markdown] tags=[]
# ### 2. Random Forest with Non-linearity

# +
rf_clf = make_pipeline(StandardScaler(), 
                       FeatureUnion([('linear_pca', PCA(n_components=40)), 
                                      ('kernel_pca', KernelPCA(n_components=40, 
                                                               kernel='rbf'))]),
                       RandomForestClassifier(random_state=0))

rf_clf.fit(X_train, Y_train)
# -

# Evaluate on train set
predictions =  pd.DataFrame(rf_clf.predict(X_train), columns=Y_train.columns)
train_scores = {score.__name__: calculate_quality(Y_train, predictions, score) 
                for score in [accuracy_score, precision_score, recall_score, f1_score]}
train_scores = pd.concat(train_scores,axis=1)

# Evaluate on test set
predictions =  pd.DataFrame(rf_clf.predict(X_test), columns=Y_test.columns)
test_scores = {score.__name__: calculate_quality(Y_test, predictions, score) 
                for score in [accuracy_score, precision_score, recall_score, f1_score]}
test_scores = pd.concat(test_scores,axis=1)
mean_test_scores = test_scores.mean()

print(mean_test_scores)
test_scores.sort_values("f1_score")

# + [markdown] tags=[]
# #### Log

# + tags=[]
# Data details
data_details = {"data_path": DF_PATH,
                "training_indices": X_train.index.tolist(),
                "test_indices":     X_test.index.tolist(), 
                "features_names":   X_train.columns.droplevel(0).tolist(),
                "targets_names":    Y_train.columns.tolist()}

with open(os.path.join(LOG_PATH, LOG_DATA_PKL), "wb") as output_file:
    pickle.dump(data_details, output_file)

# +
# Model
model = {"model_description": "Random Forest: with PCAs + RBF kernel",
         "model_details": str(rf_clf),
         "model_object": rf_clf} 

with open(os.path.join(LOG_PATH, LOG_MODEL_PKL), "wb") as output_file:
    pickle.dump(model, output_file)

# +
# Preformance details
classes_metrics = {"train_scores": train_scores, 
                   "test_scores":  test_scores}

with open(os.path.join(LOG_PATH, LOG_METRICS_PKL), "wb") as output_file:
    pickle.dump(classes_metrics, output_file)
# -

# Start a new run and track 
with mlflow.start_run(experiment_id=exp.experiment_id, 
                      run_name=model["model_description"]):
    # Log pickles
    mlflow.log_artifacts(LOG_PATH)  
    
    # Track metrics 
    for metric, score in mean_test_scores.items():
        mlflow.log_metric(metric, score) 
    

# _________

# + [markdown] tags=[]
# ### 2. Random Forest with PCAs & Hyper parameter tuning
# -

hpt_rf_clf = make_pipeline(StandardScaler(), 
                           FeatureUnion([('linear_pca', PCA()), 
                                         ('kernel_pca', KernelPCA(kernel='rbf'))]),
                           RandomForestClassifier(random_state=0, ))

# list(hpt_rf_clf.get_params().keys())
tuned_parameters = [{'featureunion__linear_pca__n_components': [5, 10, 20, 40, 60], 
                     'featureunion__kernel_pca__n_components': [5, 10, 20, 40, 60],
                     'randomforestclassifier__n_estimators':   [100, 500, 1000]}]


# +
hpt_rf_clf = GridSearchCV(hpt_rf_clf, 
                          tuned_parameters, 
                          n_jobs=-1)

hpt_rf_clf.fit(X_train, Y_train)
# -

hpt_rf_clf.best_params_

# Evaluate on test set
predictions =  pd.DataFrame(hpt_rf_clf.predict(X_train), columns=Y_train.columns)
train_scores = {score.__name__: calculate_quality(Y_train, predictions, score) 
                for score in [accuracy_score, precision_score, recall_score, f1_score]}
train_scores = pd.concat(train_scores,axis=1)

# Evaluate on test set
predictions =  pd.DataFrame(hpt_rf_clf.predict(X_test), columns=Y_test.columns)
test_scores = {score.__name__: calculate_quality(Y_test, predictions, score) 
                for score in [accuracy_score, precision_score, recall_score, f1_score]}
test_scores = pd.concat(test_scores,axis=1)

print(test_scores.mean())
test_scores.sort_values("f1_score")

# + [markdown] tags=[]
# #### Log

# + tags=[]
# Data details
data_details = {"data_path": DF_PATH,
                "training_indices": X_train.index.tolist(),
                "test_indices":     X_test.index.tolist(), 
                "features_names":   X_train.columns.droplevel(0).tolist(),
                "targets_names":    Y_train.columns.tolist()}

with open(os.path.join(LOG_PATH, LOG_DATA_PKL), "wb") as output_file:
    pickle.dump(data_details, output_file)

# +
# Model
model = {"model_description": "Random Forest: with PCAs + RBF kernel + Hyperparamter tuning",
         "model_details": str(hpt_rf_clf),
         "model_object": hpt_rf_clf} 

with open(os.path.join(LOG_PATH, LOG_MODEL_PKL), "wb") as output_file:
    pickle.dump(model, output_file)

# +
# Preformance details
classes_metrics = {"train_scores": train_scores, 
                   "test_scores":  test_scores}

with open(os.path.join(LOG_PATH, LOG_METRICS_PKL), "wb") as output_file:
    pickle.dump(classes_metrics, output_file)
# -

# Start a new run and track 
with mlflow.start_run(experiment_id=exp.experiment_id, 
                      run_name=model["model_description"]):
    # Log pickles
    mlflow.log_artifacts(LOG_PATH)  
    
    # Track metrics 
    for metric, score in mean_test_scores.items():
        mlflow.log_metric(metric, score) 
    



StarcoderdataJupyterScriptsDedupFiltered