from itertools import chain

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import missingno as msno
import pycountry_convert as pcc

from matplotlib.ticker import MaxNLocator
from matplotlib.ticker import FormatStrFormatter
from sklearn.model_selection import train_test_split
from sklearn.model_selection import ParameterGrid
from sklearn.metrics import mean_squared_error as mse
from sklearn.metrics import mean_absolute_error as mae
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import PolynomialFeatures
from sklearn.preprocessing import MinMaxScaler
from sklearn.preprocessing import MaxAbsScaler
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import FunctionTransformer
from sklearn.tree import DecisionTreeRegressor
from sklearn.base import RegressorMixin
from sklearn.base import BaseEstimator
from sklearn.base import TransformerMixin
from sklearn.linear_model import Ridge
from sklearn.neighbors import KNeighborsRegressor
from sklearn.utils import resample
from sklearn.impute import SimpleImputer
from sklearn.impute import KNNImputer
from sklearn.compose import ColumnTransformer

np.random.seed(42)

plt.style.use('ggplot')
red =  np.array((226/255, 74/255, 51/255))
blue = np.array((52/255, 138/255, 189/255))
grey = np.array((100/255, 100/255, 100/255))
cobalt = np.array((0/255, 71/255, 171/255))
main_color = cobalt

raw_data = pd.read_csv('data.csv')
raw_data.head().T

pd.DataFrame([raw_data.count(), raw_data.nunique(), raw_data.dtypes], index=['Non-Null Count', 'Unique Values', 'Dtype']).T

raw_data.describe().T

y_data = raw_data.loc[:, 'Life expectancy']
X_data = raw_data.drop('Life expectancy', axis=1)
test_ratio = 0.3
val_ratio = 0.2
X_train_val, X_test, y_train_val, y_test = train_test_split(X_data, y_data, test_size=test_ratio)
X_train, X_val, y_train, y_val = train_test_split(X_train_val, y_train_val, test_size=val_ratio/(1-test_ratio))
split_df = pd.DataFrame([X_train.shape[0], X_val.shape[0], X_test.shape[0]], index=['Train', 'Val', 'Test'], columns=['Size'])
split_df['Relative size'] = split_df['Size'] / split_df['Size'].sum()
split_df['Relative size'] = split_df['Relative size'].round(3)
split_df

fig, ax = plt.subplots(1, 1, figsize=(12, 5), layout='constrained')
fig.suptitle('Missing values (heatmap)', fontsize=20)

msno.matrix(X_train, fontsize=14, sparkline=False, ax=ax, color=main_color)
ax.get_yaxis().set_visible(False)

fig, ax = plt.subplots(1, 1, figsize=(12, 8), layout='constrained')
fig.suptitle('Missing values (counts)', fontsize=20)

missings = X_train.isna().sum().sort_values()
bars = ax.barh(missings.index, missings, color=main_color)
ax.set_xlabel('Count', fontsize=14)
ax.tick_params(axis='both', which='major', labelsize=12)
for bar, count in zip(bars, missings):
    ax.text(bar.get_width()+2, bar.get_y() + bar.get_height() / 2, f'{count}', va='center', fontsize=10)

cat_columns = ['Status', 'Country']
num_columns = ['Year', 'Adult Mortality', 'infant deaths', 'Alcohol', 'percentage expenditure', 'Hepatitis B',
               'Measles', 'BMI', 'under-five deaths', 'Polio', 'Total expenditure', 'Diphtheria', 'HIV/AIDS', 'GDP', 'Population',
               'thinness  1-19 years', 'thinness 5-9 years', 'Income composition of resources', 'Schooling']
logscaled_columns = ['infant deaths', 'percentage expenditure', 'Measles', 'under-five deaths', 'HIV/AIDS', 'GDP', 'Population']

fig, axes = plt.subplots(7, 3, figsize=(12, 16), layout='constrained')
fig.suptitle('Features', size=20)
fig.supylabel('Count', size=18)

