# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/src/feature_engineering.ipynb.

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

# %% ../nbs/src/feature_engineering.ipynb 3
from typing import Tuple

import pandas as pd
from utilsforecast.compat import DataFrame, pl, pl_DataFrame
from utilsforecast.processing import (
    drop_index_if_pandas,
    horizontal_concat,
    maybe_compute_sort_indices,
    take_rows,
    vertical_concat,
)

from . import StatsForecast
from .core import _id_as_idx
from .models import MSTL, _predict_mstl_components

# %% ../nbs/src/feature_engineering.ipynb 4
def mstl_decomposition(
    df: DataFrame,
    model: MSTL,
    freq: str,
    h: int,
) -> Tuple[DataFrame, DataFrame]:
    """Decompose the series into trend and seasonal using the MSTL model.

    Parameters
    ----------
    df : pandas or polars DataFrame
        DataFrame with columns [`unique_id`, `ds`, `y`].
    model : statsforecast MSTL
        Model to use for the decomposition.
    freq : str
        Frequency of the data (pandas alias)
    h : int
        Forecast horizon.

    Returns
    -------
    train_df : pandas or polars DataFrame
        Original dataframe with the 'trend' and 'seasonal' columns added.
    X_df : pandas or polars DataFrame
        Future dataframe to be provided to the predict method through `X_df`.
    """
    if not isinstance(model, MSTL):
        raise ValueError(f"`model` must be an MSTL instance, got {type(model)}")
    sort_idxs = maybe_compute_sort_indices(df, "unique_id", "ds")
    if sort_idxs is not None:
        df = take_rows(df, sort_idxs)
    df = drop_index_if_pandas(df)
    sf = StatsForecast(models=[model], freq=freq)
    sf.fit(df=df)
    X_df = sf._make_future_df(h=h)
    if isinstance(X_df, pd.DataFrame) and _id_as_idx():
        X_df = X_df.reset_index()
    train_features = []
    future_features = []
    df_constructor = type(df)
    seas_cols = [c for c in sf.fitted_[0, 0].model_.columns if c.startswith("seasonal")]
    for fitted_model in sf.fitted_[:, 0]:
        train_features.append(fitted_model.model_[["trend"] + seas_cols])
        seas_comp = _predict_mstl_components(
            fitted_model.model_, h, model.season_length
        )
        future_df = df_constructor(
            {
                "trend": fitted_model.trend_forecaster.predict(h)["mean"],
                **dict(zip(seas_cols, seas_comp.T)),
            }
        )
        future_features.append(future_df)
    train_features = vertical_concat(train_features, match_categories=False)
    if isinstance(df, pl_DataFrame):
        train_features = pl.from_pandas(train_features)
    train_df = horizontal_concat([df, train_features])
    future_features = vertical_concat(future_features, match_categories=False)
    X_df = horizontal_concat([X_df, future_features])
    return train_df, X_df
