import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestRegressor
from sklearn.neighbors import KNeighborsRegressor
from sklearn.svm import SVR
from sklearn.model_selection import GridSearchCV
from sklearn.model_selection import RandomizedSearchCV
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score

sns.set_style('darkgrid')

# Data extraction.  
edu_df = pd.read_csv('hdi-full.csv', index_col=0)

# Keep columns with at least 161 non-null values. 
edu_df = edu_df.dropna(axis=1, thresh=161)
edu_df.shape

# Keep rows with at least 607 non-null values.
edu_df = edu_df.dropna(axis=0, thresh=607)
edu_df.shape

# The null values in each column. 
edu_df.isnull().sum().sort_values(ascending=True)

# The null values in each row.
print(dict(edu_df.isnull().sum(axis=1).sort_values(ascending=False)))

# Basic imputation. 
edu_df = edu_df.fillna(edu_df.mean(numeric_only=True))

# The null values in each row.
print(dict(edu_df.isnull().sum(axis=1).sort_values(ascending=False)))

# Encoding an ordinal categorical variable.
scale_mapper = {"Low" : 1, 
                "Medium" : 2, 
                "High" : 3, 
                "Very High" : 4}
edu_df["Human Development Groups"] = edu_df["Human Development Groups"].replace(scale_mapper)

# The columns we don't want to include as features.
removals = ["Country", "UNDP Developing Regions"]
for i in range(1990, 2022):
    removals.append("Human Development Index (" + str(i) + ")")
    removals.append("HDI male (" + str(i) + ")")
    removals.append("HDI female (" + str(i) + ")")
    removals.append("Expected Years of Schooling, female (" + str(i) + ")")
    removals.append("Expected Years of Schooling, male (" + str(i) + ")")
    removals.append("Expected Years of Schooling (" + str(i) + ")")
    removals.append("Mean Years of Schooling, female (" + str(i) + ")")
    removals.append("Mean Years of Schooling, male (" + str(i) + ")")
    removals.append("Mean Years of Schooling (" + str(i) + ")")
    removals.append("Planetary pressures-adjusted Human Development Index (" + str(i) + ")")
    removals.append("HDI Rank (" + str(i) + ")")

# Analyzing distribution of our target.
male = 'Mean Years of Schooling, male (2021)'
print("The standard deviation of the Mean Years of Schooling, female (2021) is {}".format(round(edu_df[target].std(), 2)))
print("The mean of the Mean Years of Schooling, female (2021) is {}".format(round(edu_df[target].mean(), 2)))
print("The standard deviation Mean Years of Schooling, male (2021) is {}".format(round(edu_df[male].std(), 2)))
print("The mean of the Mean Years of Schooling, male (2021) is {}".format(round(edu_df[male].mean(), 2)))

# Analyzing the distribution of males vs. females. 
edu_df[target].plot(kind='kde', label="Female")
edu_df[male].plot(kind='kde', label="Male")
plt.title("Distribution of Mean Years of Schooling for Males and Females (2021)")
plt.legend(loc="upper left")
plt.show()

# Visualizing a potential correlation. 
sns.regplot(data=edu_df, x='Human Development Groups', y='Mean Years of Schooling, female (2021)')
plt.title("Correlation between Schooling and the Prevalence of HD Groups")
plt.show()

sns.regplot(data=edu_df, x='Labour force participation rate, male (% ages 15 and older) (2021)', y='Mean Years of Schooling (2021)')
plt.title("Correlation between Schooling and Male Labour Force Participation Rate")
plt.show()

sns.regplot(data=edu_df, x='Labour force participation rate, female (% ages 15 and older) (2021)', y='Mean Years of Schooling (2021)')
plt.title("Correlation between Schooling and Female Labour Force Participation Rate")
plt.show()

sns.regplot(data=edu_df, y='Mean Years of Schooling (2021)', x='Gross National Income Per Capita (2021)')
plt.title("Correlation between Schooling and Gross National Income per Capita")
plt.show()

# Partitioning the data into test and train, using a 75/25 split.
X = edu_df[features]
y = edu_df[target]

X_train, X_test, y_train, y_test = train_test_split(X,  #the input features
                                                    y,  #the label
                                                    test_size=0.25,  #set aside 30% of the data as the test set
                                                    random_state=5  #reproduce the results
                                                    )

# Testing our model without tuning.
knn = KNeighborsRegressor()
knn.fit(X_train, y_train)

# Predicting the labels for the test set.
y_pred = knn.predict(X_test)

mse = mean_squared_error(y_test, y_pred)

# Evaluate the Predictions
print('The STD of the target is: {}'.format(round(edu_df[target].std(), 2)))
print('The mse of the model is: {}'.format(round(mse, 2)))
print('The R-squared is: ', round(r2_score(y_test, y_pred), 2))

# Finding the best k-value.
k_values = range(3, 50)
scores = []

best_mse = (float('inf'), -1)
for k in k_values:
    knn = KNeighborsRegressor(n_neighbors=k)
    knn.fit(X_train, y_train)
    score = knn.score(X_test, y_test)
    scores.append(score)
    y_pred = knn.predict(X_test)
    mse = mean_squared_error(y_test, y_pred)

    if mse <= best_mse[0]:
        best_mse = (mse, k)

# Visualizing the different k-values and their scores.
x = np.array(k_values)
y = np.array(scores)

plt.title("K-values and their Scores")
plt.xlabel("K-value")
plt.ylabel("Score")
plt.plot(x, y, color="red")
plt.show()

# Explain the visualization.
print(f"The model with the best mse had a n_neighbors value of {best_mse[1]} with a mse of {(round(best_mse[0], 2))}.")