for col, ax in zip(num_columns, fig.axes):
    data = X_train[col].dropna()
    if col in logscaled_columns:
        hist, bins = np.histogram(data, bins=20)
        logbins = np.logspace(np.log10(max(bins[0], 0.001)), np.log10(bins[-1]), len(bins))
        ax.set_xscale('log')
        ax.hist(data, bins=logbins, color=main_color)
    else:
        ax.hist(data, color=main_color)
        if col == 'Year':
            ax.xaxis.set_major_formatter(FormatStrFormatter('%d'))
    ax.set_xlabel(col)

ax = axes[6][1]
states = X_train['Status'].value_counts()
ax.bar(states.index, states, color=main_color)
ax.set_xlabel('Status')

axes[6][2].axis('off')

plt.show()

print('List of countries:')
print(', '.join(X_train['Country'].unique()))

fig, axes = plt.subplots(2, 1, figsize=(12, 6), sharex=True)
fig.subplots_adjust(hspace=0.03)
fig.suptitle('Target variable', size=18)

ax = axes[0]
ax.hist(y_train, bins=35, color=main_color)
ax.set_ylabel('Count')
ax.tick_params(labelbottom=False)
ax.tick_params(axis='x', which='both', length=0)

medianprops = dict(linewidth=2.5, color=main_color)
flierprops = dict(marker='o', markerfacecolor='none', markersize=7, markeredgecolor=main_color)
ax = axes[1]
ax.boxplot(y_train, vert=False, widths=[0.4], showfliers=True, flierprops=flierprops, medianprops=medianprops)
ax.get_yaxis().set_visible(False)
ax.set_xlabel('Life expectancy [years]')

plt.show()

fig, axes = plt.subplots(7, 3, figsize=(12, 28), layout='constrained')
fig.suptitle('Life expectancy vs Features', size=20)
# fig.supxlabel('Life expectancy', size=18)

alpha = 0.3
for col, ax in zip(num_columns, fig.axes):
    data = X_train[col]
    if col in logscaled_columns:
        ax.set_yscale('log')
        ax.scatter(y_train, data, alpha=alpha, color=main_color)
    else:
        ax.scatter(y_train, data, alpha=alpha, color=main_color)
    ax.set_ylabel(col)
    ax.set_title(f'Life expectancy vs {col}', fontsize=12)

axes[6][1].axis('off')
axes[6][2].axis('off')

plt.show()

colors_r = plt.get_cmap('Blues')(np.linspace(0, 1, 128))
colors_l = colors_r[::-1]
ggcmap_bi = mcolors.LinearSegmentedColormap.from_list('ggplot_like', np.vstack((colors_l, colors_r)))

fig, ax = plt.subplots(figsize=(12, 8), layout='constrained')
fig.suptitle('Correlation matrix', size=20)
corr = X_train[num_columns].corr()
sns.heatmap(corr.round(1), ax=ax, square=True, annot=True, linewidths=0.3,
            annot_kws={"size": 8}, cmap=ggcmap_bi, vmin=-1, vmax=1)
ax.tick_params(axis='both', which='both',length=0)
plt.show()

high_corr_columns = ['under-five deaths', 'thinness 5-9 years', 'GDP']

X_train[cat_columns].describe().T

cat_encoder = Pipeline([
    ('ohe_encoder', OneHotEncoder(handle_unknown='ignore')),
])

X_train[num_columns].describe().T

num_encoder = Pipeline([
    # ('knn_imputer', KNNImputer(missing_values=np.nan, n_neighbors=5, weights='distance')),
    ('mean_imputer', SimpleImputer(missing_values=np.nan, strategy='mean')),
])

def log_transform(x):
    return np.log(x + 1)

log_encoder = Pipeline([
    # ('knn_imputer', KNNImputer(missing_values=np.nan, n_neighbors=5, weights='distance')),
    ('mean_imputer', SimpleImputer(missing_values=np.nan, strategy='mean')),
    ('log_transformer', FunctionTransformer(log_transform))
])

