# Data handling
# ==============================================================================
import pandas as pd
import numpy as np

# Plotting
# ==============================================================================
import matplotlib.pyplot as plt
plt.style.use('fivethirtyeight')
plt.rcParams['lines.linewidth'] = 1.5
plt.rcParams['font.size'] = 8

# Modeling
# ==============================================================================
from ngboost import NGBRegressor
from ngboost import scores
from ngboost import distns
from sklearn.tree import DecisionTreeRegressor
from scipy import stats

# Warnings configuration
# ==============================================================================
import warnings
warnings.filterwarnings('once')

# Simulated data
# ==============================================================================
# Simulation slightly modified from the example published in
# XGBoostLSS – An extension of XGBoost to probabilistic forecasting Alexander März

np.random.seed(seed=1234)
n = 3000

# Normal distribution with changing variance
x = np.linspace(start=0, stop= 24, num=n)
y = np.random.normal(
        loc   = 15,
        scale = (
            1 + 1.5 * ((4.8 < x) & (x < 7.2)) + 4 * ((7.2 < x) & (x < 12))
            + 1.5 * ((12 < x) & (x < 14.4)) + 2 * (x > 16.8)
        )
    )

# Compute the 0.1 and 0.9 quantiles for each simulated value of x
# ==============================================================================
quantile_10 = stats.norm.ppf(
                q = 0.1,
                loc   = 15,
                scale = (
                    1 + 1.5 * ((4.8 < x) & (x < 7.2)) + 4 * ((7.2 < x) & (x < 12))
                    + 1.5 * ((12 < x) & (x < 14.4)) + 2 * (x > 16.8)
                )
             )

quantile_90 = stats.norm.ppf(
                q = 0.9,
                loc   = 15,
                scale = (
                    1 + 1.5 * ((4.8 < x) & (x < 7.2)) + 4 * ((7.2 < x) & (x < 12))
                    + 1.5 * ((12 < x) & (x < 14.4)) + 2 * (x > 16.8)
                )
             )

# Plot
# ==============================================================================
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(7, 3.5))
ax.scatter(x, y, alpha = 0.3, c = "black", s = 8)
ax.plot(x, quantile_10, c = "black")
ax.plot(x, quantile_90, c = "black")
ax.fill_between(x, quantile_10, quantile_90, alpha=0.3, color='red',
                label='Theoretical quantile interval 0.1-0.9')
ax.set_xticks(range(0, 25))
ax.set_title('Electricity consumption throughout the day')
ax.set_xlabel('Hour of the day')
ax.set_ylabel('Electricity consumption (MWh)')
plt.legend();

# Model fitting
# ==============================================================================
model = NGBRegressor(
            Dist  = distns.Normal,
            Base  = DecisionTreeRegressor(criterion='friedman_mse', max_depth=3),
            Score = scores.LogScore,
            natural_gradient = True,
            n_estimators     = 500,
            learning_rate    = 0.01,
            minibatch_frac   = 1.0,
            col_sample       = 1.0,
            random_state     = 123,
        )

model.fit(X=x.reshape(-1, 1), Y=y)

# Point estimate prediction
# ==============================================================================
point_predictions = model.predict(X=x.reshape(-1, 1))
point_predictions

# Conditional distribution prediction
# ==============================================================================
dist_predictions = model.pred_dist(X=x.reshape(-1, 1))
dist_predictions

# Each distns.distn object contains the estimated parameters
dist_predictions[0].params

# Estimated parameters for each hour
# ==============================================================================
df_parameters = pd.DataFrame([dist.params for dist in dist_predictions])
df_parameters['hour'] = x
df_parameters

# Compute the theoretical quantiles to define the central interval accumulating
# 90% of probability using the predicted parameters.
# ==============================================================================
grid_X = np.linspace(start=min(x), stop=max(x), num=1000)
dist_predictions = model.pred_dist(X=grid_X.reshape(-1, 1))
pred_quantile_10 = [stats.norm.ppf(q=0.1, **dist.params) for dist in dist_predictions]
pred_quantile_90 = [stats.norm.ppf(q=0.9, **dist.params) for dist in dist_predictions]

# Plot quantile interval 10%-90%
# ==============================================================================
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(7, 3.5))
ax.scatter(x, y, alpha = 0.3, c = "black", s = 8)
ax.plot(x, quantile_10, c = "black")
ax.plot(x, quantile_90, c = "black")
ax.fill_between(x, quantile_10, quantile_90, alpha=0.3, color='red',
                label='Theoretical quantile interval 0.1-0.9')
