# Install packages
! uv add lightgbm scikit-optimize category_encoders

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from category_encoders import LeaveOneOutEncoder
from lightgbm import LGBMRegressor
import seaborn as sns
from sklearn.compose import ColumnTransformer, TransformedTargetRegressor
from sklearn.metrics import make_scorer, mean_pinball_loss, median_absolute_error
from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, PowerTransformer, RobustScaler
from sklearn import set_config
from skopt import gp_minimize, space
from skopt.utils import use_named_args

# Settings
%config InlineBackend.figure_format = 'retina'
sns.set_theme(font_scale=1.5, rc={'lines.linewidth': 2.5})

def quantile_loss(y_true, y_pred, alpha):
    """
    Quantile loss function.

    Parameters
    ----------
    y_true : ndarray
        True observations (i.e., labels).
    y_pred : ndarray
        Model predictions (same shape as `y_true`).
    alpha : float
        Evaluated quantile values. Should be between 0 and 1.
    """
    errors = y_true - y_pred
    pos = errors >= 0.0
    neg = np.logical_not(pos)
    loss_vals = alpha * (pos * errors) + (alpha - 1) * (neg * errors)
    return loss_vals

y_true = 5.0
y_pred = np.linspace(0.0, 10, 51)
errors = y_true - y_pred

# Quantile values
alphas = [0.1, 0.5, 0.9]
colors = ['sandybrown', 'chocolate', 'saddlebrown']

fig, axes = plt.subplots(1, 3, sharey=True, figsize=(15, 5))
fig.tight_layout(w_pad=1)
axline_kw = dict(ls='--', lw=1.5, c='gray')

for alpha, color, ax in zip(alphas, colors, axes):
    ax.axvline(0.0, **axline_kw)
    ax.axhline(0.0, **axline_kw)

    loss_vals = quantile_loss(y_true, y_pred, alpha)
    ax.plot(errors, loss_vals, color)
    ax.set(title=fr'$\alpha={alpha}$', xlabel=r'Error, $y - \hat{y}$')

axes[0].set_ylabel(r'Quantile loss, $\mathcal{L}$');

cars = pd.read_csv('../data/autoscout24_germany_dataset_cleaned.csv')
cars.sample(frac=1).head(n=10)

cars.drop(columns=['model'], inplace=True)
cars.head()

data = cars.drop(columns='price', inplace=False)
target = cars['price']

X_train, X_test, y_train, y_test = \
    train_test_split(data, target, test_size=0.2)

for feature in ('make', 'fuel', 'gear'):
    print(
        f'Feature: {feature} - '
        f'Number of categories: {data[feature].nunique()}')

loo_encoder = LeaveOneOutEncoder(sigma=0.05)
encoded = loo_encoder.fit_transform(X_train['make'], y_train)
encoded.head(n=10)

pt = PowerTransformer(method='box-cox')
y_orig = target.values[:, np.newaxis]
y_norm = pt.fit_transform(y_orig)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
plt.suptitle('Targets Distribution', y=1)

hist_kw = dict(bins=30, facecolor='teal')

_ = ax1.hist(y_orig, **hist_kw)
ax1.set_title('Original')

_ = ax2.hist(y_norm, **hist_kw)
ax2.set_title('Transformed (Box-Cox)');

def create_preprocessor():
    """
    Data preprocessing and feature transformation pipeline

    Returns
    -------
    out : `sklearn.compose.ColumnTransformer` object
    """
    # Numerical features
    numerical_features = ['mileage', 'hp', 'age']
    numerical_transformer = RobustScaler()

    # Categorical features
    categorical_features_a = ['fuel', 'gear']
    categorical_transformer_a = OneHotEncoder(handle_unknown='ignore')

    categorical_features_b = ['make']
    categorical_transformer_b = Pipeline(
        steps=[
            ('loo_encoder', LeaveOneOutEncoder()),
            ('scaler', RobustScaler())])

    # Preprocessing pipeline
    preprocessor = ColumnTransformer(
        transformers=[
            ('numer', numerical_transformer, numerical_features),
            ('categ_a', categorical_transformer_a, categorical_features_a),
            ('categ_b', categorical_transformer_b, categorical_features_b)])

    return preprocessor