def country_to_continent(country):
    try:
        alpha2 = pcc.country_name_to_country_alpha2(country)
        continent = pcc.country_alpha2_to_continent_code(alpha2)
        return continent
    except KeyError:
        return np.nan

class ContinentTransformer(BaseEstimator, TransformerMixin):
    def __init__(self):
        pass
        
    def fit(self, X=None, y=None):
        return self
    
    def transform(self, X, y=None):
        X['Country'] = X['Country'].apply(country_to_continent)
        return X

continent_encoder = Pipeline([
    ('continent_transformer', ContinentTransformer()),
    ('ohe_encoder', OneHotEncoder(handle_unknown='ignore')),
])

column_transformer = ColumnTransformer([
    # ('continent_encoder', continent_encoder, ['Country']),
    ('categorical_encoder', cat_encoder, cat_columns),
    ('numerical_encoder', num_encoder, num_columns),
])

preprocessor_random_forest = Pipeline([
    ('preprocessor', column_transformer),
])

preprocessor_random_forest.fit(X_train)

column_transformer = ColumnTransformer([
    ('continent_encoder', continent_encoder, ['Country']),
    ('categorical_encoder', cat_encoder, cat_columns),
    ('numerical_encoder', num_encoder, [col for col in num_columns if col not in [high_corr_columns]]),
    ('logscaled_encoder', log_encoder, logscaled_columns),
])

preprocessor_ridge = Pipeline([
    ('preprocessor', column_transformer),
    # ('basis_functions', PolynomialFeatures(include_bias=True)),
    ('scaler', MaxAbsScaler()),
])

preprocessor_ridge.fit(X_train)

column_transformer = ColumnTransformer([
    ('continent_encoder', continent_encoder, ['Country']),
    ('categorical_encoder', cat_encoder, ['Status']),
    ('numerical_encoder', num_encoder, num_columns),
])

preprocessor_knn = Pipeline([
    ('preprocessor', column_transformer),
    ('scaler', MaxAbsScaler()),
])

preprocessor_knn.fit(X_train)

class CustomRandomForest(RegressorMixin):
    """
    Třída Vašeho modelu
    Bude se jednat o model náhodného lesa, kde podmodely tvoří rozhodovací stromy pro regresi.
    Pro podmodely můžete použít implementaci DecisionTreeRegressor ze sklearn.
    """
    def __init__(self, n_estimators, max_samples_fraction, max_depth, **kwargs):
        """
        Konstruktor modelu
        Základní hyperparametery:
            n_estimators - počet podmodelů - rozhodovacích stromů.
            max_samples_fraction - vyberte si, zda tento parametr bude označovat relativní počet bodů (tj. číslo mezi 0 a 1) 
                          nebo absolutní počet bodů (tj. číslo mezi 1 a velikostí trénovací množiny), 
                          které budou pro každý podmodel rozhodovacího stromu náhodně vybrány z trénovací množiny (bootstrap) a použity k jeho trénování.
            max_depth - maximální hloubka každého z podmodelů rozhodovacího stromu.
            kwargs - (volitelně) případné další hyperparametry, které pošlete do podmodelů rozhodovacího stromu
        """
        self.n_estimators = n_estimators
        self.max_samples_fraction = max_samples_fraction
        self.max_depth = max_depth
        self.decision_tree_kwargs = kwargs
        
    def fit(self, X, y):
        """
        Natrénování modelu. Trénovací data jsou v argumentech X a y.
        Pro trénování podmodelů používejte bootstraping a velikost samplovaného vzorku vezměte z hyperparametru max_samples_fraction
        """
        self.estimators = []
        n_samples = int(X.shape[0] * self.max_samples_fraction)
        for _ in range(self.n_estimators):
            X_sample, y_sample = resample(X, y, replace=True, n_samples=n_samples)
            tree = DecisionTreeRegressor(splitter='random', max_depth=self.max_depth, **self.decision_tree_kwargs)
            tree.fit(X_sample, y_sample)
            self.estimators.append(tree)
            
    def predict(self, X):
        """
        Predikce y v zadaných bodech X
        """
        estimations = np.zeros((self.n_estimators, X.shape[0]))
        for i, estimator in enumerate(self.estimators):
            estimations[i] = estimator.predict(X)
        y_predicted = estimations.mean(axis=0)
        return y_predicted

