"""Références simples (baselines) pour comparer nos modèles prédictifs."""

from __future__ import annotations

import pandas as pd

from .features import _steps_from_minutes, DEFAULT_BASE_FREQ_MINUTES


def _aligned_target(series: pd.Series, steps_ahead: int) -> pd.DataFrame:
    """
    Construit un DataFrame avec la cible (y_true) et l'indice de base pour les prédictions.

    y_true est la série décalée dans le futur (shift négatif), ce qui aligne chaque
    valeur de référence avec le point de départ utilisé pour prédire.
    """
    y_true = series.shift(-steps_ahead)
    return pd.DataFrame({"y_true": y_true})


def persistence_baseline(
    series: pd.Series,
    *,
    horizon_minutes: int,
    base_freq_minutes: int = DEFAULT_BASE_FREQ_MINUTES,
) -> pd.DataFrame:
    """
    Baseline de persistance : on suppose que la prochaine valeur sera identique à la dernière observée.

    Retourne un DataFrame avec y_true (valeur future) et y_pred (valeur observée au temps présent),
    alignés pour l'horizon demandé.
    """
    steps = _steps_from_minutes(horizon_minutes, base_freq_minutes)
    frame = _aligned_target(series, steps)
    frame["y_pred"] = series
    return frame.dropna()


def moving_average_baseline(
    series: pd.Series,
    *,
    horizon_minutes: int,
    window_minutes: int = 60,
    base_freq_minutes: int = DEFAULT_BASE_FREQ_MINUTES,
) -> pd.DataFrame:
    """
    Baseline de moyenne mobile : on prolonge la moyenne des dernières valeurs pour prévoir la prochaine.
    """
    steps = _steps_from_minutes(horizon_minutes, base_freq_minutes)
    window_steps = _steps_from_minutes(window_minutes, base_freq_minutes)
    rolling_mean = series.rolling(window=window_steps, min_periods=max(1, window_steps // 2)).mean()

    frame = _aligned_target(series, steps)
    frame["y_pred"] = rolling_mean
    return frame.dropna()


def hourly_climatology_baseline(
    train_series: pd.Series,
    eval_index: pd.DatetimeIndex,
    *,
    horizon_minutes: int,
) -> pd.Series:
    """
    Baseline de climatologie horaire : on prédit la moyenne observée (sur la partie train)
    pour l'heure du jour correspondant à l'horizon visé.

    Retourne une série alignée sur eval_index, avec une prédiction pour chaque ligne.
    """
    if not isinstance(eval_index, pd.DatetimeIndex):
        raise TypeError("eval_index doit être un DatetimeIndex.")

    climatology_by_hour = train_series.groupby(train_series.index.hour).mean()
    # Heure cible : on ajoute l'horizon aux timestamps pour récupérer l'heure de la cible
    target_hours = (eval_index + pd.to_timedelta(horizon_minutes, "minutes")).hour
    preds = pd.Series(target_hours, index=eval_index).map(climatology_by_hour)
    return preds