def create_quantile_regressor(alpha=0.5):
    """
    LightGBM regressor w/ quantile loss applied on transformed targets.
    In the regression problem, targets (prices) are transformed by using
    Box-Cox power transform method, `PowerTransform(method='box-cox')`.

    Parameters
    ----------
    alpha: float
        The coefficient used in quantile-based loss, should be between 0
        and 1. Default: 0.5 (median).

    Returns
    -------
    Quantile regressor to apply on transformed targets.
    """
    ttq_regressor = TransformedTargetRegressor(
        regressor=LGBMRegressor(objective='quantile', alpha=alpha),
        transformer=PowerTransformer(method='box-cox'))

    return ttq_regressor

def create_full_pipeline(alpha=0.5):
    """
    Full prediction pipeline.

    Parameters
    ----------
    alpha: float
        The coefficient used in quantile-based loss, should be between 0
        and 1. Default: 0.5 (median).

    Returns
    -------
    Pipeline of data transforms with the final estimator returned as
    `sklearn.pipeline.Pipeline` object.
    """
    preprocessor = create_preprocessor()
    ttq_regressor = create_quantile_regressor(alpha=alpha)
    model = Pipeline(
        steps=[
            ('preprocessor', preprocessor),
            ('ttq_regressor', ttq_regressor)])

    return model

set_config(display='diagram')
create_full_pipeline()

def hyperparameter_tuner(model, X, y, search_space, cvs_params, gp_params):
    """
    Parameters
    ----------
    model : estimator object implementing `fit()` method
        The object to use to fit the data.
    X : array-like of shape (n_samples, n_features)
        The data to fit.
    y : array-like of shape (n_samples,) or (n_samples, n_outputs)
        The target variable to try to predict.
    search_space : list, shape (n_dims,)
        List of hyperparameter search space dimensions.
    cvs_params : dict
        Parameters for Scikit-Learn's `cross_val_score()` method.
    gp_params : dict
        Parameters for Scikit-Optimize's `gp_minimize()` method.

    Returns
    -------
    The Gaussian Processes optimization result with the required
    information, returned as SciPy's `OptimizeResult` object.
    """
    @use_named_args(search_space)
    def objective_func(**params):
        """
        Function to *minimize*. Decorated by the `use_named_args()`, it
        takes named (keyword) arguments and returns the objective value.
        """
        model.set_params(**params)
        scores = cross_val_score(model, X, y, **cvs_params)
        return -1.0 * np.mean(scores)

    # We are now ready for sequential model-based optimization
    optim_result = gp_minimize(objective_func, search_space, **gp_params)
    return optim_result

%%time

# Parameters of a Scikit-Learn estimator in a pipeline are accessed
# using the `<estimator>__<parameter>` syntax. Valid parameter keys
# can be listed with model's `get_params()` method.

# Hyperparameters search space
search_space = [
    space.Real(
        0.01, 0.6, prior='log-uniform',
        name='preprocessor__categ_b__loo_encoder__sigma'),
    space.Integer(10, 50, name='ttq_regressor__regressor__num_leaves'),
    space.Integer(2, 30, name='ttq_regressor__regressor__max_depth'),
    space.Real(
        0.01, 0.5, prior='log-uniform',
        name='ttq_regressor__regressor__learning_rate')]

# Parameters that'll be passed to sklearn's `cross_val_score()`
cvs_params = {'cv': 3, 'n_jobs': 4}

# Parameters that'll be passed to skopt's `gp_minimize()`
gp_params = {'n_calls': 50,
             'n_initial_points': 10,
             'acq_func': 'EI',
             'n_jobs': 4,
             'verbose': False}

# Lower, middle and higher quantiles to evaluate: lower/upper bounds + median
quantiles = {'q_lower': 0.025, 'q_middle': 0.5, 'q_upper': 0.975}

