# Tratamiento de datos
# ==============================================================================
import pandas as pd
import numpy as np

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

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

# Configuración warnings
# ==============================================================================
import warnings
warnings.filterwarnings('once')

# Datos simulados
# ==============================================================================
# Simulación ligeramente modificada del ejemplo publicado en
# XGBoostLSS – An extension of XGBoost to probabilistic forecasting Alexander März

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

# Distribución normal con varianza cambiante
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)
        )
    )

# Cálculo del cuantil 0.1 y 0.9 para cada posición de x simulada
# ==============================================================================
cuantil_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)
                )
             )

cuantil_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)
                )
             )

# Gráfico
# ==============================================================================
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, cuantil_10, c = "black")
ax.plot(x, cuantil_90, c = "black")
ax.fill_between(x, cuantil_10, cuantil_90, alpha=0.3, color='red',
                label='Intervalo cuantílico teórico 0.1-0.9')
ax.set_xticks(range(0, 25))
ax.set_title('Evolución del consumo eléctrico a lo largo del día')
ax.set_xlabel('Hora del día')
ax.set_ylabel('Consumo eléctrico (MWh)')
plt.legend();

# Ajuste del modelo
# ==============================================================================
modelo = 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,
        )

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

# Predicción point estimate
# ==============================================================================
predicciones_point = modelo.predict(X=x.reshape(-1, 1))
predicciones_point

# Predicción distribución condicional
# ==============================================================================
predicciones_dist = modelo.pred_dist(X=x.reshape(-1, 1))
predicciones_dist

# Dentro de cada objeto distns.distn se encuentran los parámetros estimados
predicciones_dist[0].params

# Parámetros estimados para cada hora
# ==============================================================================
df_parametros = pd.DataFrame([dist.params for dist in predicciones_dist])
df_parametros['hora'] = x
df_parametros

# Cálculo de los cuantiles teóricos para establecer el intervalo central que 
# acumula un 90% de probabilidad empleando los parámetros predichos.
# ==============================================================================
grid_X = np.linspace(start=min(x), stop=max(x), num=1000)
predicciones_dist = modelo.pred_dist(X=grid_X.reshape(-1, 1))
pred_cuantil_10 = [stats.norm.ppf(q=0.1, **dist.params) for dist in predicciones_dist]
pred_cuantil_90 = [stats.norm.ppf(q=0.9, **dist.params) for dist in predicciones_dist]

# Gráfico intervalo cuantílico 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, cuantil_10, c = "black")
ax.plot(x, cuantil_90, c = "black")
ax.fill_between(x, cuantil_10, cuantil_90, alpha=0.3, color='red',
                label='Intervalo cuantílico teórico 0.1-0.9')
ax.plot(grid_X, pred_cuantil_10, c = "blue", label='Intervalo cuantílico predicho')
ax.plot(grid_X, pred_cuantil_90, c = "blue")

ax.set_xticks(range(0, 25))
ax.set_title('Evolución del consumo eléctrico a lo largo del día')
ax.set_xlabel('Hora del día')
ax.set_ylabel('Consumo eléctrico (MWh)')
plt.legend();

# Predicción intervalos para las observaciones disponibles
# ==============================================================================
predicciones_dist = modelo.pred_dist(X=x.reshape(-1, 1))
pred_cuantil_10 = [stats.norm.ppf(q=0.1, **dist.params) for dist in predicciones_dist]
pred_cuantil_90 = [stats.norm.ppf(q=0.9, **dist.params) for dist in predicciones_dist]

# Cobertura del intervalo predicho
# ==============================================================================
dentro_intervalo = np.where((y >= pred_cuantil_10) & (y <= pred_cuantil_90), True, False)
cobertura = dentro_intervalo.mean()
print(f"Cobertura del intervalo predicho: {100 * cobertura:.2f}%")

# Predicción de los parámetros de la distribución cuando x=8
# ==============================================================================
distribucion = modelo.pred_dist(X=np.array([8]).reshape(-1, 1))
distribucion.params

# Probabilidad de que el consumo supere un valor de 22.5
# ==============================================================================
100 * (1 - stats.norm.cdf(x=22.5, **distribucion.params))