param_grid = ParameterGrid({
    'n_estimators': range(20, 300, 20),
})

X_train_np = preprocessor_random_forest.transform(X_train)
X_val_np = preprocessor_random_forest.transform(X_val)

log_rf = pd.DataFrame(columns=['n_estimators', 'train_rmse', 'val_rmse'])
estimators_rf = []
for params in param_grid:
    reg = CustomRandomForest(max_depth=None, max_samples_fraction=1, **params)
    reg.fit(X_train_np, y_train)
    train_rmse = mse(y_train, reg.predict(X_train_np), squared=False)
    val_rmse = mse(y_val, reg.predict(X_val_np), squared=False)
    log_rf.loc[len(log_rf.index)] = [params['n_estimators'], train_rmse, val_rmse]
    estimators_rf.append(reg)

log_rf.sort_index(inplace=True)

fig, axes = plt.subplots(2, 1, figsize=(12, 5), layout='constrained', sharex=True)
fig.suptitle('Random forest learning curve', size=20)
fig.supylabel('RMSE', size=16)
fig.supxlabel('Number of estimators', size=16)

ax = axes[0]
ax.plot(log_rf['n_estimators'], log_rf['train_rmse'], label='train', color=main_color, linewidth=2.5)
ax.legend()

ax = axes[1]
ax.plot(log_rf['n_estimators'], log_rf['val_rmse'], label='validation', color='green', linewidth=2.5)
ax.legend()

plt.show()

log_rf = log_rf.sort_values('val_rmse')
log_rf.head(5)

best_rf = estimators_rf[log_rf.sort_values('val_rmse').index[0]]

X_val_np = preprocessor_random_forest.transform(X_val)
rf_eval = pd.DataFrame([
   mse(y_val, best_rf.predict(X_val_np), squared=False),
   mae(y_val, best_rf.predict(X_val_np)),
], index=['RMSE', 'MAE'], columns=['random forest'])
rf_eval

param_grid = ParameterGrid({
    'alpha': np.linspace(0, 0.01, num=100),
})

X_train_np = preprocessor_ridge.transform(X_train)
X_val_np = preprocessor_ridge.transform(X_val)

log_ridge = pd.DataFrame(columns=['alpha', 'train_rmse', 'val_rmse'])
estimators_ridge = []
for params in param_grid:
    reg = Ridge(alpha=params['alpha'])
    reg.fit(X_train_np, y_train)
    train_rmse = mse(y_train, reg.predict(X_train_np), squared=False)
    val_rmse = mse(y_val, reg.predict(X_val_np), squared=False)
    log_ridge.loc[len(log_ridge.index)] = [params['alpha'], train_rmse, val_rmse]
    estimators_ridge.append(reg)

log_ridge.sort_index(inplace=True)

fig, axes = plt.subplots(2, 1, figsize=(12, 5), layout='constrained', sharex=True)
fig.suptitle('Ridge learning curve', size=20)
fig.supylabel('RMSE', size=16)
fig.supxlabel('Alpha', size=16)

ax = axes[0]
ax.plot(log_ridge['alpha'], log_ridge['train_rmse'], label='train', color=main_color, linewidth=2.5)
ax.legend(loc='lower right')

ax = axes[1]
ax.plot(log_ridge['alpha'], log_ridge['val_rmse'], label='validation', color='darkorchid', linewidth=2.5)
ax.legend(loc='lower right')

