# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/modules/03_model_selection.ipynb.

# %% auto #0
__all__ = ['SplitTimeSeries', 'hyperopt_tune', 'optuna_tune', 'forward_feature_selection', 'backward_feature_selection',
           'mv_forward_feature_selection', 'mv_backward_feature_selection', 'ms_arr_forward_feature_selection',
           'ms_arr_backward_feature_selection']

# %% ../nbs/modules/03_model_selection.ipynb #623ad0b6
import numpy as np
from typing import List, Dict, Optional, Callable, Tuple, Any, Union

#------------------------------------------------------------------------------
# Parametric Time Series Split
#------------------------------------------------------------------------------
class SplitTimeSeries:

    def __init__(self,
                 n_splits: int,
                 test_size: int,
                 step_size: Optional[int] = None
                 ) -> None:

        """
        A time series cross-validator that generates train/test splits with a fixed test size and a configurable step size.

        Parameters
        ----------
        n_splits : int
            Number of splits to generate.
        test_size : int
            The number of samples in each test set.
        step_size : int, optional
            The number of samples to move the test set forward for each split. If None, it defaults to the test_size, meaning non-overlapping test sets.
        """
        self.test_size = test_size
        self.step_size = test_size if step_size is None else step_size
        self.n_splits = n_splits

    def split(self, X):
        n_samples = len(X)
        split_starts = []
        # Start the last test set at the last possible position
        last_test_start = n_samples - self.test_size
        # Build test starts, moving backward with fixed step_size
        current = last_test_start
        while current >= 0:
            split_starts.append(current)
            current -= self.step_size
        split_starts = split_starts[::-1]  # Reverse to start from earliest

        # Use only the last n_splits
        split_starts = split_starts[-self.n_splits:]

        for test_start in split_starts:
            test_end = test_start + self.test_size
            train_index = np.arange(0, test_start)
            test_index = np.arange(test_start, test_end)
            yield train_index, test_index


# %% ../nbs/modules/03_model_selection.ipynb #da8122f3
from pyexpat import model
import pandas as pd
import numpy as np
from scipy import stats
from numba import jit
from scipy.stats import boxcox
from scipy.special import inv_boxcox
from sklearn.linear_model import LinearRegression
from numba import jit
##Stationarity Check
from statsmodels.tsa.stattools import adfuller, kpss
from hyperopt import fmin, tpe, hp, Trials, STATUS_OK, space_eval
from statsmodels.tsa.holtwinters import ExponentialSmoothing
import matplotlib.pyplot as plt
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statistics import NormalDist
import warnings
warnings.filterwarnings("ignore")
from statsmodels.tsa.statespace.sarimax import SARIMAX
from tqdm import tqdm_notebook
from itertools import product
from typing import List, Dict, Optional, Callable, Tuple, Any, Union

def hyperopt_tune(
    model: object,
    df: pd.DataFrame,
    cv_split: int,
    test_size: int,
    eval_metric: Callable,
    param_space: dict,
    step_size: int = None,
    eval_num=100,
    verbose=False
    ) -> Tuple[Dict[str, Any], List[int], List[str]]:

    """
    Tune forecasting model hyperparameters using time series cross-validation and hyperopt.

    Parameters
    ----------
    model : object
        Forecasting model object with .fit and .forecast methods and relevant attributes.
    df : pd.DataFrame
        Time series data with a datetime index and a target column and optionally exogenous features.
    cv_split : int
        Number of cross-validation splits.
    test_size : int
        Number of samples in each test set.
    eval_metric : Callable
        Evaluation metric function.
    step_size : int, optional
        Step size to move the test window forward in each split.
    param_space : dict
        Hyperparameter search space for the forecasting model.
    eval_num : int, optional
        Number of hyperparameter combinations to evaluate. Default is 100.
    verbose : bool, optional
        Whether to print the evaluation metric for each hyperparameter combination. Default is False.
    
    Returns
    -------
    Tuple[Dict[str, Any], List[int], List[str]]
        A tuple containing the best hyperparameters, selected lags, and selected transforms.
    """
    tscv = SplitTimeSeries(n_splits=cv_split, test_size=test_size, step_size=step_size)

    def _set_model_params(params):
        # Handle special model parameters that are not passed to model constructor
        # and must be set directly on the forecasting model object
        if "lags" in params:
            if isinstance(params["lags"], int):
                model.n_lag = list(range(1, params["lags"] + 1))
            elif isinstance(params["lags"], list):
                model.n_lag = params["lags"]
        if "box_cox" in params:
            model.box_cox = params["box_cox"]
        if "box_cox_lmda" in params:
            model.lamda = params["box_cox_lmda"]
        if "box_cox_biasadj" in params:
            model.biasadj = params["box_cox_biasadj"]

    def _get_model_params_for_fit(params):
        # Exclude special parameters that should not be passed to the model constructor
        skip_keys = {"box_cox", "lags", "box_cox_lmda", "box_cox_biasadj"}

        return {k: v for k, v in params.items() if k not in skip_keys}
    

    def objective(params):
        _set_model_params(params)
        if isinstance(model.model, LinearRegression):
            # For LinearRegression, we don't need to set model_params
            model_params = None
        else:
            # For other models, get the parameters to set
            model_params = _get_model_params_for_fit(params)

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

            if model_params is not None:
                model.model.set_params(**model_params)
            model.fit(train)
            
            y_pred = model.forecast(len(y_test), x_test)

            #Evaluate using the specified metric
            if eval_metric.__name__ in ["MASE", "SMAE", "SRMSE", "RMSSE"]:
                score = eval_metric(y_test,
                                    y_pred,
                                    train[model.target_col])
            else:
                score = eval_metric(y_test,y_pred)
            metrics.append(score)

        mean_score = np.mean(metrics)
        if verbose:
            print("Score:", mean_score)
        return {"loss": mean_score, "status": STATUS_OK}

    trials = Trials()
    best_hyperparams = fmin(
        fn=objective,
        space=param_space,
        algo=tpe.suggest,
        max_evals=eval_num,
        trials=trials,
    )

    # if lags are in the param space, extract the best lags and extract remain parameters for the model
    if "lags" in param_space:
        best_lags = space_eval(param_space, best_hyperparams)["lags"]
        model_parameters = {k: v for k, v in space_eval(param_space, best_hyperparams).items() if k != "lags"}
    else:
        best_lags = None
        model_parameters = space_eval(param_space, best_hyperparams)

    # model.tuned_params = [
    #     space_eval(param_space, {k: v[0] for k, v in t["misc"]["vals"].items()})
    #     for t in trials.trials]

    return model_parameters, best_lags


