import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from matplotlib.colors import LinearSegmentedColormap
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler

# The lesson's ColumnTransformer on California housing: 8 raw test districts in (left),
# the 13 columns the model receives out (right). Fitted on the training districts only.
url = "https://raw.githubusercontent.com/hsandmann/biblio/refs/heads/main/ml/aula02/housing.csv"
housing = pd.read_csv(url)
X = housing.drop(columns='median_house_value')
y = housing['median_house_value']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
num = ['longitude', 'latitude', 'housing_median_age', 'total_rooms',
       'total_bedrooms', 'population', 'households', 'median_income']
preprocess = ColumnTransformer([
    ('num', make_pipeline(SimpleImputer(strategy='median'), StandardScaler()), num),
    ('cat', OneHotEncoder(handle_unknown='ignore', sparse_output=False), ['ocean_proximity']),
]).fit(X_train)

missing = X_test.index[X_test['total_bedrooms'].isna()][0]        # one district with a hole
rows = list(X_test.index[:7]) + [missing]
raw = X_test.loc[rows]
out = preprocess.transform(raw)
names = [n.split('__')[1].replace('ocean_proximity_', '') for n in preprocess.get_feature_names_out()]

BLUE, ORANGE, GREEN = '#58a6ff', '#f0883e', '#3fb950'
plt.rcParams.update({'text.color': 'gray', 'axes.labelcolor': 'gray', 'axes.edgecolor': 'gray',
                     'xtick.color': 'gray', 'ytick.color': 'gray'})
fig, (axl, axr) = plt.subplots(1, 2, figsize=(10, 4.3), gridspec_kw={'width_ratios': [1, 1.9]})

# left: a few raw columns, as text
show = ['median_income', 'total_bedrooms', 'ocean_proximity']
axl.set_xlim(0, 3); axl.set_ylim(len(rows) - 0.5, -1.4)
axl.axis('off')
for j, c in enumerate(show):
    axl.text(j + 0.5, -0.9, c, ha='center', fontsize=7.5, fontweight='bold')
    for i, r in enumerate(rows):
        v = raw.loc[r, c]
        txt = 'NaN' if (isinstance(v, float) and np.isnan(v)) else (f'{v:g}' if not isinstance(v, str) else v)
        axl.text(j + 0.5, i, txt, ha='center', va='center', fontsize=8,
                 color='#ff7b72' if txt == 'NaN' else 'gray')
axl.set_title('raw test rows (3 of the 9 columns)', fontsize=9)

# right: the output matrix
cmap = LinearSegmentedColormap.from_list('bo', [BLUE, '#d0d7de', ORANGE])
numblock = np.ma.masked_array(out, mask=np.zeros_like(out, bool))
numblock.mask[:, 8:] = True
axr.imshow(np.clip(numblock, -2.5, 2.5), cmap=cmap, vmin=-2.5, vmax=2.5, aspect='auto')
catblock = np.ma.masked_array(out, mask=np.zeros_like(out, bool))
catblock.mask[:, :8] = True
axr.imshow(catblock, cmap=LinearSegmentedColormap.from_list('g', ['#ffffff00', GREEN]),
           vmin=0, vmax=1, aspect='auto')
for i in range(out.shape[0]):
    for j in range(out.shape[1]):
        axr.text(j, i, f'{out[i, j]:.1f}' if j < 8 else f'{out[i, j]:.0f}', ha='center', va='center',
                 fontsize=7.5, color='#24292f' if j < 8 else ('white' if out[i, j] else 'gray'))
axr.set_xticks(range(len(names)))
axr.set_xticklabels(names, rotation=40, ha='right', fontsize=7.5)
axr.set_yticks([])
axr.axvline(7.5, color='gray', lw=1.5)
axr.set_title('what the model receives: 13 columns', fontsize=9)
axr.text(3.5, -1.0, "'num': median, then (x − mean) / std  (8 → 8)", ha='center', fontsize=8, color=BLUE)
axr.text(10, -1.0, "'cat': one-hot  (1 → 5)", ha='center', fontsize=8, color=GREEN)
axr.set_ylim(len(rows) - 0.5, -1.4)
i_miss = len(rows) - 1
axr.add_patch(plt.Rectangle((3.5, i_miss - 0.5), 1, 1, fill=False, ec='#ff7b72', lw=2))
med = preprocess.named_transformers_['num'][0].statistics_[num.index('total_bedrooms')]
axl.text(0.5, -0.06, f'NaN → training median {med:g} → scaled {out[i_miss, 4]:.1f}',
         ha='center', va='top', fontsize=7.5, color='#ff7b72', transform=axl.transAxes)
for s in axr.spines.values():
    s.set_visible(False)

plt.tight_layout()

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

plt.close()