# Import necessary libraries/functions

# Data wrangling and computation
import pandas as pd
import numpy as np
# Machine learning models
from sklearn.linear_model import LogisticRegression
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier
import xgboost as xgb
from catboost import CatBoostClassifier
from sklearn import model_selection
import sklearn.metrics as metrics
# Plotting
import matplotlib.pyplot as plt
from plotnine import ggplot, aes, facet_wrap, geom_histogram, labs
import seaborn as sns

# Load in train and test sets
train = pd.read_csv('train.csv')
test = pd.read_csv('test.csv')

# Look at the first few rows of the training data
train.head()

# Check column types, see if any rows have null values
train.info()

# Check how many missing values are in the training data
train.apply(lambda x: x.isna().sum())

# Basic summary statistics for each column in train
train.drop(columns=['id', 'CustomerId', 'Surname']).describe()

# Look at means and standard deviations of different variables for the people that closed their accounts and for those who did not
train.drop(columns = ['id', 'CustomerId', 'Surname', 'Geography', 'Gender']).groupby('Exited').agg(['mean', 'std']).round(2)

# Make a df for plotting, change "Exited" to a "category" type to assist with plotting
train_plot = train.astype({'Exited': 'category'})

# Plot quantitative variables
(
    # Reshape the data to facilitate plotting with ggplot
    pd.melt(train_plot[['CreditScore', 'Age' ,'Balance', 'EstimatedSalary', 'Exited']], id_vars = 'Exited')
    >> ggplot() +
        geom_histogram(aes('value', fill = 'Exited')) +
        facet_wrap('variable', scales = 'free') +
        labs(x = 'Category', y = 'Count', title = 'Histograms of Quantitative Variables')
)

# Plot categorical variables and quantitative variables with few categories (NumOfProducts, Tenure)
(
    pd.melt(train_plot[['Gender', 'Geography', 'HasCrCard', 'IsActiveMember', 'NumOfProducts', 'Exited', 'Tenure']], id_vars = 'Exited')
    >> ggplot() +
        geom_histogram(aes('value', fill = 'Exited'), binwidth = .5) +
        facet_wrap('variable', scales = 'free') +
        labs(x = 'Value', y = 'Count', title = 'Histograms of Categorical Variables', subtitle = '(and Quantitative Variables with Few Categories)')
)

sns.heatmap(train.drop(columns = ['id', 'CustomerId', 'Surname', 'Geography', 'Gender', 'HasCrCard', 'IsActiveMember', 'Exited']).corr(), annot = True)
plt.title('Correlation Plot')
plt.show()

train['EstimatedSalary2'] = train['EstimatedSalary']**2
train['CreditScore2'] = train['CreditScore']**2
train['Balance2'] = train['Balance']**2

# The new variables are now part of the training data
train.head()

# Import additional functions for preprocessing
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import StandardScaler

# Creating the X_train matrix:

# Remove columns that are not predictors
X_train = train.drop(columns=['id', 'CustomerId','Surname', 'Exited'])

# Separate out the categorical columns for one hot encoding
X_train_encoding = X_train[['Geography', 'Gender', 'HasCrCard', 'IsActiveMember']]

# Create a matrix of the one hot encoded categorical columns
encoder = OneHotEncoder(sparse_output = False)
X_train_encoded = encoder.fit_transform(X_train_encoding)

# Check the names of the columns that were created via one hot encoding
encoder.get_feature_names_out()

# The encoded matrix is currently a numpy array. Change this to a data frame
X_train_encoded = pd.DataFrame(X_train_encoded, columns = encoder.get_feature_names_out())

# Combine the numerical columns with the categorical columns to create a single X_train data frame
X_train = pd.concat([X_train.drop(columns = ['Geography', 'Gender', 'HasCrCard', 'IsActiveMember']).reset_index(drop = True), X_train_encoded], axis = 1)

# Make sure everything looks correct
X_train.head()

y_train = train['Exited']

X_train.head()