plt.show()

log_ridge = log_ridge.sort_values('val_rmse')
log_ridge.head(5)

best_ridge = estimators_ridge[log_ridge.sort_values('val_rmse').index[0]]

X_val_np = preprocessor_ridge.transform(X_val)
ridge_eval = pd.DataFrame([
   mse(y_val, best_ridge.predict(X_val_np), squared=False),
   mae(y_val, best_ridge.predict(X_val_np)),
], index=['RMSE', 'MAE'], columns=['ridge'])
ridge_eval

param_grid = ParameterGrid({
    'n_neighbors': range(1, 20),
    'weights': ['uniform', 'distance'],
})

X_train_np = preprocessor_knn.transform(X_train)
X_val_np = preprocessor_knn.transform(X_val)

log_knn = pd.DataFrame(columns=['n_neighbors', 'weights', 'train_rmse', 'val_rmse'])
estimators_knn = []
for params in param_grid:
    reg = KNeighborsRegressor(n_neighbors=params['n_neighbors'], weights=params['weights'])
    reg.fit(X_train_np, y_train)
    train_rmse = mse(y_train, reg.predict(X_train_np), squared=False)
    val_rmse = mse(y_val, reg.predict(X_val_np), squared=False)
    log_knn.loc[len(log_knn.index)] = [params['n_neighbors'], params['weights'], train_rmse, val_rmse]
    estimators_knn.append(reg)

log_knn.sort_index(inplace=True)
df = log_knn[log_knn['weights'] == 'distance']

fig, axes = plt.subplots(2, 1, figsize=(12, 5), layout='constrained', sharex=True)
fig.suptitle('KNN learning curve (distance)', size=20)
fig.supylabel('RMSE', size=16)
fig.supxlabel('Number of neighbors', size=16)

ax = axes[0]
ax.plot(df['n_neighbors'], df['train_rmse'], label='train', color=main_color, linewidth=2.5)
ax.legend(loc='lower right')

ax = axes[1]
ax.plot(df['n_neighbors'], df['val_rmse'], label='validation', color='chocolate', linewidth=2.5)
ax.legend(loc='lower right')

ax.xaxis.set_major_locator(MaxNLocator(integer=True))
plt.show()

log_knn = log_knn.sort_values('val_rmse')
log_knn.head(5)

best_knn = estimators_knn[log_knn.sort_values('val_rmse').index[0]]

X_val_np = preprocessor_knn.transform(X_val)
knn_eval = pd.DataFrame([
   mse(y_val, best_knn.predict(X_val_np), squared=False),
   mae(y_val, best_knn.predict(X_val_np)),
], index=['RMSE', 'MAE'], columns=['knn'])
knn_eval

eval = pd.concat([rf_eval.T, ridge_eval.T, knn_eval.T])
eval

fig, axes = plt.subplots(1, 2, figsize=(12, 5))

ax = axes[0]

for metric, ax in zip(['RMSE', 'MAE'], fig.axes):
    ax.set_title(metric)
    bars = ax.bar(eval.index, eval.loc[:, metric], color=['green', 'darkorchid', 'chocolate'], width=0.5)
    ax.axhline(bars[0].get_height(), color = 'black', linestyle = '--', alpha=0.5, linewidth=1)
    ax.set_yticks([0, *[b.get_height() for b in bars]])

X_train_val_np = preprocessor_random_forest.fit_transform(X_train_val)
best_rf.fit(X_train_val_np, y_train_val)

X_test_np = preprocessor_random_forest.transform(X_test)
print(mse(y_test, best_rf.predict(X_test_np), squared=False))

eval_data = pd.read_csv('evaluation.csv')
eval_data_np = preprocessor_random_forest.transform(eval_data)
eval_data['Life expectancy'] = best_rf.predict(eval_data_np)
eval_data.to_csv('results.csv', columns=['Country', 'Year', 'Life expectancy'], header=True, index=False)