# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/modules/02_models/07_ets.ipynb.

# %% ../../nbs/modules/02_models/07_ets.ipynb #928028da
from __future__ import annotations
from typing import List, Dict, Optional, Callable, Tuple, Any, Union
from xml.parsers.expat import model
from sklearn.base import clone
from tabnanny import verbose
import numpy as np
import pandas as pd
import copy
import statsmodels.api as sm
from statsmodels.tsa.holtwinters import ExponentialSmoothing
from ..transformations import (box_cox_transform, back_box_cox_transform)
from ..model_selection import SplitTimeSeries
# dot not show warnings
import warnings
warnings.filterwarnings("ignore")
import copy
import statsmodels.api as sm
from statsmodels.tsa.holtwinters import ExponentialSmoothing
from scipy.special import logsumexp
from scipy.stats import t
import re # for regex escaping to build drop patterns
import warnings
warnings.filterwarnings("ignore")

class ets:
    def __init__(
        self,
        target_col: str,
        trend: Optional[str] = None,
        damped_trend: bool = False,
        seasonal: Optional[str] = None,
        seasonal_periods: Optional[int] = None,
        initialization_method: Optional[str] = "estimated",
        initial_level: Optional[float] = None,
        initial_trend: Optional[float] = None,
        initial_seasonal: Optional[list] = None,
        use_boxcox: Union[bool, str] = False,
        bounds: Optional[dict] = None,
        dates=None,
        freq: Optional[str] = None,
        missing: str = "none",
        optimized: bool = True,
        smoothing_level: Optional[float] = None,
        smoothing_trend: Optional[float] = None,
        smoothing_seasonal: Optional[float] = None,
        damping_trend: Optional[float] = None,
        remove_bias: bool = False,
        start_params=None,
        method: Optional[str] = None,
        minimize_kwargs: Optional[dict] = None,
        use_brute: bool = True,
        box_cox: bool = False,
        box_cox_lmda: Optional[float] = None,
        box_cox_biasadj: bool = False,
        fit_kwargs: Optional[dict] = None,
    ) -> None:
        
        """
        Holt-Winters Exponential Smoothing forecaster.

        A thin wrapper around ``statsmodels.tsa.holtwinters.ExponentialSmoothing``.

        Parameters
        ----------
        target_col : str
            Name of the target variable column.
        trend : {"add", "mul", "additive", "multiplicative"} or None
            Trend component type.
        damped_trend : bool, default False
            Whether to damp the trend.  Only meaningful when ``trend`` is not `None`.
        seasonal : {"add", "mul", "additive", "multiplicative"} or None
            Seasonal component type.
        seasonal_periods : int or None
            Number of periods in a complete seasonal cycle — e.g. 12 for monthly data with an annual cycle.  Required when ``seasonal`` is not ``None``.
        initialization_method : {"estimated", "heuristic", "legacy-heuristic", "known"} or None, default "estimated"
            How to initialise the recursions.  When `"known"` is chosen, `initial_level` (and `initial_trend` / `initial_seasonal` where applicable) must also be provided.
        initial_level : float or None, default None
            Initial level value.  Required when `initialization_method=`"known"`.
        initial_trend : float or None, default None
            Initial trend value.  Required when `initialization_method=`"known"` and the model has a trend component.
        initial_seasonal : list or None, default None
            Initial seasonal factors (length `seasonal_periods` or ``seasonal_periods - 1``).  Required when ``initialization_method="known"`` and the model is seasonal.
        bounds : dict or None, default None
            Parameter bounds passed to ``ExponentialSmoothing``, e.g. `{"smoothing_level": (0, 1)}`.
        dates : array-like of datetime or None, default None
            Datetime index for the series.  Inferred automatically when `endog` is a Pandas object with a `DatetimeIndex`.
        freq : str or None, default None
            Frequency of the time series (e.g. ``"M"``, ``"D"``).  Optional when `dates` is provided.
        missing : {"none", "drop", "raise"}, default "none"
            How to handle ``NaN`` values in the input series.
        optimized : bool, default True
            Estimate smoothing parameters by maximising the log-likelihood.
        smoothing_level : float or None, default None
            Fixed alpha value.  When set, this value is used directly and not optimised.
        smoothing_trend : float or None, default None
            Fixed beta value.  Only used when the model has a trend component.
        smoothing_seasonal : float or None, default None
            Fixed gamma value.  Only used when the model is seasonal.
        damping_trend : float or None, default None
            Fixed phi (damping) value.  Only used when ``damped_trend=True``.
        remove_bias : bool, default False
            Remove bias from forecast values by enforcing that the mean residual is zero.
        start_params : array-like or None, default None
            Starting parameter values for the optimiser.
        method : str or None, default None
            Optimisation method — one of `"L-BFGS-B"` (default), `"TNC"`, `"SLSQP"`, `"Powell"`, `"trust-constr"`, `"basinhopping"` (alias `"bh"`), or `"least_squares"` (alias `"ls"`).
        minimize_kwargs : dict or None, default None
            Extra keyword arguments forwarded to the chosen SciPy minimiser.
        use_brute : bool, default True
            Search for good starting values with a brute-force grid search before running the main optimiser.
        box_cox : bool, default False
            Apply a manual Box-Cox transformation to the target.
        box_cox_lmda : float or None, default None
            Lambda for the manual transform.  ``None`` means auto-estimate.
        box_cox_biasadj : bool, default False
            Bias adjustment when inverting the manual Box-Cox on forecasts.
        fit_kwargs : dict or None, default None
            Any additional keyword arguments forwarded verbatim to `ExponentialSmoothing.fit`.
        
        Returns
        -------
        ets
            Fitted model object with a ``forecast`` method for making predictions and properties for information criteria scores (AIC, BIC, etc.)
        """

        self.target_col    = target_col
        self.box_cox       = box_cox
        self.lamda         = box_cox_lmda
        self.biasadj       = box_cox_biasadj
        self.fit_kwargs    = fit_kwargs or {}

        # ── ExponentialSmoothing constructor params ────────────────────────────
        self.trend                = trend
        self.damped_trend         = damped_trend
        self.seasonal             = seasonal
        self.seasonal_periods     = seasonal_periods
        self.initialization_method = initialization_method
        self.initial_level        = initial_level
        self.initial_trend        = initial_trend
        self.initial_seasonal     = initial_seasonal
        self.use_boxcox           = use_boxcox
        self.bounds               = bounds
        self.dates                = dates
        self.freq                 = freq
        self.missing              = missing

        # ── ExponentialSmoothing.fit params ───────────────────────────────────
        self.optimized          = optimized
        self.smoothing_level    = smoothing_level
        self.smoothing_trend    = smoothing_trend
        self.smoothing_seasonal = smoothing_seasonal
        self.damping_trend      = damping_trend
        self.remove_bias        = remove_bias
        self.start_params       = start_params
        self.method             = method
        self.minimize_kwargs    = minimize_kwargs
        self.use_brute          = use_brute

        # ── placeholders set during fit ────────────────────────────────────────
        self.tuned_params   = None
        self.actuals        = None
        self.prob_forecasts = None

    # ─────────────────────────────────────────────────────────────────────────
    # DATA PREPARATION
    # ─────────────────────────────────────────────────────────────────────────

    def data_prep(self,
                  df: pd.DataFrame
                  ) -> pd.DataFrame:
        """
        Apply preprocessing and return a cleaned DataFrame.

        Parameters
        ----------
        df : pd.DataFrame
                Input DataFrame containing just the target column.

        Returns
        -------
        pd.DataFrame
        """
        dfc = df.copy()

        # ── manual Box-Cox ────────────────────────────────────────────────────
        if self.box_cox:
            self.is_zero = np.any(np.array(dfc[self.target_col]) < 1)
            trans_data, self.lamda = box_cox_transform(
                x=dfc[self.target_col], shift=self.is_zero,
                box_cox_lmda=self.lamda,
            )
            dfc[self.target_col] = trans_data

        return dfc[[self.target_col]].dropna()

    # ─────────────────────────────────────────────────────────────────────────
    # FIT
    # ─────────────────────────────────────────────────────────────────────────

    def fit(self,
            df: pd.DataFrame
            ) -> None:
        """
        Fit ``ExponentialSmoothing`` to the training data.

        Parameters
        ----------
        df : pd.DataFrame
            Training DataFrame containing the target column
        
        Returns
        -------
        None
        """

        model_df  = self.data_prep(df)
        self.y    = model_df[self.target_col].to_numpy()

        # ── Resolve locally — never mutate self.* so copy() preserves user intent ─
        #
        # trend=None  →  trend-related fit params are meaningless; zero them out
        #   • smoothing_trend  must be None  (no beta to optimise)
        #   • damped_trend     must be False (no trend to damp)
        #   • damping_trend    must be None  (no phi to optimise)
        #   • initial_trend    must be None  (no initial trend state)
        #
        # seasonal=None  →  seasonal-related fit params are meaningless; zero them
        #   • smoothing_seasonal must be None  (no gamma to optimise)
        #   • seasonal_periods   must be None  (no cycle length needed)
        #   • initial_seasonal   must be None  (no initial seasonal state)
        #
        # damped_trend=False  →  damping_trend must be None (no phi to optimise)
        #
        # These rules prevent statsmodels from raising errors when the parameter
        # space is explored during hyperparameter optimisation and conflicting
        # combinations are passed in.

        _trend    = self.trend
        _seasonal = self.seasonal

        # trend-related
        _damped_trend      = self.damped_trend      if _trend     is not None else False
        _smoothing_trend   = self.smoothing_trend   if _trend     is not None else None
        _damping_trend     = self.damping_trend     if (_trend    is not None and _damped_trend) else None
        _initial_trend     = self.initial_trend     if _trend     is not None else None

        # seasonal-related
        _smoothing_seasonal = self.smoothing_seasonal if _seasonal is not None else None
        _seasonal_periods   = self.seasonal_periods   if _seasonal is not None else None
        _initial_seasonal   = self.initial_seasonal   if _seasonal is not None else None

        # ── ExponentialSmoothing constructor kwargs ────────────────────────────
        hw_kwargs: dict = dict(
            endog                 = self.y,
            trend                 = _trend,
            damped_trend          = _damped_trend,
            seasonal              = _seasonal,
            seasonal_periods      = _seasonal_periods,
            initialization_method = self.initialization_method,
            use_boxcox            = self.use_boxcox,
            missing               = self.missing,
        )
        # optional — only pass when explicitly set to avoid overriding statsmodels defaults
        if self.initial_level  is not None: hw_kwargs["initial_level"]    = self.initial_level
        if _initial_trend      is not None: hw_kwargs["initial_trend"]    = _initial_trend
        if _initial_seasonal   is not None: hw_kwargs["initial_seasonal"] = _initial_seasonal
        if self.bounds         is not None: hw_kwargs["bounds"]           = self.bounds
        if self.dates          is not None: hw_kwargs["dates"]            = self.dates
        if self.freq           is not None: hw_kwargs["freq"]             = self.freq

        self.model = ExponentialSmoothing(**hw_kwargs)

        # ── ExponentialSmoothing.fit kwargs ───────────────────────────────────
        fit_kw: dict = dict(
            optimized    = self.optimized,
            remove_bias  = self.remove_bias,
            use_brute    = self.use_brute,
        )
        # optional smoothing coefficients — only pass when fixed by the user
        # (resolved values already account for trend/seasonal being None)
        if self.smoothing_level  is not None: fit_kw["smoothing_level"]    = self.smoothing_level
        if _smoothing_trend      is not None: fit_kw["smoothing_trend"]    = _smoothing_trend
        if _smoothing_seasonal   is not None: fit_kw["smoothing_seasonal"] = _smoothing_seasonal
        if _damping_trend        is not None: fit_kw["damping_trend"]      = _damping_trend
        if self.start_params     is not None: fit_kw["start_params"]       = self.start_params
        if self.method           is not None: fit_kw["method"]             = self.method
        if self.minimize_kwargs  is not None: fit_kw["minimize_kwargs"]    = self.minimize_kwargs
        # merge with any extra user-supplied kwargs (fit_kwargs wins on conflict)
        fit_kw.update(self.fit_kwargs)

        self.model_fit = self.model.fit(**fit_kw)

    # ─────────────────────────────────────────────────────────────────────────
    # INFORMATION CRITERIA  (delegated to statsmodels fitted result)
    # ─────────────────────────────────────────────────────────────────────────

    @property
    def aic(self) -> float:
        """Akaike Information Criterion."""
        return self.model_fit.aic

    @property
    def aicc(self) -> float:
        """Corrected AIC."""
        return self.model_fit.aicc

    @property
    def bic(self) -> float:
        """Bayesian Information Criterion."""
        return self.model_fit.bic

    @property
    def hqc(self) -> float:
        """Hannan-Quinn Criterion (computed from fitted log-likelihood)."""
        llf = self.model_fit.llf
        k   = self.model_fit.df_model
        n   = len(self.y)
        return -2 * llf + 2 * k * np.log(np.log(n))

    def copy(self):
        return copy.deepcopy(self)

    # ─────────────────────────────────────────────────────────────────────────
    # FORECAST
    # ─────────────────────────────────────────────────────────────────────────

    def forecast(
        self,
        H: int,
        exog: Optional[pd.DataFrame] = None,
    ) -> np.ndarray:
        """
        Multi-step forecast.

        Parameters
        ----------
        H : int
            Forecast horizon.
        exog : pd.DataFrame or None
            Accepted for API consistency with other models but silently ignored — ETS forecasts do not use exogenous variables.

        Returns
        -------
        np.ndarray
            Forecast values of length `H`.
        """
        forecasts = np.array(self.model_fit.forecast(H))

        # ── non-negativity ────────────────────────────────────────────────────
        forecasts = np.array([max(0, x) for x in forecasts])

        # ── invert manual Box-Cox ─────────────────────────────────────────────
        if self.box_cox:
            forecasts = back_box_cox_transform(
                y_pred=forecasts, lmda=self.lamda,
                shift=self.is_zero, box_cox_biasadj=self.biasadj,
            )

        return forecasts

    # ─────────────────────────────────────────────────────────────────────────
    # CROSS-VALIDATION  (identical signature and logic to arima.cross_validate)
    # ─────────────────────────────────────────────────────────────────────────

    def cross_validate(
        self,
        df: pd.DataFrame,
        cv_split: int,
        test_size: int,
        metrics: List[Callable],
        step_size: int = 1,
        h_split_point: Optional[int] = None,
        cv_df: bool = False,
    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        Run time-series cross-validation.

        Parameters
        ----------
        df : pd.DataFrame
            Full dataset
        cv_split : int
            Number of CV folds.
        test_size : int
            Test window size per fold.
        metrics : list of callable
            Metric functions (e.g. ``[MAE, RMSE]``).
        step_size : int, default 1
            Step size to advance the test window each fold.
        h_split_point : int or None, default None
            Split the test window into two sub-horizons for separate short- and long-term evaluation.
        cv_df : bool, default False
            If `True``, return a fold-level prediction DataFrame alongside the summary.

        Returns
        -------
        overall_performance : pd.DataFrame and cv_predictions : pd.DataFrame
            Summary DataFrame with mean metric scores across folds, and (optionally) a fold-level DataFrame with true vs. predicted values for each fold.
        """
        cv_df_ = pd.DataFrame()
        tscv = SplitTimeSeries(
            n_splits=cv_split, test_size=test_size, step_size=step_size
        )
        metrics_dict = {m.__name__: [] for m in metrics}
        if h_split_point is not None:
            metrics_dict1 = {m.__name__: [] for m in metrics}
            metrics_dict2 = {m.__name__: [] for m in metrics}

        for idx, (train_index, test_index) in enumerate(tscv.split(df)):
            train, test = df.iloc[train_index], df.iloc[test_index]
            x_test = test.drop(columns=[self.target_col])
            y_test = np.array(test[self.target_col])

            self.fit(train)
            bb_forecast = np.array(self.forecast(test_size, x_test))

            for m in metrics:
                if m.__name__ in ["MASE", "SMAE", "SRMSE", "RMSSE"]:
                    eval_val = m(y_test, bb_forecast, train[self.target_col])
                else:
                    eval_val = m(y_test, bb_forecast)
                metrics_dict[m.__name__].append(eval_val)

            if h_split_point is not None and isinstance(h_split_point, int):
                y_test_1, y_test_2 = y_test[:h_split_point], y_test[h_split_point:]
                bb_forecast_1, bb_forecast_2 = (
                    bb_forecast[:h_split_point], bb_forecast[h_split_point:]
                )
                for m in metrics:
                    if m.__name__ in ["MASE", "SMAE", "SRMSE", "RMSSE"]:
                        eval_val1 = m(y_test_1, bb_forecast_1, np.array(train[self.target_col]))
                        eval_val2 = m(y_test_2, bb_forecast_2, np.array(train[self.target_col]))
                    else:
                        eval_val1 = m(y_test_1, bb_forecast_1)
                        eval_val2 = m(y_test_2, bb_forecast_2)
                    metrics_dict1[m.__name__].append(eval_val1)
                    metrics_dict2[m.__name__].append(eval_val2)

            if cv_df:
                split_results = {
                    "cutoff": np.repeat(test.index[0], len(test)),
                    "index":  test.index,
                    "split":  np.repeat(f"fold_{idx + 1}", len(test)),
                    "y_true": y_test,
                    "y_pred": bb_forecast,
                }
                cv_df_ = pd.concat(
                    [cv_df_, pd.DataFrame(split_results)], ignore_index=True
                )

        overall_performance = pd.DataFrame(
            [[m.__name__, np.mean(metrics_dict[m.__name__])] for m in metrics],
            columns=["eval_metric", "score"],
        )

        if h_split_point is not None and isinstance(h_split_point, int):
            perf_1_df = pd.DataFrame(
                [[m.__name__, np.mean(metrics_dict1[m.__name__])] for m in metrics],
                columns=["eval_metric", f"score_before_{h_split_point}"],
            )
            perf_2_df = pd.DataFrame(
                [[m.__name__, np.mean(metrics_dict2[m.__name__])] for m in metrics],
                columns=["eval_metric", f"score_after_{h_split_point}"],
            )
            overall_performance = (
                overall_performance
                .merge(perf_1_df, on="eval_metric")
                .merge(perf_2_df, on="eval_metric")
            )

        return overall_performance, cv_df_

# %% auto #0
__all__ = ['ets']