# %% ../nbs/modules/03_model_selection.ipynb #42ad67aa
import optuna
def optuna_tune(
    model: object,
    df: pd.DataFrame,
    cv_split: int,
    test_size: int,
    eval_metric: Callable,
    param_space: Dict[str, Any],
    step_size: int = None,
    eval_num: int = 100,
    verbose: bool = False,
) -> Tuple[Dict[str, Any], Any]:
    """
    Tune forecasting model hyperparameters using time series cross-validation and Optuna.
 
    Parameters
    ----------
    model : object
        Forecasting model with .fit and .forecast methods.
    df : pd.DataFrame
        Time series data (datetime index, target column, optional exogenous features).
    cv_split : int
        Number of cross-validation splits.
    test_size : int
        Number of samples in each test fold.
    eval_metric : Callable
        Metric function to minimise.
    param_space : dict
        Each value must be a callable that accepts an Optuna `trial` and returns a value.
    step_size : int, optional
        Step size between CV folds.
    eval_num : int, optional
        Number of Optuna trials. Default 100.
    verbose : bool, optional
        Print score for every trial. Default False.
 
    Returns
    -------
    Tuple[Dict[str, Any], Any]
        Best hyperparameters and best lags (if 'lags' is in param_space).

    """
    tscv = SplitTimeSeries(n_splits=cv_split, test_size=test_size, step_size=step_size)
 
    _skip = {"box_cox", "lags", "box_cox_lmda", "box_cox_biasadj"}
 
    def _set_model_params(params: dict):
        if "lags" in params:
            lags = params["lags"]
            model.n_lag = list(range(1, lags + 1)) if isinstance(lags, int) else list(lags)
        if "box_cox"         in params: model.box_cox = params["box_cox"]
        if "box_cox_lmda"    in params: model.lamda   = params["box_cox_lmda"]
        if "box_cox_biasadj" in params: model.biasadj = params["box_cox_biasadj"]
 
    def _fit_params(params: dict) -> Optional[dict]:
        if isinstance(model.model, LinearRegression):
            return None
        return {k: v for k, v in params.items() if k not in _skip}
 
    def objective(trial: optuna.Trial) -> float:
        params = {name: suggest_fn(trial) for name, suggest_fn in param_space.items()}
 
        _set_model_params(params)
        fit_params = _fit_params(params)
 
        scores = []
        for train_idx, test_idx in tscv.split(df):
            train, test = df.iloc[train_idx], df.iloc[test_idx]
            x_test = test.drop(columns=[model.target_col])
            y_test = np.array(test[model.target_col])
 
            if fit_params is not None:
                model.model.set_params(**fit_params)
            model.fit(train)
            y_pred = model.forecast(len(y_test), x_test)
 
            if eval_metric.__name__ in ("MASE", "SMAE", "SRMSE", "RMSSE"):
                score = eval_metric(y_test, y_pred, train[model.target_col])
            else:
                score = eval_metric(y_test, y_pred)
            scores.append(score)
 
        mean_score = float(np.mean(scores))
        if verbose:
            print(f"Trial {trial.number:>4d} | score={mean_score:.6f} | {params} ")
        return mean_score
 
    study = optuna.create_study(direction="minimize")
    study.optimize(objective, n_trials=eval_num)
 
    best = study.best_params
 
    if "lags" in param_space:
        best_lags = best.pop("lags")
        if not isinstance(best_lags, list):
            best_lags = list(best_lags)
        model_parameters = best
    else:
        best_lags = None
        model_parameters = best
 
    return model_parameters, best_lags

# %% ../nbs/modules/03_model_selection.ipynb #855c8f86
#------------------------------------------------------------------------------
# Feature Selection Algorithms
# ------------------------------------------------------------------------------
 