# Initialize the logistic regression model (L1 regularization)
log_reg = LogisticRegression(solver = 'liblinear', penalty = 'l1')

# Set up 10-fold cross-validation so that we use the same 10 folds to evaluate all models.
kfold = model_selection.KFold(n_splits = 10, shuffle = True, random_state = 24)

# Perform 10-fold CV to evaluate logistic regression model performance
log_reg_auc = model_selection.cross_val_score(log_reg, X_train, y_train, cv = kfold, scoring = 'roc_auc')

# Logistic regression achieves a ROC AUC of about 81.8%
np.mean(log_reg_auc)

knn = KNeighborsClassifier()
# Set up a grid of hyperparameters to tune
knn_parameters = {'n_neighbors': [500, 1000, 2500, 5000],
                  'weights': ['uniform', 'distance']}

# Perform 10-fold cross-validation to obtain ROC AUC scores for each combination of hyperparameters
knn_cv = model_selection.GridSearchCV(knn, knn_parameters, scoring = 'roc_auc', cv = kfold, n_jobs = 4) # use 4 cores in parallel to expedite tuning
knn_cv.fit(X_train, y_train)

# Make a dataframe to display the roc auc score for each combination of hyperparameters (sorted from best to worst)
pd.concat([pd.DataFrame(knn_cv.cv_results_['params']),
           pd.DataFrame(knn_cv.cv_results_['mean_test_score'],
                        columns = ['roc auc'])], axis = 1).sort_values('roc auc', ascending = False).head(5)

# The best KNN model uses n_neighbors = 500, weights = uniform, and achieves a ROC AUC score of 60.1%

svc = SVC()
svc_parameters = [
    {'kernel': ['linear'], 'C': [0.1, 1, 10]},
    {'kernel': ['rbf'], 'C': [0.1, 1, 10], 'gamma': [0.1, 1, 10]}
]

svc_cv = model_selection.GridSearchCV(svc, svc_parameters, scoring = 'roc_auc', cv = kfold, n_jobs = 4)

# SVC is sensitive to the scale of the data, so let's first scale our continuous predictors
scaler = StandardScaler()

# Select the continuous features
X_train_continuous = X_train[['CreditScore', 'Age', 'Tenure', 'Balance',
                              'NumOfProducts', 'EstimatedSalary', 'EstimatedSalary2', 'CreditScore2', 'Balance2']]

# Apply scaling, turn into a dataframe
X_train_continuous_scaled = pd.DataFrame(scaler.fit_transform(X_train_continuous),
                                         columns = ['CreditScore', 'Age', 'Tenure', 'Balance',
                                                    'NumOfProducts', 'EstimatedSalary', 'EstimatedSalary2', 'CreditScore2', 'Balance2'])

# Combine all columns together to create the final training data
X_train_svc = pd.concat([X_train_continuous_scaled,
                         X_train.drop(columns = ['CreditScore', 'Age', 'Tenure', 'Balance',
                                                 'NumOfProducts', 'EstimatedSalary', 'EstimatedSalary2', 'CreditScore2', 'Balance2'])], axis = 1)

# Fit the SVC
svc_cv.fit(X_train_svc, y_train)

pd.concat([pd.DataFrame(svc_cv.cv_results_['params']),
           pd.DataFrame(svc_cv.cv_results_['mean_test_score'],
                        columns = ['roc auc'])], axis = 1).sort_values('roc auc', ascending = False).head(5)

# The best SVM model uses C = 0.1, kernel = rbf, gamma = 0.1
# and achieves a ROC AUC score of 82.5%

rf = RandomForestClassifier()
rf_parameters = {'n_estimators': [500, 1000], 'max_features': [2, 4, 6], 'max_depth': [6, 9, 12]}

rf_cv = model_selection.GridSearchCV(rf, rf_parameters, scoring = 'roc_auc', cv = kfold, n_jobs = 4)

rf_cv.fit(X_train, y_train)