ax.plot(grid_X, pred_quantile_10, c = "blue", label='Predicted quantile interval')
ax.plot(grid_X, pred_quantile_90, c = "blue")

ax.set_xticks(range(0, 25))
ax.set_title('Electricity consumption throughout the day')
ax.set_xlabel('Hour of the day')
ax.set_ylabel('Electricity consumption (MWh)')
plt.legend();

# Interval predictions for the available observations
# ==============================================================================
dist_predictions = model.pred_dist(X=x.reshape(-1, 1))
pred_quantile_10 = [stats.norm.ppf(q=0.1, **dist.params) for dist in dist_predictions]
pred_quantile_90 = [stats.norm.ppf(q=0.9, **dist.params) for dist in dist_predictions]

# Coverage of the predicted interval
# ==============================================================================
within_interval = np.where((y >= pred_quantile_10) & (y <= pred_quantile_90), True, False)
coverage = within_interval.mean()
print(f"Coverage of the predicted interval: {100 * coverage:.2f}%")

# Predict distribution parameters when x=8
# ==============================================================================
distribution = model.pred_dist(X=np.array([8]).reshape(-1, 1))
distribution.params

# Probability that consumption exceeds 22.5
# ==============================================================================
100 * (1 - stats.norm.cdf(x=22.5, **distribution.params))

# Probability prediction across the full hourly range
# ==============================================================================
grid_X = np.linspace(start=0, stop=24, num=500)
dist_predictions = model.pred_dist(X=grid_X.reshape(-1, 1))
probability_predictions = [(1 - stats.norm.cdf(x=22.5, **dist.params)) for dist in dist_predictions]

# Probability plot
# ==============================================================================
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(7, 3))
ax.plot(grid_X, probability_predictions)
ax.set_xticks(range(0, 25))
ax.set_title('Probability that consumption exceeds 22.5 MWh',
             fontdict={'fontsize':13})
ax.set_xlabel('Hour of the day')
ax.set_ylabel('Probability');

# Data handling
# ==============================================================================
import pandas as pd
import numpy as np
from tqdm import tqdm

# Plotting
# ==============================================================================
import matplotlib.pyplot as plt
plt.style.use('fivethirtyeight')
plt.rcParams['lines.linewidth'] = 1.5
plt.rcParams['font.size'] = 8

# Modeling
# ==============================================================================
from ngboost import NGBRegressor
from ngboost import scores
from ngboost import distns
from sklearn.tree import DecisionTreeRegressor
from sklearn.model_selection import GridSearchCV
from scipy import stats
import multiprocessing

# Warnings configuration
# ==============================================================================
import warnings
warnings.filterwarnings('once')

# Load data
# ==============================================================================
url = (
    'https://raw.githubusercontent.com/JoaquinAmatRodrigo/'
    'Estadistica-machine-learning-python/master/data/dbbmi.csv'
)
data = pd.read_csv(url)
data.head(3)

# Plot distribution of bmi values
# ==============================================================================
fig, ax = plt.subplots(figsize=(6, 3))
ax.hist(data.bmi, bins=30, density=True, color='#3182bd', alpha=0.8)
ax.set_title('Distribution of body mass index values')
ax.set_xlabel('bmi')
ax.set_ylabel('density');

# Plot distribution of body mass index as a function of age
# ==============================================================================
fig, axs = plt.subplots(ncols=2, nrows=1, figsize=(8, 4), sharey=True)
data[data.age <= 2.5].plot(
    x    = 'age',
    y    = 'bmi',
    c    = 'black',
    kind = "scatter",
    alpha = 0.1,
    ax   = axs[0]
)
axs[0].set_ylim([10, 30])
axs[0].set_title('Age <= 2.5 years')

data[data.age > 2.5].plot(
    x    = 'age',
    y    = 'bmi',
    c    = 'black',
    kind = "scatter",
    alpha = 0.1,
    ax   = axs[1]
)
axs[1].set_title('Age > 2.5 years')

fig.tight_layout()
plt.subplots_adjust(top = 0.85)
fig.suptitle('Distribution of body mass index as a function of age', fontsize = 12);