def forward_feature_selection(
    model: object,
    df: pd.DataFrame,
    cv_split: int,
    H: int,
    step_size: Optional[int] = None,
    metrics: Union[Callable, List[Callable]] = None,
    lags_to_consider: Optional[int] = None,
    candidate_features: Optional[List[str]] = None,
    transformations: Optional[List] = None,
    starting_lags: Optional[List[int]] = None,
    starting_transforms: Optional[List] = None,
    best_start_score: Optional[List[float]] = None,
    verbose=False,
):
    """
    Forward stepwise feature selection for `ml_forecaster` models.
 
    At each iteration every remaining candidate (lag, exogenous column, or
    lag-transform) is tested individually by adding it to the current best
    feature set.  The candidate that produces the largest cross-validation
    improvement is permanently added.  The loop continues until no remaining
    candidate improves any of the evaluation metrics.
 
    Parameters
    ----------
    model : object
        A *configured but unfitted* `ml_forecaster` instance.  The function works exclusively on deep copies and never mutates the object passed in.
    df : pd.DataFrame
        Full training DataFrame. Must contain the target column and any candidate exogenous columns.
    cv_split : int
        Number of time-series cross-validation folds.
    H : int
        Forecast horizon (test window size for each fold).
    step_size : int, optional
        Step size between consecutive CV folds.  If `None` (default) the step equals `H`, producing non-overlapping folds — consistent with the default behaviour of `ml_forecaster.cross_validate`.
    metrics : callable or list of callable.
        One or more metric functions accepted by `ml_forecaster.cross_validate` (e.g. `[MAE, RMSE]`). Selection is driven by the **first** metric in the list; a candidate is only accepted when it improves **all** metrics simultaneously.
    lags_to_consider : int, optional
        Consider lags `1, 2, ..., lags_to_consider` as candidates.  If `None`, lag selection is skipped.
    candidate_features : list of str, optional
        Column names in `df` that are exogenous feature candidates.  The function never modifies this list.  If `None`, exogenous feature selection is skipped.
    transformations : list, optional
        Lag-transform objects to test as candidates (e.g. `[rolling_mean(3, 1), expanding_std(1)]`).  The function never modifies this list.  If `None`, transform selection is skipped.
    starting_lags : list of int, optional
        Lags to include in the initial feature set before the search begins. These are *not* candidates — they are always included.  Must be a list (e.g. `[1]` or `[1, 2, 3]`).
    starting_transforms : list, optional
        Lag-transform objects to include in the initial feature set before the search begins.  Must be a list.
    best_start_score : list of float, optional
        Initial best scores for each metric. If not provided, the function will compute the baseline score using the model with the starting features (if any) before beginning the search.
    verbose : bool, default False
        Print a message each time a candidate is accepted.
 
    Returns
    -------
    dict
        A dictionary with keys `best_lags`, `best_exogs`, and `best_transforms` containing the selected features.
    """
 
    # ── Normalise metrics to always be a list ─────────────────────────────────
    # cross_validate always expects a list; wrap a lone callable for the user.
    if metrics is None:
        raise ValueError("metrics must be provided (a callable or list of callables).")
    if callable(metrics):
        metrics = [metrics]
 
    # ── step_size default: non-overlapping folds (step = H) ───────────────────
    _step_size = step_size if step_size is not None else H
 
    # ── Validate starting_* arguments ────────────────────────────────────────
    if starting_lags is not None and not isinstance(starting_lags, list):
        raise ValueError("starting_lags must be a list of integers, e.g. [1] or [1, 2, 3].")
    if starting_transforms is not None and not isinstance(starting_transforms, list):
        raise ValueError("starting_transforms must be a list of transformation instances.")
 
    # ── Build candidate pools (local copies — never mutate caller's lists) ────
    remaining_lags = (
        [x for x in range(1, lags_to_consider + 1)
         if x not in (starting_lags or [])]
        if lags_to_consider is not None else []
    )
    remaining_feats = list(candidate_features) if candidate_features is not None else []
    remaining_transforms = list(transformations) if transformations is not None else []
 
    # ── Working df — start without exog candidates, add back as selected ──────
    if candidate_features is not None:
        df_work = df.drop(columns=candidate_features)
        df_orig = df.copy()
    else:
        df_work = df.copy()
        df_orig = df.copy()
 
    # ── Best-feature state ─────────────────────────────────────────────────────
    best_features = {
        "best_lags":       list(starting_lags)      if starting_lags      is not None else [],
        "best_exogs":      [],
        "best_transforms": list(starting_transforms) if starting_transforms is not None else [],
    }
 
    # Current best lags / transforms — kept in sync as candidates are accepted
    current_lags       = list(starting_lags)      if starting_lags      is not None else []
    current_transforms = list(starting_transforms) if starting_transforms is not None else []
 
    best_score = [float('inf')] * len(metrics)
 
    # ── Inner helpers ─────────────────────────────────────────────────────────
 
    def _scores_improved(new_score, ref_score):
        """True only when every metric strictly improves."""
        return all(n < r for n, r in zip(new_score, ref_score))
 
    def _validate(m, df_test):
        """
        Fit-and-score one candidate model via cross-validation.
        Returns a list of mean scores, one per metric.
        """
        result_df, _ = m.cross_validate(
            df=df_test,
            cv_split=cv_split,
            test_size=H,
            metrics=metrics,
            step_size=_step_size,
        )
        # result_df has columns ['eval_metric', 'score']
        return result_df["score"].tolist()
 
    # Columns in candidate_features that are also cat_variables on the model.
    # These must be hidden from the model until they are explicitly selected
    # as exog candidates — otherwise fit() crashes trying to build cat_var
    # for columns that have been dropped from df_work.
    _cat_exog_candidates = set()
    if candidate_features is not None and model.cat_variables is not None:
        _cat_exog_candidates = set(candidate_features) & set(model.cat_variables)
 
    def _make_candidate_model(lags, transforms, active_exogs=None):
        """
        Return a deep copy of the base model configured with the given lags
        and transforms.  active_exogs is the list of exogenous columns
        currently present in the df being passed to cross_validate; any
        cat_variables that are still pending as candidates are stripped from
        the copy so fit() does not look for columns missing from df_work.
        """
        m = model.copy()
        m.n_lag         = sorted(lags)     if lags       else None
        m.lag_transform = list(transforms) if transforms else None
 
        # Remove cat_variables that are still candidate columns not yet added
        if m.cat_variables is not None and _cat_exog_candidates:
            active = set(active_exogs) if active_exogs else set()
            present_cats = [
                c for c in m.cat_variables
                if c not in _cat_exog_candidates or c in active]
            m.cat_variables = present_cats if present_cats else None
 
        return m
 
    # ── Baseline score with starting features (if any) ───────────────────────
    if starting_lags is not None or starting_transforms is not None:
        m_start = _make_candidate_model(current_lags, current_transforms, best_features["best_exogs"])
        best_score = _validate(m_start, df_work)
        if verbose:
            print(f"Baseline score with starting features: {best_score}")
 
    # ── Stepwise forward loop ─────────────────────────────────────────────────
    while True:
        improvement    = False
        best_candidate = {'type': None, 'value': None}

        if best_start_score is not None:
            running_score = best_start_score
        else:
            running_score  = best_score  # tracks the best score seen this iteration
 
        # ── Test lag candidates ───────────────────────────────────────────────
        if remaining_lags:
            for lag in remaining_lags:
                m = _make_candidate_model(current_lags + [lag], current_transforms, best_features["best_exogs"])
                score = _validate(m, df_work)
                if _scores_improved(score, running_score):
                    running_score  = score
                    best_candidate = {'type': 'lag', 'value': lag}
                    improvement    = True
 
        # ── Test exogenous feature candidates ─────────────────────────────────
        if remaining_feats:
            for feat in remaining_feats:
                df_test = df_work.copy()
                df_test[feat] = df_orig[feat]
                m = _make_candidate_model(current_lags, current_transforms, best_features["best_exogs"] + [feat])
                score = _validate(m, df_test)
                if _scores_improved(score, running_score):
                    running_score  = score
                    best_candidate = {'type': 'exog', 'value': feat}
                    improvement    = True
 
        # ── Test lag-transform candidates ─────────────────────────────────────
        if remaining_transforms:
            for trans in remaining_transforms:
                m = _make_candidate_model(current_lags, current_transforms + [trans], best_features["best_exogs"])
                score = _validate(m, df_work)
                if _scores_improved(score, running_score):
                    running_score  = score
                    best_candidate = {'type': 'transform', 'value': trans}
                    improvement    = True
 
        # ── Accept the best candidate found this round ────────────────────────
        if improvement:
            best_score = running_score
            ctype = best_candidate['type']
            cval  = best_candidate['value']
 
            if ctype == 'lag':
                current_lags.append(cval)
                current_lags.sort()
                best_features["best_lags"].append(cval)
                best_features["best_lags"].sort()
                remaining_lags.remove(cval)
 
            elif ctype == 'exog':
                best_features["best_exogs"].append(cval)
                remaining_feats.remove(cval)
                df_work[cval] = df_orig[cval]
 
            elif ctype == 'transform':
                current_transforms.append(cval)
                best_features["best_transforms"].append(cval)
                remaining_transforms.remove(cval)
 
            if verbose:
                label = cval.get_name() if ctype == 'transform' else cval
                print(f"Added {ctype}: {label} | score: {best_score}")
 
        else:
            break  # No candidate improved any metric — search is complete
 
    # ── Finalise output ───────────────────────────────────────────────────────
    # Convert transform objects to their string names for the return value
    best_features["best_transforms"] = [
        t.get_name() for t in best_features["best_transforms"]
    ]
 
    return best_features