# Cache Gaussian Process optimization results
cache_gp_results = dict()

for q_level, alpha in quantiles.items():
    print(f'\nHyperparameter optimization for quantile regressor with '
          f'alpha={alpha:.3f} ..... ', end='')

    # Create model pipeline
    model = create_full_pipeline(alpha=alpha)

    # Cross-validation scoring strategy
    scorer = make_scorer(
        mean_pinball_loss, alpha=alpha, greater_is_better=False)
    cvs_params['scoring'] = scorer

    # Find optimal hyperparameters
    gp_result = hyperparameter_tuner(
        model, X_train, y_train, search_space, cvs_params, gp_params)

    # Cache current results
    cache_gp_results[q_level] = gp_result
    print('Done')

    # Report GP optimization result
    print('Best parameters:')
    for i_param, param_value in enumerate(gp_result.x):
        param_name = search_space[i_param].name.split('__')[-1]
        print(''.rjust(4) + f'{param_name}: {param_value:.6f}')
    print('=' * 40)

n_models = len(quantiles)
fig, axes = plt.subplots(1, n_models, figsize=(15, 5))
fig.tight_layout(w_pad=1)

colors = ['sandybrown', 'chocolate', 'saddlebrown']

for i, (q_level, gp_result) in enumerate(cache_gp_results.items()):
    call_idx = np.arange(1, gp_result.func_vals.size + 1)
    min_vals = [gp_result.func_vals[:i].min() for i in call_idx]

    ax = axes[i]
    ax.plot(call_idx, min_vals, '-o', color=colors[i])

    alpha = quantiles[q_level]
    ax.set(xlabel=r'Number of calls, $n$', title=fr'$\alpha$={alpha}')

axes[0].set_ylabel(r'$f_{min}(x)$ after $n$ calls')
plt.suptitle('GP Convergence Traces', y=1.1);

trained_models = dict()

for q_level, alpha in quantiles.items():
    # Create model for this quantile
    model = create_full_pipeline(alpha=alpha)

    # Retrieve best (optimized) parameters for this model
    gp_res = cache_gp_results[q_level]
    n_params = len(search_space)
    best_params = dict(
        (search_space[i].name, gp_res.x[i]) for i in range(n_params))

    # Set the parameters of this model
    model.set_params(**best_params)

    model.fit(X_train, y_train)
    trained_models[q_level] = model

def coverage_percentage(
        y: np.ndarray,
        y_lower: np.ndarray,
        y_upper: np.ndarray) -> float:
    """
    Function to calculate the coverage percentage of the predictive
    interval, i.e. the percentage of observations that fall between the
    lower and upper bounds of the predictions.

    Parameters
    ----------
    y : ndarray of shape (n_samples,)
        Target values.
    y_lower : ndarray of shape (n_samples,)
        Lower predictive intervals.
    y_upper : ndarray of shape (n_samples,)
        Upper predictive intervals.
    """
    return np.mean(np.logical_and(y > y_lower, y < y_upper)) * 100.0

# Compute predictive intervals

y_pred_lower = trained_models['q_lower'].predict(X_test)
y_pred_upper = trained_models['q_upper'].predict(X_test)

ctop = coverage_percentage(y_test, y_pred_lower, np.inf)
cbot = coverage_percentage(y_test, -np.inf, y_pred_upper)
cint = coverage_percentage(y_test, y_pred_lower, y_pred_upper)

print(f'Coverage of top 95%: {ctop:.1f}%',
      f'Coverage of bottom 95%: {cbot:.1f}%',
      f'Coverage of 95% predictive interval: {cint:.1f}%', sep='\n')

y_pred_middle = trained_models['q_middle'].predict(X_test)

print(
    f'Median Absolute Error: ',
    f'\tLower-quantile model: {int(median_absolute_error(y_test, y_pred_lower))}',
    f'\tMiddle-quantile model: {int(median_absolute_error(y_test, y_pred_middle))}',
    f'\tUpper-quantile model: {int(median_absolute_error(y_test, y_pred_upper))}',
    sep='\n')