import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.experimental import enable_iterative_imputer  # noqa: F401
from sklearn.impute import IterativeImputer, KNNImputer, SimpleImputer

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

url = "https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula04/titanic/train.csv"
X = pd.read_csv(url)[['Age', 'Pclass', 'SibSp', 'Parch']]
miss = X['Age'].isna().to_numpy()
versions = [('only the 714 observed ages', X['Age'].dropna().to_numpy(), None),
            ('mean imputation', SimpleImputer(strategy='mean').fit_transform(X)[:, 0], '#ff7b72'),
            ('KNNImputer (5 neighbors)', KNNImputer(n_neighbors=5).fit_transform(X)[:, 0], '#f0883e'),
            ('IterativeImputer', IterativeImputer(random_state=0).fit_transform(X)[:, 0], '#58a6ff')]
bins = np.arange(-10, 82, 2.5)

fig, axes = plt.subplots(1, 4, figsize=(10, 2.9), sharey=True)
for ax, (name, v, col) in zip(axes, versions):
    if col is None:
        ax.hist(v, bins=bins, color='gray', alpha=0.7)
    else:
        ax.hist([v[~miss], v[miss]], bins=bins, stacked=True, color=['gray', col], alpha=0.85,
                label=['observed', '177 imputed'])
        ax.legend(fontsize=7.5, frameon=False, labelcolor='gray',
                  loc='upper right' if name == 'IterativeImputer' else 'upper left')
    ax.set_title(f'{name}\nstd {np.std(v, ddof=1):.2f} years', fontsize=9)
    ax.tick_params(labelsize=7.5)
    ax.set_xlabel('Age', fontsize=8.5)
    for side in ['left', 'right', 'top']:
        ax.spines[side].set_visible(False)
axes[0].set_yticks([])
ages = versions[3][1]
axes[3].annotate(f'{(ages < 0).sum()} negative ages\n(min {ages.min():.1f})', xy=(ages.min(), 3),
                 xytext=(-8, 95), fontsize=8, color='#58a6ff', arrowprops=dict(arrowstyle='->', color='#58a6ff'))
axes[1].annotate('177 copies\nof 29.7', xy=(29.7, 170), xytext=(45, 150), fontsize=8, color='#ff7b72',
                 arrowprops=dict(arrowstyle='->', color='#ff7b72'))
fig.suptitle('Titanic ages: every imputation shrinks the spread; the mean does it with a single spike',
             fontsize=10)
plt.tight_layout()

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

plt.close()