# %% ../nbs/modules/03_model_selection.ipynb #62f1fb3f
def backward_feature_selection(
    model: object,
    df: pd.DataFrame,
    cv_split: int,
    H: int,
    step_size: Optional[int] = None,
    metrics: Union[Callable, List[Callable]] = None,
    lags_to_consider: Optional[List[int]] = None,
    candidate_features: Optional[List[str]] = None,
    transformations: Optional[List] = None,
    verbose=False,
):
    """
    Backward stepwise feature selection for `ml_forecaster` models.

    Starts with the full feature set (all provided lags, exogenous columns, and
    lag-transforms) and at each iteration tries removing each current feature
    individually.  The feature whose removal produces the largest cross-validation
    improvement is permanently dropped.  The loop continues until no remaining
    feature can be removed without hurting any of the evaluation metrics.

    Parameters
    ----------
    model : ml_forecaster
        A *configured but unfitted* `ml_forecaster` instance.  The function works exclusively on deep copies and never mutates the object passed in.
    df : pd.DataFrame
        Full training DataFrame. Must contain the target column and any candidate exogenous columns.
    cv_split : int
        Number of time-series cross-validation folds.
    H : int
        Forecast horizon (test window size for each fold).
    step_size : int, optional
        Step size between consecutive CV folds.  If `None` (default) the step equals `H`, producing non-overlapping folds.
    metrics : callable or list of callable
        One or more metric functions accepted by `ml_forecaster.cross_validate` (e.g. `[MAE, RMSE]`). Selection is driven by the **first** metric in the list; a feature is only removed when doing so improves **all** metrics simultaneously.
    lags_to_consider : list of int, optional
        Lags to include in the initial feature set and test for removal (e.g. `[1, 2, 3, 4]`).  If `None`, no lag removal is attempted.
    candidate_features : list of str, optional
        Column names in `df` that start in the model and are tested for removal.  If `None`, exogenous feature removal is skipped.
    transformations : list, optional
        Lag-transform objects that start in the model and are tested for removal (e.g. `[rolling_mean(3, 1), expanding_std(1)]`).  If `None`, transform removal is skipped.
    verbose : bool, default False
        Print a message each time a feature is removed.

    Returns
    -------
    dict
        A dictionary with keys `best_lags`, `best_exogs`, and `best_transforms` containing the surviving features after backward selection.
    """

    # ── Normalise metrics to always be a list ─────────────────────────────────
    if metrics is None:
        raise ValueError("metrics must be provided (a callable or list of callables).")
    if callable(metrics):
        metrics = [metrics]

    # ── step_size default: non-overlapping folds (step = H) ───────────────────
    _step_size = step_size if step_size is not None else H

    # ── Build initial feature sets (local copies — never mutate caller's lists) ─
    if lags_to_consider is None:
        current_lags = []
    elif isinstance(lags_to_consider, int):
        current_lags = list(range(1, lags_to_consider + 1))
    else:
        current_lags = sorted(lags_to_consider)
    current_exogs      = list(candidate_features)  if candidate_features is not None else []
    current_transforms = list(transformations)      if transformations    is not None else []


    # ── Working df includes all exog candidates from the start ────────────────
    df_work = df.copy()

    # ── Columns in candidate_features that are also cat_variables on the model ─
    _cat_exog_candidates = set()
    if candidate_features is not None and model.cat_variables is not None:
        _cat_exog_candidates = set(candidate_features) & set(model.cat_variables)

    # ── Inner helpers ─────────────────────────────────────────────────────────

    def _scores_improved(new_score, ref_score):
        """True only when every metric strictly improves."""
        return all(n < r for n, r in zip(new_score, ref_score))

    def _make_candidate_model(lags, transforms, active_exogs=None):
        """
        Return a deep copy of the base model configured with the given lags and
        transforms.  Cat_variables that are no longer in active_exogs are stripped
        from the copy so fit() does not look for missing columns.
        """
        m = model.copy()
        m.n_lag         = sorted(lags)     if lags       else None
        m.lag_transform = list(transforms) if transforms else None

        if m.cat_variables is not None and _cat_exog_candidates:
            active = set(active_exogs) if active_exogs else set()
            present_cats = [
                c for c in m.cat_variables
                if c not in _cat_exog_candidates or c in active]
            m.cat_variables = present_cats if present_cats else None

        return m

    def _validate(m, df_test):
        """Fit-and-score one candidate model via cross-validation."""
        result_df, _ = m.cross_validate(
            df=df_test,
            cv_split=cv_split,
            test_size=H,
            metrics=metrics,
            step_size=_step_size,
        )
        return result_df["score"].tolist()

    # ── Baseline score with all features ──────────────────────────────────────
    m_start = _make_candidate_model(current_lags, current_transforms, current_exogs)
    best_score = _validate(m_start, df_work)
    if verbose:
        print(f"Baseline score with all features: {best_score}")

    # ── Stepwise backward loop ────────────────────────────────────────────────
    while True:
        improvement    = False
        best_candidate = {'type': None, 'value': None}
        running_score  = best_score

        # ── Test lag removals ─────────────────────────────────────────────────
        if current_lags:
            for lag in current_lags:
                reduced_lags = [l for l in current_lags if l != lag]
                m = _make_candidate_model(reduced_lags, current_transforms, current_exogs)
                score = _validate(m, df_work)
                if _scores_improved(score, running_score):
                    running_score  = score
                    best_candidate = {'type': 'lag', 'value': lag}
                    improvement    = True

        # ── Test exogenous feature removals ───────────────────────────────────
        if current_exogs:
            for feat in current_exogs:
                reduced_exogs = [f for f in current_exogs if f != feat]
                df_test = df_work.drop(columns=[feat])
                m = _make_candidate_model(current_lags, current_transforms, reduced_exogs)
                score = _validate(m, df_test)
                if _scores_improved(score, running_score):
                    running_score  = score
                    best_candidate = {'type': 'exog', 'value': feat}
                    improvement    = True

        # ── Test lag-transform removals ───────────────────────────────────────
        if current_transforms:
            for trans in current_transforms:
                reduced_transforms = [t for t in current_transforms if t is not trans]
                m = _make_candidate_model(current_lags, reduced_transforms, current_exogs)
                score = _validate(m, df_work)
                if _scores_improved(score, running_score):
                    running_score  = score
                    best_candidate = {'type': 'transform', 'value': trans}
                    improvement    = True

        # ── Accept the best removal found this round ──────────────────────────
        if improvement:
            best_score = running_score
            ctype = best_candidate['type']
            cval  = best_candidate['value']

            if ctype == 'lag':
                current_lags.remove(cval)
                if verbose:
                    print(f"Removed lag: {cval} | score: {best_score}")

            elif ctype == 'exog':
                current_exogs.remove(cval)
                df_work = df_work.drop(columns=[cval])
                if verbose:
                    print(f"Removed exog: {cval} | score: {best_score}")

            elif ctype == 'transform':
                current_transforms = [t for t in current_transforms if t is not cval]
                if verbose:
                    label = cval.get_name()
                    print(f"Removed transform: {label} | score: {best_score}")

        else:
            break  # No removal improved any metric — search is complete

    # ── Finalise output ───────────────────────────────────────────────────────
    return {
        "best_lags":       current_lags,
        "best_exogs":      current_exogs,
        "best_transforms": [t.get_name() for t in current_transforms],
    }