# Predicción de probabilidad en todo el rango horario
# ==============================================================================
grid_X = np.linspace(start=0, stop=24, num=500)
predicciones_dist = modelo.pred_dist(X=grid_X.reshape(-1, 1))
predicciones_proba = [(1 - stats.norm.cdf(x=22.5, **dist.params)) for dist in predicciones_dist]

# Gráfico probabilidades
# ==============================================================================
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(7, 3))
ax.plot(grid_X, predicciones_proba)
ax.set_xticks(range(0, 25))
ax.set_title('Probabilidad de que el consumo supere los 22.5 MWh',
             fontdict={'fontsize':13})
ax.set_xlabel('Hora del día')
ax.set_ylabel('Probabilidad');

# Tratamiento de datos
# ==============================================================================
import pandas as pd
import numpy as np
from tqdm import tqdm

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

# Modelado
# ==============================================================================
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

# Configuración warnings
# ==============================================================================
import warnings
warnings.filterwarnings('once')

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

# Gráfico distribución de los valores de bmi
# ==============================================================================
fig, ax = plt.subplots(figsize=(6, 3))
ax.hist(datos.bmi, bins=30, density=True, color='#3182bd', alpha=0.8)
ax.set_title('Distribución valores índice masa corporal')
ax.set_xlabel('bmi')
ax.set_ylabel('densidad');

# Gráfico distribución índice masa corporal en función de la edad
# ==============================================================================
fig, axs = plt.subplots(ncols=2, nrows=1, figsize=(8, 4), sharey=True)
datos[datos.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('Edad <= 2.5 años')

datos[datos.age > 2.5].plot(
    x    = 'age',
    y    = 'bmi',
    c    = 'black',
    kind = "scatter",
    alpha = 0.1,
    ax   = axs[1]
)
axs[1].set_title('Edad > 2.5 años')

fig.tight_layout()
plt.subplots_adjust(top = 0.85)
fig.suptitle('Distribución índice masa corporal en función de la edad', fontsize = 12);

def fit_plot_distribution(
    x: np.ndarray, 
    nombre_distribucion, 
    ax=None, 
    verbose: bool = True,
):
    '''
    Esta función superpone la(s) curva(s) de densidad de una o varias 
    distribuciones con el histograma de los datos.
    
    Parameters
    ----------
    x : array_like
        datos con los que ajustar la distribución.
        
    nombre_distribucion : str o list
        nombre de una distribución o lista de nombres de distribuciones 
        disponibles en `scipy.stats`.
    
    ax : matplotlib.axes.Axes, optional
        eje sobre el que realizar el gráfico. Si es None, se crea uno nuevo.
    
    verbose : bool, optional
        Si True, imprime los resultados del ajuste. Para múltiples distribuciones,
        imprime advertencias de errores (default True).
        
    Returns
    -------
    ax: matplotlib.axes.Axes
        gráfico creado
        
    '''
    
    # Validación de entrada
    x = np.asarray(x)
    if x.size == 0:
        raise ValueError("El array 'x' no puede estar vacío")
    if not np.all(np.isfinite(x)):
        raise ValueError("El array 'x' contiene valores NaN o infinitos")
    
    if isinstance(nombre_distribucion, str):
        nombre_distribucion = [nombre_distribucion]
    if not isinstance(nombre_distribucion, (list, tuple)):
        raise TypeError("El parámetro 'nombre_distribucion' debe ser un string, lista o tupla")
    
    # Crear gráfico
    if ax is None:
        fig, ax = plt.subplots(figsize=(7,4))
    
    # Graficar histograma
    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('Ajuste distribución' if len(nombre_distribucion) == 1 else 
                 'Ajuste distribuciones')
    ax.set_xlabel('x')
    ax.set_ylabel('Densidad de probabilidad')
    
    for nombre in tqdm(nombre_distribucion):
        if not hasattr(stats, nombre):
            if verbose:
                print(f"⚠️  La distribución '{nombre}' no existe en scipy.stats")
            continue
        
        try:

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

    return ax

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

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

# Grid de hiperparámetros
# ==============================================================================
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]
}

# Búsqueda por validación cruzada
# ==============================================================================
grid = GridSearchCV(
        estimator = NGBRegressor(
                        Dist    = distns.LogNormal,
                        Score   = scores.LogScore,
                        verbose = False
                    ),
        param_grid = param_grid,
        scoring    = 'neg_root_mean_squared_error',  # Evalúa la predicción puntual (valor esperado). Para evaluar la calidad distribucional completa se usaría neg-log-likelihood o CRPS.
        n_jobs     = multiprocessing.cpu_count() - 1,
        cv         = 5, 
        verbose    = 0
       )