def fit_plot_distribution(
    x: np.ndarray, 
    distribution_name, 
    ax=None, 
    verbose: bool = True,
):
    '''
    This function overlays the density curve(s) of one or more distributions
    on top of the data histogram.
    
    Parameters
    ----------
    x : array_like
        Data used to fit the distribution.
        
    distribution_name : str or list
        Name of a distribution or list of distribution names available
        in `scipy.stats`.
    
    ax : matplotlib.axes.Axes, optional
        Axes on which to draw the plot. If None, a new one is created.
    
    verbose : bool, optional
        If True, prints fitting results. For multiple distributions,
        prints error warnings (default True).
        
    Returns
    -------
    ax: matplotlib.axes.Axes
        Generated plot
        
    '''
    
    # Input validation
    x = np.asarray(x)
    if x.size == 0:
        raise ValueError("Array 'x' cannot be empty")
    if not np.all(np.isfinite(x)):
        raise ValueError("Array 'x' contains NaN or infinite values")
    
    if isinstance(distribution_name, str):
        distribution_name = [distribution_name]
    if not isinstance(distribution_name, (list, tuple)):
        raise TypeError("Parameter 'distribution_name' must be a string, list, or tuple")
    
    # Create plot
    if ax is None:
        fig, ax = plt.subplots(figsize=(7,4))
    
    # Plot histogram
    ax.hist(x=x, density=True, bins=30, color="#3182bd", alpha=0.5)
    ax.plot(x, np.full_like(x, -0.01), '|k', markeredgewidth=1)
    ax.set_title('Distribution fit' if len(distribution_name) == 1 else 
                 'Distribution fits')
    ax.set_xlabel('x')
    ax.set_ylabel('Probability density')
    
    for name in tqdm(distribution_name):
        if not hasattr(stats, name):
            if verbose:
                print(f"⚠️  Distribution '{name}' does not exist in scipy.stats")
            continue
        
        try:

            distribution = getattr(stats, name)
            parameters = distribution.fit(data=x)
            
            is_continuous = isinstance(distribution, stats.rv_continuous)
            pdf_func = distribution.pdf if is_continuous else distribution.pmf
            
            x_hat = np.linspace(min(x), max(x), num=100)
            y_hat = pdf_func(x_hat, *parameters)
            ax.plot(x_hat, y_hat, linewidth=2, label=distribution.name)
            
        except Exception as e:
            if verbose:
                print(f"⚠️  Error fitting distribution '{name}': {type(e).__name__}: {e}")
    
    ax.legend(loc='best')

    return ax

fig, ax = plt.subplots(figsize=(6, 3))
fit_plot_distribution(
    x=data.bmi.to_numpy(),
    distribution_name=['lognorm', 'norm', 'expon'],
    ax=ax
);

X = data.age.to_numpy()
y = data.bmi.to_numpy()

# Hyperparameter grid
# ==============================================================================
b1 = DecisionTreeRegressor(criterion='friedman_mse', max_depth=2)
b2 = DecisionTreeRegressor(criterion='friedman_mse', max_depth=4)
param_grid = {
    'Base': [b1, b2],
    'n_estimators': [100, 500, 1000]
}

# Cross-validation search
# ==============================================================================
grid = GridSearchCV(
        estimator = NGBRegressor(
                        Dist    = distns.LogNormal,
                        Score   = scores.LogScore,
                        verbose = False
                    ),
        param_grid = param_grid,
        scoring    = 'neg_root_mean_squared_error',  # Evaluates point prediction (expected value). To evaluate full distributional quality, neg-log-likelihood or CRPS would be used.
        n_jobs     = multiprocessing.cpu_count() - 1,
        cv         = 5, 
        verbose    = 0
       )

# Assign result to _ to suppress printed output
_ = grid.fit(X = X.reshape(-1, 1), y = y)

# Grid results
# ==============================================================================
results = pd.DataFrame(grid.cv_results_)
results.filter(regex = '(param.*|mean_t|std_t)')\
    .drop(columns = 'params')\
    .sort_values('mean_test_score', ascending = False)

# Best hyperparameters
# ==============================================================================
print("-----------------------------------")
print("Best hyperparameters found")
print("-----------------------------------")
print(grid.best_params_, ":", grid.best_score_, grid.scoring)

# Best model found
# ==============================================================================
model = grid.best_estimator_

# Compute the quantiles to define the central interval accumulating
# 90% of probability using the predicted parameters.
# ==============================================================================
grid_X = np.linspace(start=min(data.age), stop=max(data.age), num=2000)
dist_predictions = model.pred_dist(X=grid_X.reshape(-1, 1))
pred_quantile_10 = [stats.lognorm.ppf(q=0.1, **dist.params) for dist in dist_predictions]
pred_quantile_90 = [stats.lognorm.ppf(q=0.9, **dist.params) for dist in dist_predictions]
df_intervals = pd.DataFrame({
                    'age': grid_X,
                    'pred_quantile_10': pred_quantile_10,
                    'pred_quantile_90': pred_quantile_90
                })