# %% ../nbs/modules/03_model_selection.ipynb #46039723
def mv_forward_feature_selection(
    model: object,
    df: pd.DataFrame,
    target_col: str,
    cv_split: int,
    H: int,
    step_size=None,
    metrics=None,
    lags_to_consider=None,
    candidate_features=None,
    transformations=None,
    starting_lags=None,
    starting_transforms=None,
    verbose=False,
):
    """
    Forward stepwise feature selection for `ml_mv_forecaster`.

    Parameters
    ----------
    model : ml_mv_forecaster
        Template model — never mutated.
    df : pd.DataFrame
        DataFrame containing the target variable and any candidate features.
    target_col : str
        Target variable used to evaluate cross-validation score.
    cv_split : int
        Number of time-series cross-validation folds.
    H : int
        Forecast horizon / test size per fold.
    step_size : int, optional
        Rolling-window step size (defaults to H).
    metrics : callable or list of callable
        One or more metric functions (e.g. `[MAE, RMSE]`). Selection is driven by the **first** metric in the list; a candidate is only accepted when it improves **all** metrics simultaneously.
    lags_to_consider : dict, optional
        ``{col: max_lag}`` — lags 1..max_lag are candidates.
    candidate_features : list of str, optional
        Exogenous columns to consider adding.
    transformations : dict, optional
        ``{col: [transform_objects]}`` — transform candidates per target.
    starting_lags : dict, optional
        Lags already included before search begins.
    starting_transforms : dict, optional
        Transforms already included before search begins.
    verbose : bool, default False

    Returns
    -------
    dict
        `{"best_lags": {col: [...]}, "best_exogs": [...],
           "best_transforms": {col: [name_str, ...]}}`
    """
    if metrics is None:
        raise ValueError("metrics must be provided.")
    if callable(metrics):
        metrics = [metrics]

    _step_size = step_size if step_size is not None else H

    if lags_to_consider is None:
        lags_to_consider = {}
    if transformations is None:
        transformations = {}

    remaining_lags = {col: list(range(1, ml + 1)) for col, ml in lags_to_consider.items()}
    remaining_transforms = {col: list(tl) for col, tl in transformations.items()}
    remaining_feats = list(candidate_features) if candidate_features is not None else []

    best_lags = {col: [] for col in lags_to_consider}
    if starting_lags is not None:
        for col, lags in starting_lags.items():
            best_lags[col] = list(lags)
            remaining_lags[col] = [x for x in remaining_lags.get(col, []) if x not in lags]

    best_transforms = {col: [] for col in transformations}
    if starting_transforms is not None:
        for col, tl in starting_transforms.items():
            best_transforms[col] = list(tl)
            remaining_transforms[col] = [t for t in remaining_transforms.get(col, []) if t not in tl]

    best_features = {
        "best_lags":       best_lags,
        "best_exogs":      [],
        "best_transforms": best_transforms,
    }

    if candidate_features is not None:
        df_work = df.drop(columns=candidate_features)
        df_orig = df.copy()
    else:
        df_work = df.copy()
        df_orig = df.copy()

    best_score = [float('inf')] * len(metrics)

    _cat_exog_candidates = set()
    if candidate_features is not None and model.cat_variables is not None:
        _cat_exog_candidates = set(candidate_features) & set(model.cat_variables)

    def _scores_improved(new_score, ref_score):
        return all(n < r for n, r in zip(new_score, ref_score))

    def _validate(m, df_test):
        # cross_validate returns (metrics_df, cv_df_); metrics_df has columns
        # ["eval_metric", "score"] — one row per metric.
        result_df, _ = m.cross_validate(
            df=df_test, target_col=target_col, cv_split=cv_split,
            test_size=H, metrics=metrics, step_size=_step_size,
        )
        return result_df['score'].tolist()

    def _make_candidate_model(lags_dict, transforms_dict, active_exogs=None):
        m = model.copy()
        # Use {} (not None) when no lags — ml_mv_forecaster.data_prep iterates m.n_lag
        m.n_lag = {col: sorted(v) for col, v in lags_dict.items() if v}
        at = {col: list(v) for col, v in transforms_dict.items() if v}
        m.lag_transform = at if at else None
        if m.cat_variables is not None and _cat_exog_candidates:
            active = set(active_exogs) if active_exogs else set()
            present = [c for c in m.cat_variables
                       if c not in _cat_exog_candidates or c in active]
            m.cat_variables = present if present else None
        return m

    # Evaluate starting point if warm-start provided
    if starting_lags is not None or starting_transforms is not None:
        m_start = _make_candidate_model(
            best_features["best_lags"],
            best_features["best_transforms"],
            best_features["best_exogs"],
        )
        best_score = _validate(m_start, df_work)
        if verbose:
            print(f"Baseline score: {best_score}")

    while True:
        improvement    = False
        best_candidate = {'target': None, 'type': None, 'value': None}
        running_score  = best_score

        # --- try adding each candidate lag ---
        for col, lags in remaining_lags.items():
            for lg in lags:
                trial_lags = {c: list(v) for c, v in best_features["best_lags"].items()}
                trial_lags[col] = sorted(trial_lags[col] + [lg])
                m = _make_candidate_model(trial_lags, best_features["best_transforms"],
                                          best_features["best_exogs"])
                score = _validate(m, df_work)
                if _scores_improved(score, running_score):
                    running_score  = score
                    best_candidate = {'target': col, 'type': 'lag', 'value': lg}
                    improvement    = True

        # --- try adding each candidate exog feature ---
        for feat in remaining_feats:
            df_test = df_work.copy()
            df_test[feat] = df_orig[feat]
            active_exogs = best_features["best_exogs"] + [feat]
            m = _make_candidate_model(best_features["best_lags"],
                                      best_features["best_transforms"], active_exogs)
            score = _validate(m, df_test)
            if _scores_improved(score, running_score):
                running_score  = score
                best_candidate = {'target': None, 'type': 'exog', 'value': feat}
                improvement    = True

        # --- try adding each candidate transform ---
        for col, tl in remaining_transforms.items():
            for trans in tl:
                trial_trans = {c: list(v) for c, v in best_features["best_transforms"].items()}
                trial_trans[col] = trial_trans[col] + [trans]
                m = _make_candidate_model(best_features["best_lags"], trial_trans,
                                          best_features["best_exogs"])
                score = _validate(m, df_work)
                if _scores_improved(score, running_score):
                    running_score  = score
                    best_candidate = {'target': col, 'type': 'transform', 'value': trans}
                    improvement    = True

        if improvement:
            best_score = running_score
            ctype  = best_candidate['type']
            cval   = best_candidate['value']
            ctargt = best_candidate['target']
            if ctype == 'lag':
                best_features["best_lags"][ctargt].append(cval)
                best_features["best_lags"][ctargt].sort()
                remaining_lags[ctargt].remove(cval)
            elif ctype == 'exog':
                best_features["best_exogs"].append(cval)
                remaining_feats.remove(cval)
                df_work[cval] = df_orig[cval]
            elif ctype == 'transform':
                best_features["best_transforms"][ctargt].append(cval)
                remaining_transforms[ctargt].remove(cval)
            if verbose:
                label = cval.get_name() if ctype == 'transform' else cval
                tgt_str = f" [{ctargt}]" if ctargt else ""
                print(f"Added {ctype}{tgt_str}: {label} | score: {best_score}")
        else:
            break

    for col in best_features["best_lags"]:
        best_features["best_lags"][col].sort()
    best_features["best_transforms"] = {
        col: [t.get_name() for t in tl]
        for col, tl in best_features["best_transforms"].items()
    }
    return best_features