# Se asigna el resultado a _ para que no se imprima por pantalla
_ = grid.fit(X = X.reshape(-1, 1), y = y)

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

# Mejores hiperparámetros
# ==============================================================================
print("-----------------------------------")
print("Mejores hiperparámetros encontrados")
print("-----------------------------------")
print(grid.best_params_, ":", grid.best_score_, grid.scoring)

# Mejor modelo encontrado
# ==============================================================================
modelo = grid.best_estimator_

# Cálculo de los cuantiles teóricos para establecer el intervalo central que 
# acumula un 90% de probabilidad empleando los parámetros predichos.
# ==============================================================================
grid_X = np.linspace(start=min(datos.age), stop=max(datos.age), num=2000)
predicciones_dist = modelo.pred_dist(X=grid_X.reshape(-1, 1))
pred_cuantil_10 = [stats.lognorm.ppf(q=0.1, **dist.params) for dist in predicciones_dist]
pred_cuantil_90 = [stats.lognorm.ppf(q=0.9, **dist.params) for dist in predicciones_dist]
df_intervalos = pd.DataFrame({
                    'age': grid_X,
                    'pred_cuantil_10': pred_cuantil_10,
                    'pred_cuantil_90': pred_cuantil_90
                })

# Gráfico intervalo cuantílico 10%-90%
# ==============================================================================
fig, axs = plt.subplots(ncols=2, nrows=1, figsize=(10, 4.5), sharey=True)

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

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

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

axs[0].set_ylim([10, 30])
axs[0].set_title('Edad <= 2.5 años')

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

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

df_intervalos[df_intervalos.age > 2.5].plot(
    x    = 'age',
    y    = 'pred_cuantil_90',
    c    = 'red',
    kind = "line",
    ax   = axs[1]
)
axs[1].set_title('Edad > 2.5 años')
axs[1].get_legend().remove()

fig.tight_layout()
plt.subplots_adjust(top = 0.85)
fig.suptitle('Distribución índice masa corporal en función de la edad', fontsize = 14);

# Predicción intervalos para las observaciones disponibles
# ==============================================================================
predicciones_dist = modelo.pred_dist(X=X.reshape(-1, 1))
pred_cuantil_10 = [stats.lognorm.ppf(q=0.1, **dist.params) for dist in predicciones_dist]
pred_cuantil_90 = [stats.lognorm.ppf(q=0.9, **dist.params) for dist in predicciones_dist]

# Cobertura del intervalo predicho
# ==============================================================================
dentro_intervalo = np.where((y >= pred_cuantil_10) & (y <= pred_cuantil_90), True, False)
cobertura = dentro_intervalo.mean()
print(f"Cobertura del intervalo predicho: {100 * cobertura:.2f}%")

# Predicción intervalos del 98% para las observaciones disponibles
# ==============================================================================
X = datos.age.to_numpy()
y = datos.bmi.to_numpy()

predicciones_dist = modelo.pred_dist(X=X.reshape(-1, 1))
pred_cuantil_01 = [stats.lognorm.ppf(q=0.01, **dist.params) for dist in predicciones_dist]
pred_cuantil_99 = [stats.lognorm.ppf(q=0.99, **dist.params) for dist in predicciones_dist]
dentro_intervalo = np.where((y >= pred_cuantil_01) & (y <= pred_cuantil_99), True, False)

df_intervalos = pd.DataFrame({
                    'age': X,
                    'bmi': y,
                    'pred_cuantil_01': pred_cuantil_01,
                    'pred_cuantil_99': pred_cuantil_99,
                    'dentro_intervalo': dentro_intervalo
                })

df_intervalos.head()

# Gráfico anomalías
# ==============================================================================
fig, axs = plt.subplots(ncols=2, nrows=1, figsize=(10, 4.5), sharey=True)

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

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

axs[0].set_ylim([10, 30])
axs[0].set_title('Edad <= 2.5 años')

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

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

axs[1].set_title('Edad > 2.5 años')
fig.tight_layout()
plt.subplots_adjust(top = 0.85)
fig.suptitle('Distribución índice masa corporal en función de la edad', fontsize = 14);

import session_info
session_info.show(html=False)