pd.concat([pd.DataFrame(rf_cv.cv_results_['params']),
           pd.DataFrame(rf_cv.cv_results_['mean_test_score'],
                        columns = ['roc auc'])], axis = 1).sort_values('roc auc', ascending = False).head(5)

# The best RF model uses max_depth = 12, max_features = 6, and n_estimators = 1000
# and achieves a ROC AUC score of 88.8%

xgb_clf = xgb.XGBClassifier()
xgb_parameters = {'learning_rate': [0.05, 0.1, 0.15, 0.3], 'max_depth': [5, 6, 7], 'colsample_bytree': [0.25, 0.5, 1]}

xgb_cv = model_selection.GridSearchCV(xgb_clf, xgb_parameters, scoring = 'roc_auc', cv = kfold, n_jobs = 4)

xgb_cv.fit(X_train, y_train)

pd.concat([pd.DataFrame(xgb_cv.cv_results_['params']),
           pd.DataFrame(xgb_cv.cv_results_['mean_test_score'],
                        columns = ['roc auc'])], axis = 1).sort_values('roc auc', ascending = False).head(5)

# The best XGBoost model uses colsample_bytree = 0.5, learning_rate = 0.15, and max_depth = 5
# and achieves a ROC AUC score of about 89.0%

cbc_clf = CatBoostClassifier(cat_features = [1, 2, 6, 7, 8], od_type = "Iter", od_wait = 20)
cbc_parameters = {'learning_rate': [0.1, 0.2, 0.3], 'depth': [4, 6, 8], 'iterations': [50, 100, 150]}

cbc_cv = model_selection.GridSearchCV(cbc_clf, cbc_parameters, scoring = 'roc_auc', cv = kfold, n_jobs = 4)

# CatBoost does not require one hot encoding of categorical variables, so we'll use the original training dataset here,
# changing some columns we'll treat as categories to "string" type to be processed properly by the model
catboost_train = train.drop(columns = ['id', 'CustomerId', 'Surname', 'Exited'])
catboost_train = catboost_train.astype({'NumOfProducts': 'string', 'HasCrCard': 'string', 'IsActiveMember': 'string'})


cbc_cv.fit(catboost_train, y_train)

pd.concat([pd.DataFrame(cbc_cv.cv_results_['params']),
           pd.DataFrame(cbc_cv.cv_results_['mean_test_score'],
                        columns = ['roc auc'])], axis = 1).sort_values('roc auc', ascending = False).head(5)

# The best catboost model uses depth = 6, iterations = 150, and learning_rate = 0.2,
# and achieves a ROC AUC score of just under 89.0%

from xgboost import plot_importance
plot_importance(xgb_cv.best_estimator_)
plt.show()

# Add engineered features to testing data
test['EstimatedSalary2'] = test['EstimatedSalary']**2
test['CreditScore2'] = test['CreditScore']**2
test['Balance2'] = test['Balance']**2

# Next, we need to one hot encode the categorical variables in the testing data
X_test = test.drop(columns=['id', 'CustomerId','Surname'])

# Separate out the categorical columns for one hot encoding, apply one hot encoding
X_test_encoding = X_test[['Geography', 'Gender', 'HasCrCard', 'IsActiveMember']]
X_test_encoded = encoder.transform(X_test_encoding)

# The encoded matrix is currently a numpy array. Change this to a data frame
X_test_encoded = pd.DataFrame(X_test_encoded, columns = encoder.get_feature_names_out())

# Combine the numerical columns with the categorical columns to create a single X_test data frame
X_test = pd.concat([X_test.drop(columns = ['Geography', 'Gender', 'HasCrCard', 'IsActiveMember']).reset_index(drop = True), X_test_encoded], axis = 1)

# Make sure everything looks correct
X_test.head()

# Obtain predictions
final_predictions = pd.DataFrame(dict(id = test['id'],
                  Exited = xgb_cv.best_estimator_.predict_proba(X_test)[:, 1]))

# Look at the first few predictions
final_predictions.head()

# Save predictions to .csv
final_predictions.to_csv('churn_predictions.csv', index = False)