# %% ../nbs/modules/03_model_selection.ipynb #8f3f6822
def mv_backward_feature_selection(
    model: object,
    df: pd.DataFrame,
    target_col: str,
    cv_split: int,
    H: int,
    step_size=None,
    metrics=None,
    lags_to_consider=None,
    candidate_features=None,
    transformations=None,
    verbose=False,
):
    """
    Backward stepwise feature selection for ``ml_mv_forecaster``.

    Starts with all candidate features included and iteratively removes
    the one whose removal most improves cross-validation score.

    Parameters
    ----------
    model : ml_mv_forecaster
        Template model — never mutated.
    df : pd.DataFrame
        All candidate exog columns must already be present.
    target_col : str
        Target variable used to evaluate cross-validation score.
    cv_split : int
    H : int
        Forecast horizon / test size per fold.
    step_size : int, optional
        Rolling-window step size (defaults to H).
    metrics : callable or list of callable
        One or more metric functions (e.g. `[MAE, RMSE]`). A feature is only removed when its removal improves **all** metrics simultaneously.
    lags_to_consider : dict, optional
        ``{col: max_lag}`` — all lags 1..max_lag start as selected.
    candidate_features : list of str, optional
        Exogenous columns that start as selected.
    transformations : dict, optional
        ``{col: [transform_objects]}`` — all transforms start as selected.
    verbose : bool, default False

    Returns
    -------
    dict
        ``{"best_lags": {col: [...]}, "best_exogs": [...],
           "best_transforms": {col: [name_str, ...]}}``
    """
    if metrics is None:
        raise ValueError("metrics must be provided.")
    if callable(metrics):
        metrics = [metrics]

    _step_size = step_size if step_size is not None else H

    best_features = {
        "best_lags": {col: list(range(1, ml + 1))
                      for col, ml in (lags_to_consider or {}).items()},
        "best_exogs": list(candidate_features) if candidate_features is not None else [],
        "best_transforms": {col: list(tl) for col, tl in (transformations or {}).items()},
    }

    df_work = df.copy()

    _cat_exog_candidates = set()
    if candidate_features is not None and model.cat_variables is not None:
        _cat_exog_candidates = set(candidate_features) & set(model.cat_variables)

    def _scores_improved(new_score, ref_score):
        return all(n < r for n, r in zip(new_score, ref_score))

    def _validate(m, df_test):
        result_df, _ = m.cross_validate(
            df=df_test, target_col=target_col, cv_split=cv_split,
            test_size=H, metrics=metrics, step_size=_step_size,
        )
        return result_df['score'].tolist()

    def _make_candidate_model(lags_dict, transforms_dict, active_exogs=None):
        m = model.copy()
        m.n_lag = {col: sorted(v) for col, v in lags_dict.items() if v}
        at = {col: list(v) for col, v in transforms_dict.items() if v}
        m.lag_transform = at if at else None
        if m.cat_variables is not None and _cat_exog_candidates:
            active = set(active_exogs) if active_exogs else set()
            present = [c for c in m.cat_variables
                       if c not in _cat_exog_candidates or c in active]
            m.cat_variables = present if present else None
        return m

    # Evaluate baseline (all features included) before starting removal
    m_baseline = _make_candidate_model(
        best_features["best_lags"],
        best_features["best_transforms"],
        best_features["best_exogs"],
    )
    best_score = _validate(m_baseline, df_work)
    if verbose:
        print(f"Baseline score (all features): {best_score}")

    while True:
        improvement    = False
        best_candidate = {'target': None, 'type': None, 'value': None}
        running_score  = best_score

        # --- try removing each selected lag ---
        for targ_col, lags in best_features["best_lags"].items():
            for lg in lags:
                trial_lags = {col: list(v) for col, v in best_features["best_lags"].items()}
                trial_lags[targ_col] = sorted([x for x in lags if x != lg])
                m = _make_candidate_model(trial_lags, best_features["best_transforms"],
                                          best_features["best_exogs"])
                score = _validate(m, df_work)
                if _scores_improved(score, running_score):
                    running_score  = score
                    best_candidate = {'target': targ_col, 'type': 'lag', 'value': lg}
                    improvement    = True

        # --- try removing each selected transform ---
        for targ_col, tl in best_features["best_transforms"].items():
            for tr in tl:
                trial_trans = {col: list(v) for col, v in best_features["best_transforms"].items()}
                trial_trans[targ_col] = [x for x in tl if x is not tr]
                m = _make_candidate_model(best_features["best_lags"], trial_trans,
                                          best_features["best_exogs"])
                score = _validate(m, df_work)
                if _scores_improved(score, running_score):
                    running_score  = score
                    best_candidate = {'target': targ_col, 'type': 'transform', 'value': tr}
                    improvement    = True

        # --- try removing each selected exog feature ---
        for feat in best_features["best_exogs"]:
            df_test      = df_work.drop(columns=[feat])
            active_after = [x for x in best_features["best_exogs"] if x != feat]
            m = _make_candidate_model(best_features["best_lags"],
                                      best_features["best_transforms"], active_after)
            score = _validate(m, df_test)
            if _scores_improved(score, running_score):
                running_score  = score
                best_candidate = {'target': None, 'type': 'exog', 'value': feat}
                improvement    = True

        if improvement:
            best_score = running_score
            ctype  = best_candidate['type']
            cval   = best_candidate['value']
            ctargt = best_candidate['target']
            if ctype == 'lag':
                best_features["best_lags"][ctargt].remove(cval)
            elif ctype == 'exog':
                best_features["best_exogs"].remove(cval)
                df_work = df_work.drop(columns=[cval])
            elif ctype == 'transform':
                best_features["best_transforms"][ctargt].remove(cval)
            if verbose:
                label = cval.get_name() if ctype == 'transform' else cval
                tgt_str = f" [{ctargt}]" if ctargt else ""
                print(f"Removed {ctype}{tgt_str}: {label} | score: {best_score}")
        else:
            break

    for col in best_features["best_lags"]:
        best_features["best_lags"][col].sort()
    best_features["best_transforms"] = {
        col: [t.get_name() for t in tl]
        for col, tl in best_features["best_transforms"].items()
    }
    return best_features

