import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from io import StringIO
from sklearn.base import BaseEstimator, OneToOneFeatureMixin, TransformerMixin
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsRegressor
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.utils.validation import check_is_fitted, validate_data

# One test district, with its total_rooms multiplied by 1x ... 1000x (a unit error in new data).
# The same pipeline with and without the lesson's Clipper (percentiles 1 and 99 learned on train).
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']


class Clipper(OneToOneFeatureMixin, TransformerMixin, BaseEstimator):
    def __init__(self, low=1, high=99):
        self.low = low
        self.high = high

    def fit(self, X, y=None):
        X = validate_data(self, X)
        self.lower_ = np.percentile(X, self.low, axis=0)
        self.upper_ = np.percentile(X, self.high, axis=0)
        return self

    def transform(self, X):
        check_is_fitted(self)
        X = validate_data(self, X, reset=False)
        return np.clip(X, self.lower_, self.upper_)


def build(clip):
    steps = [SimpleImputer(strategy='median')] + ([Clipper()] if clip else []) + [StandardScaler()]
    return make_pipeline(
        ColumnTransformer([
            ('num', make_pipeline(*steps), num),
            ('cat', OneHotEncoder(handle_unknown='ignore', sparse_output=False), ['ocean_proximity'])]),
        KNeighborsRegressor(n_neighbors=10)).fit(X_train, y_train)


plain, clipped = build(False), build(True)
upper = clipped[0].named_transformers_['num'][1].upper_[num.index('total_rooms')]
row = X_test.head(1)
rooms0 = row['total_rooms'].iloc[0]
factors = np.logspace(0, 3, 61)
batch = pd.concat([row] * len(factors), ignore_index=True)
batch['total_rooms'] = rooms0 * factors
p_plain = plain.predict(batch) / 1000
p_clip = clipped.predict(batch) / 1000

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, ax = plt.subplots(figsize=(8, 3.8))
ax.plot(rooms0 * factors, p_plain, color=ORANGE, lw=2, label='pipeline without the Clipper')
ax.plot(rooms0 * factors, p_clip, color=BLUE, lw=2, label='pipeline with the Clipper')
ax.axhline(y_test.iloc[0] / 1000, color=GREEN, ls='--', lw=1.2)
ax.text(rooms0 * 1.1, y_test.iloc[0] / 1000 - 6, f'true value \\${y_test.iloc[0] / 1000:.1f}k',
        color=GREEN, fontsize=8, va='top')
ax.axvline(upper, color='gray', ls=':', lw=1.2)
ax.text(upper * 1.08, 6, f'upper_ = {upper:,.0f}\n(99th percentile of\nthe training rows)', fontsize=8, va='bottom')
ax.axvline(rooms0, color='gray', lw=0.8)
ax.text(rooms0 * 1.08, 1.05 * p_plain.max(), f'real value: {rooms0:,.0f} rooms', fontsize=8)
ax.set_xscale('log')
ax.set_ylim(0, 1.12 * p_plain.max())
ax.set_xlabel('total_rooms fed to the model for this district (log scale)', fontsize=9)
ax.set_ylabel('predicted value (\\$ thousands)', fontsize=9)
ax.tick_params(labelsize=8)
ax.legend(fontsize=8, frameon=False, loc='upper left', bbox_to_anchor=(0.0, 0.9))
for s in ('top', 'right'):
    ax.spines[s].set_visible(False)
ax.set_title(f'A typo in one input: prediction \\${p_plain[-1]:.0f}k without the Clipper, '
             f'\\${p_clip[-1]:.0f}k with it', fontsize=10)

plt.tight_layout()

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

plt.close()