# Plot quantile interval 10%-90%
# ==============================================================================
fig, axs = plt.subplots(ncols=2, nrows=1, figsize=(10, 4.5), sharey=True)

data[data.age <= 2.5].plot(
    x    = 'age',
    y    = 'bmi',
    c    = 'black',
    kind = "scatter",
    alpha = 0.1,
    ax   = axs[0]
)

df_intervals[df_intervals.age <= 2.5].plot(
    x    = 'age',
    y    = 'pred_quantile_10',
    c    = 'red',
    kind = "line",
    ax   = axs[0]
)

df_intervals[df_intervals.age <= 2.5].plot(
    x    = 'age',
    y    = 'pred_quantile_90',
    c    = 'red',
    kind = "line",
    ax   = axs[0]
)

axs[0].set_ylim([10, 30])
axs[0].set_title('Age <= 2.5 years')

data[data.age > 2.5].plot(
    x    = 'age',
    y    = 'bmi',
    c    = 'black',
    kind = "scatter",
    alpha = 0.1,
    ax   = axs[1]
)

df_intervals[df_intervals.age > 2.5].plot(
    x    = 'age',
    y    = 'pred_quantile_10',
    c    = 'red',
    kind = "line",
    ax   = axs[1]
)

df_intervals[df_intervals.age > 2.5].plot(
    x    = 'age',
    y    = 'pred_quantile_90',
    c    = 'red',
    kind = "line",
    ax   = axs[1]
)
axs[1].set_title('Age > 2.5 years')
axs[1].get_legend().remove()

fig.tight_layout()
plt.subplots_adjust(top = 0.85)
fig.suptitle('Distribution of body mass index as a function of age', fontsize = 14);

# Interval predictions for the available observations
# ==============================================================================
dist_predictions = model.pred_dist(X=X.reshape(-1, 1))
pred_quantile_10 = [stats.lognorm.ppf(q=0.1, **dist.params) for dist in dist_predictions]
pred_quantile_90 = [stats.lognorm.ppf(q=0.9, **dist.params) for dist in dist_predictions]

# Coverage of the predicted interval
# ==============================================================================
within_interval = np.where((y >= pred_quantile_10) & (y <= pred_quantile_90), True, False)
coverage = within_interval.mean()
print(f"Coverage of the predicted interval: {100 * coverage:.2f}%")

# 98% interval predictions for the available observations
# ==============================================================================
X = data.age.to_numpy()
y = data.bmi.to_numpy()

dist_predictions = model.pred_dist(X=X.reshape(-1, 1))
pred_quantile_01 = [stats.lognorm.ppf(q=0.01, **dist.params) for dist in dist_predictions]
pred_quantile_99 = [stats.lognorm.ppf(q=0.99, **dist.params) for dist in dist_predictions]
within_interval = np.where((y >= pred_quantile_01) & (y <= pred_quantile_99), True, False)

df_intervals = pd.DataFrame({
                    'age': X,
                    'bmi': y,
                    'pred_quantile_01': pred_quantile_01,
                    'pred_quantile_99': pred_quantile_99,
                    'within_interval': within_interval
                })

df_intervals.head()

# Anomaly plot
# ==============================================================================
fig, axs = plt.subplots(ncols=2, nrows=1, figsize=(10, 4.5), sharey=True)

df_intervals[df_intervals.age <= 2.5].plot(
    x    = 'age',
    y    = 'bmi',
    c    = 'black',
    kind = "scatter",
    alpha = 0.1,
    ax   = axs[0]
)

df_intervals[(df_intervals.age <= 2.5) & (df_intervals.within_interval == False)].plot(
    x    = 'age',
    y    = 'bmi',
    c    = 'red',
    kind = "scatter",
    label = 'anomaly',
    ax   = axs[0]
)

axs[0].set_ylim([10, 30])
axs[0].set_title('Age <= 2.5 years')

df_intervals[df_intervals.age > 2.5].plot(
    x    = 'age',
    y    = 'bmi',
    c    = 'black',
    kind = "scatter",
    alpha = 0.1,
    ax   = axs[1]
)

df_intervals[(df_intervals.age > 2.5) & (df_intervals.within_interval == False)].plot(
    x    = 'age',
    y    = 'bmi',
    c    = 'red',
    kind = "scatter",
    ax   = axs[1]
)

axs[1].set_title('Age > 2.5 years')
fig.tight_layout()
plt.subplots_adjust(top = 0.85)
fig.suptitle('Distribution of body mass index as a function of age', fontsize = 14);

import session_info
session_info.show(html=False)