# %% ../nbs/modules/03_model_selection.ipynb #8364f2e7
def ms_arr_forward_feature_selection(
    model: object,
    df: pd.DataFrame,
    cv_split: int,
    H: int,
    step_size: Optional[int] = None,
    metrics: Union[Callable, List[Callable]] = None,
    lags_to_consider: Optional[Union[int, List[int]]] = None,
    candidate_features: Optional[List[str]] = None,
    transformations: Optional[List] = None,
    starting_lags: Optional[List[int]] = None,
    starting_transforms: Optional[List] = None,
    validation_type: str = "cv",
    iterations: int = 10,
    verbose: bool = False,
):
    """
    Forward stepwise feature selection for `ms_arr` models.

    At each iteration every remaining candidate (lag, exogenous column, or
    lag-transform) is tested individually by adding it to the current best
    feature set.  The candidate that produces the largest improvement is
    permanently added.  The loop continues until no remaining candidate
    improves the evaluation criterion.

    The HMM state (A, pi, stds, coeffs) is warm-started from the round
    winner and propagated to subsequent rounds for consistent initialisation.

    Parameters
    ----------
    model : ms_arr
        A configured `ms_arr` instance with `fit_em()` already called (recommended to use few EM iterations for this initial fit, e.g. `iterations=10`) or a template model with the same configuration but not yet fitted.  The model is copied internally and never mutated, so the caller's instance remains unchanged.
    df : pd.DataFrame
        Full training DataFrame. Must contain the target column and any candidate exogenous columns.
    cv_split : int
        Number of time-series cross-validation folds.
    H : int
        Forecast horizon (test window size for each fold).
    step_size : int, optional
        Step size between consecutive CV folds. Defaults to H.
    metrics : callable or list of callable
        Required when validation_type='cv'. Selection driven by first metric; a candidate is accepted only when it improves all metrics.
    lags_to_consider : int or list of int, optional
        Candidate lags. Int → 1..n; list → specific lags.
    candidate_features : list of str, optional
        Exogenous column names to test as candidates.
    transformations : list, optional
        Lag-transform objects to test as candidates.
    starting_lags : list of int, optional
        Lags always included in the initial set (not candidates).
    starting_transforms : list, optional
        Transforms always included in the initial set (not candidates).
    validation_type : str, default 'cv'
        Criterion for selection: 'cv', 'AIC', 'BIC', or 'AIC_BIC'. When 'cv', metrics must be provided and drive selection. When 'AIC' or 'BIC', the respective information criterion is used. When 'AIC_BIC', a candidate is accepted only if it improves both AIC and BIC.
    iterations : int, default 100
        EM iterations used inside fit_em() for each candidate evaluation.
    verbose : bool, default False
        Print a message each time a candidate is accepted.

    Returns
    -------
    dict
         `{"best_lags": [...], "best_exogs": [...], "best_transforms": [...]}`

    """

    if metrics is not None and callable(metrics):
        metrics = [metrics]
    if validation_type == "cv" and metrics is None:
        raise ValueError("metrics must be provided when validation_type='cv'.")
    if starting_lags is not None and not isinstance(starting_lags, list):
        raise ValueError("starting_lags must be a list of integers.")
    if starting_transforms is not None and not isinstance(starting_transforms, list):
        raise ValueError("starting_transforms must be a list of transform instances.")

    _step_size = step_size if step_size is not None else H

    # ── Local working copy — caller's model NEVER touched ─────────────────────
    local_model = model.copy()

    if not local_model.is_fitted:
        local_model.fit(df)
    
    # local_model.iter = iterations

    # ── Candidate pools ───────────────────────────────────────────────────────
    if lags_to_consider is not None:
        _all = (list(range(1, lags_to_consider + 1))
                if isinstance(lags_to_consider, int) else list(lags_to_consider))
        remaining_lags = [x for x in _all if x not in (starting_lags or [])]
    else:
        remaining_lags = []

    remaining_feats      = list(candidate_features) if candidate_features is not None else []
    remaining_transforms = list(transformations)    if transformations    is not None else []

    # ── Apply starting features to local_model ────────────────────────────────
    local_model.lags          = list(starting_lags)       if starting_lags       is not None else None
    local_model.lag_transform = list(starting_transforms) if starting_transforms is not None else None

    # ── Working df ────────────────────────────────────────────────────────────
    if candidate_features is not None:
        df_work = df.drop(columns=candidate_features)
        df_orig = df.copy()
    else:
        df_work = df.copy()
        df_orig = df.copy()

    best_features = {
        "best_lags":       list(starting_lags)       if starting_lags       is not None else [],
        "best_exogs":      [],
        "best_transforms": list(starting_transforms) if starting_transforms is not None else [],
    }

    # ── Initial best score ────────────────────────────────────────────────────
    if validation_type == "cv":
        best_score = [float('inf')] * len(metrics)
    elif validation_type in ("AIC", "BIC"):
        best_score = float('inf')
    else:  # AIC_BIC
        best_score = [float('inf'), float('inf')]

    # ── Cat variable handling ─────────────────────────────────────────────────
    _cat_exog_candidates = set()
    if candidate_features is not None and model.cat_variables is not None:
        _cat_exog_candidates = set(candidate_features) & set(model.cat_variables)

    def _strip_pending_cats(m, active_exogs):
        if m.cat_variables is not None and _cat_exog_candidates:
            active = set(active_exogs) if active_exogs else set()
            present = [c for c in m.cat_variables
                       if c not in _cat_exog_candidates or c in active]
            m.cat_variables = present if present else None

    # ── Inner helpers ─────────────────────────────────────────────────────────

    def _scores_improved(new_score, ref_score):
        if isinstance(ref_score, list):
            return all(n < r for n, r in zip(new_score, ref_score))
        return new_score < ref_score

    def _train(m, df_test):
        """Full EM fit — resets coeffs so data_prep reinitialises for new feature set."""
        m.coeffs = None  # shape changes with each candidate — must reinitialise
        m.iter = iterations
        m.fit(df_test)
        return m

    def _validate(m, df_test):
        if validation_type == "cv":
            result, _ = m.cross_validate(
                df=df_test, cv_split=cv_split, test_size=H,
                metrics=metrics, step_size=_step_size, n_iter=iterations,
            )
            return result["score"].tolist()
        elif validation_type == "BIC":
            return m.bic
        elif validation_type == "AIC":
            return m.aic
        else:  # AIC_BIC
            return [m.aic, m.bic]

    def _propagate_hmm_state(src, dst):
        """Copy fitted HMM state from src → dst for warm-starting next round."""
        dst.A      = src.A.copy()
        dst.pi     = src.pi.copy()
        dst.stds   = src.stds.copy()
        # dst.coeffs = src.coeffs.copy()

    # ── Baseline score with starting features ────────────────────────────────
    if starting_lags is not None or starting_transforms is not None:
        m_start = local_model.copy()
        _strip_pending_cats(m_start, best_features["best_exogs"])
        _train(m_start, df_work)
        best_score = _validate(m_start, df_work)
        _propagate_hmm_state(m_start, local_model)
        if verbose:
            print(f"Baseline score: {best_score}")

    # ── Stepwise forward loop ─────────────────────────────────────────────────
    while True:
        improvement    = False
        best_candidate = {'type': None, 'value': None, 'model': None}
        running_score  = best_score

        # ── Lag candidates ────────────────────────────────────────────────────
        for lag in remaining_lags:
            m = local_model.copy()
            m.lags = sorted(best_features["best_lags"] + [lag])
            _strip_pending_cats(m, best_features["best_exogs"])
            _train(m, df_work)
            score = _validate(m, df_work)
            if _scores_improved(score, running_score):
                running_score  = score
                best_candidate = {'type': 'lag', 'value': lag, 'model': m}
                improvement    = True

        # ── Exogenous feature candidates ──────────────────────────────────────
        for feat in remaining_feats:
            df_test = df_work.copy()
            df_test[feat] = df_orig[feat]
            active_exogs = best_features["best_exogs"] + [feat]
            m = local_model.copy()
            _strip_pending_cats(m, active_exogs)
            _train(m, df_test)
            score = _validate(m, df_test)
            if _scores_improved(score, running_score):
                running_score  = score
                best_candidate = {'type': 'exog', 'value': feat, 'model': m}
                improvement    = True

        # ── Transform candidates ──────────────────────────────────────────────
        for trans in remaining_transforms:
            m = local_model.copy()
            m.lag_transform = list(best_features["best_transforms"]) + [trans]
            _strip_pending_cats(m, best_features["best_exogs"])
            _train(m, df_work)
            score = _validate(m, df_work)
            if _scores_improved(score, running_score):
                running_score  = score
                best_candidate = {'type': 'transform', 'value': trans, 'model': m}
                improvement    = True

        # ── Accept best candidate ─────────────────────────────────────────────
        if improvement:
            best_score = running_score
            ctype = best_candidate['type']
            cval  = best_candidate['value']
            m_won = best_candidate['model']

            if ctype == 'lag':
                best_features["best_lags"].append(cval)
                best_features["best_lags"].sort()
                remaining_lags.remove(cval)
                local_model.lags = list(best_features["best_lags"])

            elif ctype == 'exog':
                best_features["best_exogs"].append(cval)
                remaining_feats.remove(cval)
                df_work[cval] = df_orig[cval]

            elif ctype == 'transform':
                best_features["best_transforms"].append(cval)
                remaining_transforms.remove(cval)
                local_model.lag_transform = list(best_features["best_transforms"])

            _propagate_hmm_state(m_won, local_model)

            if verbose:
                label = cval.get_name() if ctype == 'transform' else cval
                print(f"Added {ctype}: {label} | score: {best_score}")
        else:
            break

    # ── Finalise — refit on full df_work with selected features ───────────────
    model_ = local_model.copy()
    model_.lags          = list(best_features["best_lags"])       if best_features["best_lags"]       else None
    model_.lag_transform = list(best_features["best_transforms"]) if best_features["best_transforms"] else None
    _strip_pending_cats(model_, best_features["best_exogs"])  # ← add this
    _train(model_, df_work)

    best_features["best_transforms"] = [t.get_name() for t in best_features["best_transforms"]]
    return best_features, model_

