import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import KFold, learning_curve, train_test_split
from sklearn.neighbors import KNeighborsRegressor
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import FunctionTransformer, OneHotEncoder, StandardScaler
from sklearn.tree import DecisionTreeRegressor

plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})

# The lab's step 4: learning curves of the three default candidates on California housing.
url = 'https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula02/housing.csv'
df = pd.read_csv(url)
X, y = df.drop(columns='median_house_value'), df['median_house_value']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
counts = ['total_rooms', 'total_bedrooms', 'population', 'households']
others = ['longitude', 'latitude', 'housing_median_age', 'median_income']
pre = ColumnTransformer([
    ('counts', make_pipeline(SimpleImputer(strategy='median'), FunctionTransformer(np.log1p),
                             StandardScaler()), counts),
    ('others', StandardScaler(), others),
    ('ocean', OneHotEncoder(handle_unknown='ignore'), ['ocean_proximity']),
])
cv = KFold(n_splits=5, shuffle=True, random_state=42)
candidates = {
    'linear: high bias': make_pipeline(pre, LinearRegression()),
    'k-NN (k=5): variance': make_pipeline(pre, KNeighborsRegressor()),
    'tree (no limits): high variance': make_pipeline(pre, DecisionTreeRegressor(random_state=42)),
}
fractions = [0.02, 0.05, 0.1, 0.2, 0.4, 0.7, 1.0]

fig, axes = plt.subplots(1, 3, figsize=(10.5, 3.5), sharey=True)
for ax, (name, model) in zip(axes, candidates.items()):
    sizes, train, val = learning_curve(model, X_train, y_train, train_sizes=fractions, cv=cv,
                                       scoring='neg_root_mean_squared_error')
    train, val = np.abs(train.mean(axis=1)) / 1000, np.abs(val.mean(axis=1)) / 1000
    ax.plot(sizes, train, 'o-', color='#f0883e', ms=4, lw=2, label='training (CV folds)')
    ax.plot(sizes, val, 'o-', color='#58a6ff', ms=4, lw=2, label='validation (CV folds)')
    ax.annotate(f'{val[-1]:.1f}', xy=(sizes[-1], val[-1]), xytext=(0, 7), textcoords='offset points', ha='right',
                fontsize=8, color='#58a6ff')
    ax.set_xscale('log')
    ax.set_xticks([300, 1000, 3000, 10000], ['300', '1k', '3k', '10k'])
    ax.set_xlabel('training rows (log scale)')
    ax.set_title(name, fontsize=10)
    ax.set_ylim(-3, 100)
axes[0].set_ylabel('RMSE (thousand US$)')
axes[0].legend(fontsize=8, loc='lower left', frameon=False)
axes[0].text(1200, 78, 'converged by ~2,600 rows:\nthe gap is gone,\nthe error is not', fontsize=7.5)
axes[1].text(700, 20, 'gap of 11k at 13,209 rows,\nvalidation still falling', fontsize=7.5)
axes[2].text(600, 30, 'training error 0 at every size;\nvalidation falls by 24k and\nhas not leveled off', fontsize=7.5)
plt.tight_layout()

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()