# Building our model with the best hyperparameters.
knn = KNeighborsRegressor(n_neighbors=10)
knn.fit(X_train, y_train)

# Predicting the labels for the test set.
y_pred = knn.predict(X_test)

mse = mean_squared_error(y_test, y_pred)

# Evaluate the Predictions
print('The STD of the target is: {}'.format(round(edu_df[target].std(), 2)))
print('The mse of the model is: {}'.format(round(mse, 2)))
print('The R-squared is: ', round(r2_score(y_test, y_pred), 2))

# Initialize our scaler, so we can normalize the numeric features. 
scaler = StandardScaler()
scaler.fit(X_train)

# Normalizing. 
X_train_scaled = scaler.transform(X_train)

# Scaling the test data.
X_test_scaled = scaler.transform(X_test)

# Building our model with the default hyperparamters.
svr = SVR()
svr.fit(X_train_scaled, y_train)

# Predicting the labels for the test set.
y_pred = svr.predict(X_test_scaled)

mse = mean_squared_error(y_test, y_pred)

# Evaluate the Predictions.
print('The mse of the model is: {}'.format(round(mse, 2)))

# Try values of C ranging from 1 to 50 and different gamma. 
params = {'C': np.arange(1, 50),
         'gamma' : ["scale", "auto"]}

svr = SVR()

# Initialize Random Search with Cross Validation.
svr_grid = RandomizedSearchCV(estimator=svr, param_distributions=params, cv=5, verbose=1)
svr_grid.fit(X_train_scaled, y_train)

print(f"Best Params: {svr_grid.best_params_}")
print("Our best score is {}.".format(round(svr_grid.best_score_, 2)))

# Predicting the labels for the test set.
y_pred = svr_grid.predict(X_test_scaled)

mse = mean_squared_error(y_test, y_pred)

# Evaluate the Predictions
print('The mse of the model is: {}.'.format(round(mse, 2)))

# Building our model with the best params.
svr = SVR(C=12, gamma='scale')
svr.fit(X_train_scaled, y_train)

# Predicting the labels for the test set.
y_pred = svr.predict(X_test_scaled)

mse = mean_squared_error(y_test, y_pred)

# Evaluate the Predictions
print('The STD of the target is: {}'.format(round(edu_df[target].std(), 2)))
print('The mse of the model is: {}'.format(round(mse, 2)))
print('The R-squared is: ', round(r2_score(y_test, y_pred), 2))

# Building our model with the default hyperparameters.
rf = RandomForestRegressor()
rf.fit(X_train, y_train)

# Predicting the labels for the test set.
y_pred = rf.predict(X_test)

mse = mean_squared_error(y_test, y_pred)

# Evaluate the Predictions.
print('The STD of the target is: {}'.format(round(edu_df[target].std(), 2)))
print('The mse of the model is: {}'.format(round(mse, 2)))
print('The R-squared is: ', round(r2_score(y_test, y_pred), 2))

# Uncomment to try tuning the hyperparameters.
# Setup the parameter grid to try values various values.
n_estimators = [5, 20]  # Number of trees.
max_features = [1.0]  # Number of features to consider at each split.
max_depth = np.arange(50, 70, 2)  # Maximum height of each decision tree.
min_samples_split = np.arange(2, 12, 1)  # Minimum sample number to split a node.
min_samples_leaf = np.arange(2, 10, 2)  # Minimum sample number that can be stored in a leaf node.
bootstrap = [True, False]  # Method used to sample data points.

params = {'n_estimators': n_estimators,
          'max_features': max_features,
          'max_depth': max_depth,
          'min_samples_split': min_samples_split,
          'min_samples_leaf': min_samples_leaf,
          'bootstrap': bootstrap}

# Initialize Grid Search with Cross Validation.
rf = RandomForestRegressor(random_state=5)

rf_grid = GridSearchCV(estimator=rf, param_grid=params, cv=4, verbose=1, n_jobs=-1, return_train_score=True)
rf_grid.fit(X_train, y_train)

print(rf_grid.best_params_)  #print the best parameters
print(round(rf_grid.best_score_, 2))  #print the best score

mse = mean_squared_error(y_test, rf_grid.predict(X_test))

# Evaluate the Predictions
print('The mse of the model is: {}'.format(round(mse, 2)))

# Building our model with the above best params.
rf = RandomForestRegressor(random_state=5, n_estimators=20, min_samples_split=2, min_samples_leaf=2, max_features=1.0,
                           max_depth=50, bootstrap=True)
rf.fit(X_train, y_train)

# Predicting the labels for the test set.
y_pred = rf.predict(X_test)

mse = mean_squared_error(y_test, y_pred)

# Evaluate the Predictions.
print('The STD of the target is: {}'.format(round(edu_df[target].std(), 2)))
print('The mse of the model is: {}'.format(round(mse, 2)))
print('The R-squared is: ', round(r2_score(y_test, y_pred), 2))

# Importances of the different features in our model.
importances = rf.feature_importances_
forest_importances = pd.Series(importances, index=features)

top_20 = forest_importances.sort_values(ascending=False)[:20]
fig, ax = plt.subplots()
top_20.plot.bar()
ax.set_title("Feature Importances")
ax.set_ylabel("Importance")
ax.set_xlabel("Features")
plt.show()

# Evaluating the scores.
cv_results = pd.DataFrame(rf_grid.cv_results_)
cv_results = cv_results.sort_values('mean_test_score', ascending=False)
cv_results[['mean_train_score', 'std_train_score', 'mean_test_score', 'std_test_score']].head()