# %% ../nbs/modules/03_model_selection.ipynb #ec0db62f
def ms_arr_backward_feature_selection(
    df: pd.DataFrame,
    cv_split: int,
    H: int,
    step_size: Optional[int] = None,
    model: object = None,
    metrics: Union[Callable, List[Callable]] = None,
    lags_to_consider: Optional[Union[int, List[int]]] = None,
    candidate_features: Optional[List[str]] = None,
    transformations: Optional[List] = None,
    validation_type: str = "cv",
    iterations: int = 100,
    verbose: bool = False,
):
    """
    Backward stepwise feature selection for `ms_arr` models.

    Starts with the full feature set and at each iteration tries removing each
    current feature individually.  The feature whose removal produces the
    largest improvement is permanently dropped.  The loop continues until no
    removal improves the evaluation criterion.

    The HMM state (A, pi, stds, coeffs) is warm-started from the round
    winner and propagated to subsequent rounds.

    Parameters
    ----------
    df : pd.DataFrame
        Full training DataFrame. All candidate exogenous columns must be present.
    cv_split : int
        Number of time-series cross-validation folds.
    H : int
        Forecast horizon (test window size for each fold).
    step_size : int, optional
        Step size between consecutive CV folds. Defaults to H.
    model : ms_arr
        A configured but unfitted ms_arr instance. Never mutated.
    metrics : callable or list of callable
        Required when validation_type='cv'. A feature is only removed when
        doing so improves all metrics simultaneously.
    lags_to_consider : int or list of int, optional
        Initial lag set. Int → 1..n; list → specific lags.
    candidate_features : list of str, optional
        Exogenous columns that start in the model and are tested for removal.
    transformations : list, optional
        Lag-transform objects that start in the model and are tested for removal.
    validation_type : str, default 'cv'
        Criterion for selection: 'cv', 'AIC', 'BIC', or 'AIC_BIC'.
    iterations : int, default 100
        EM iterations used inside fit_em() for each candidate evaluation.
    verbose : bool, default False
        Print a message each time a feature is removed.

    Returns
    -------
    dict
        `{"best_lags": [...], "best_exogs": [...], "best_transforms": [...]}`
    """

    if metrics is not None and callable(metrics):
        metrics = [metrics]
    if validation_type == "cv" and metrics is None:
        raise ValueError("metrics must be provided when validation_type='cv'.")

    _step_size = step_size if step_size is not None else H

    # ── Local working copy ────────────────────────────────────────────────────
    local_model = model.copy()

    if not local_model.is_fitted:
        local_model.fit(df)

    # ── Initial feature sets ──────────────────────────────────────────────────
    if lags_to_consider is not None:
        current_lags = (list(range(1, lags_to_consider + 1))
                        if isinstance(lags_to_consider, int) else sorted(lags_to_consider))
    else:
        current_lags = []

    current_exogs      = list(candidate_features) if candidate_features is not None else []
    current_transforms = list(transformations)    if transformations    is not None else []

    local_model.lags          = list(current_lags)       if current_lags       else None
    local_model.lag_transform = list(current_transforms) if current_transforms else None

    df_work = df.copy()

    # ── Cat variable handling ─────────────────────────────────────────────────
    _cat_exog_candidates = set()
    if candidate_features is not None and model.cat_variables is not None:
        _cat_exog_candidates = set(candidate_features) & set(model.cat_variables)

    def _strip_pending_cats(m, active_exogs):
        if m.cat_variables is not None and _cat_exog_candidates:
            active = set(active_exogs) if active_exogs else set()
            present = [c for c in m.cat_variables
                       if c not in _cat_exog_candidates or c in active]
            m.cat_variables = present if present else None

    # ── Inner helpers ─────────────────────────────────────────────────────────

    def _scores_improved(new_score, ref_score):
        if isinstance(ref_score, list):
            return all(n < r for n, r in zip(new_score, ref_score))
        return new_score < ref_score

    # def _train(m, df_test):
    #     m.iter = iterations
    #     m.fit_em(df_test)
    #     return m
    
    def _train(m, df_test):
        """Full EM fit — resets coeffs so data_prep reinitialises for new feature set."""
        m.coeffs = None  # shape changes with each candidate — must reinitialise
        m.iter = iterations
        m.fit(df_test)
        return m

    def _validate(m, df_test):
        if validation_type == "cv":
            result, _ = m.cross_validate(
                df=df_test, cv_split=cv_split, test_size=H,
                metrics=metrics, step_size=_step_size, n_iter=iterations,
            )
            return result["score"].tolist()
        elif validation_type == "BIC":
            return m.bic
        elif validation_type == "AIC":
            return m.aic
        else:  # AIC_BIC
            return [m.aic, m.bic]

    def _propagate_hmm_state(src, dst):
        dst.A      = src.A.copy()
        dst.pi     = src.pi.copy()
        dst.stds   = src.stds.copy()
        dst.coeffs = src.coeffs.copy()

    # ── Baseline score with all features ─────────────────────────────────────
    m_start = local_model.copy()
    _strip_pending_cats(m_start, current_exogs)
    _train(m_start, df_work)
    best_score = _validate(m_start, df_work)
    _propagate_hmm_state(m_start, local_model)
    if verbose:
        print(f"Baseline score with all features: {best_score}")

    # ── Stepwise backward loop ────────────────────────────────────────────────
    while True:
        improvement    = False
        best_candidate = {'type': None, 'value': None, 'model': None}
        running_score  = best_score

        # ── Test lag removals ─────────────────────────────────────────────────
        for lag in current_lags:
            reduced_lags = [l for l in current_lags if l != lag]
            m = local_model.copy()
            m.lags = reduced_lags if reduced_lags else None
            _strip_pending_cats(m, current_exogs)
            _train(m, df_work)
            score = _validate(m, df_work)
            if _scores_improved(score, running_score):
                running_score  = score
                best_candidate = {'type': 'lag', 'value': lag, 'model': m}
                improvement    = True

        # ── Test exog removals ────────────────────────────────────────────────
        for feat in current_exogs:
            reduced_exogs = [f for f in current_exogs if f != feat]
            df_test = df_work.drop(columns=[feat])
            m = local_model.copy()
            _strip_pending_cats(m, reduced_exogs)
            _train(m, df_test)
            score = _validate(m, df_test)
            if _scores_improved(score, running_score):
                running_score  = score
                best_candidate = {'type': 'exog', 'value': feat, 'model': m}
                improvement    = True

        # ── Test transform removals ───────────────────────────────────────────
        for trans in current_transforms:
            reduced_transforms = [t for t in current_transforms if t is not trans]
            m = local_model.copy()
            m.lag_transform = reduced_transforms if reduced_transforms else None
            _strip_pending_cats(m, current_exogs)
            _train(m, df_work)
            score = _validate(m, df_work)
            if _scores_improved(score, running_score):
                running_score  = score
                best_candidate = {'type': 'transform', 'value': trans, 'model': m}
                improvement    = True

        # ── Accept best removal ───────────────────────────────────────────────
        if improvement:
            best_score = running_score
            ctype = best_candidate['type']
            cval  = best_candidate['value']
            m_won = best_candidate['model']

            if ctype == 'lag':
                current_lags.remove(cval)
                local_model.lags = current_lags if current_lags else None
                if verbose:
                    print(f"Removed lag: {cval} | score: {best_score}")

            elif ctype == 'exog':
                current_exogs.remove(cval)
                df_work = df_work.drop(columns=[cval])
                if verbose:
                    print(f"Removed exog: {cval} | score: {best_score}")

            elif ctype == 'transform':
                current_transforms = [t for t in current_transforms if t is not cval]
                local_model.lag_transform = current_transforms if current_transforms else None
                if verbose:
                    label = cval.get_name()
                    print(f"Removed transform: {label} | score: {best_score}")

            _propagate_hmm_state(m_won, local_model)

        else:
            break

    # ── Finalise ──────────────────────────────────────────────────────────────
    model_ = local_model.copy()
    model_.lags          = current_lags       if current_lags       else None
    model_.lag_transform = current_transforms if current_transforms else None
    _strip_pending_cats(model_, current_exogs)  # ← add this
    _train(model_, df_work)


    return {
        "best_lags":       current_lags,
        "best_exogs":      current_exogs,
        "best_transforms": [t.get_name() for t in current_transforms],
    }, model_

