"""CTF-specific wrapper around the expanding-window OOS prediction loop.

Provides ``ctf_expanding_window_oos()``, a convenience entry point that
accepts the three raw CTF DataFrames (chars, features, daily_ret), applies
feature selection and rank normalization, constructs the long-format panel
expected by ``expanding_window_oos()``, and returns an ``OOSResult`` whose
``to_ctf_weights()`` is immediately usable.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

import pandas as pd

from eapctf.ctf.prep import ctf_rank_normalize, ctf_select_features
from eapctf.predict.models import BasePredictor
from eapctf.predict.oos import OOSResult, expanding_window_oos

if TYPE_CHECKING:
    pass

_VALID_SUBSETS = frozenset({"jkp153", "all", "available"})


def ctf_expanding_window_oos(
    chars: pd.DataFrame,
    features: pd.DataFrame,
    models: list[BasePredictor] | None = None,
    window_length: int = 120,
    rank_normalize: bool = True,
    feature_subset: str = "all",
    test_only: bool = True,
    ctff_test_col: str = "ctff_test",
    eom_col: str = "eom",
    id_col: str = "id",
    ret_col: str = "ret_exc_lead1m",
    **kwargs: Any,
) -> OOSResult:
    """CTF-specific wrapper around ``expanding_window_oos()``.

    Handles the three-step preprocessing pipeline (feature selection →
    rank normalization → panel construction) before delegating to the
    standard OOS loop.  The returned ``OOSResult.predictions`` contains
    only test-period dates by default (``test_only=True``), so
    ``oos.to_ctf_weights()`` is immediately usable without further filtering.

    Typical usage inside a CTF ``main()`` function::

        from eapctf.ctf.predict import ctf_expanding_window_oos
        from eapctf.predict import make_predictor

        def main(chars, features, daily_ret):
            oos = ctf_expanding_window_oos(
                chars, features,
                models=[make_predictor("ridge")],
            )
            return oos.to_ctf_weights(signal_col="pred_RidgePredictor", mode="zscore")

    Parameters
    ----------
    chars : pd.DataFrame
        Full CTF characteristics panel (firms × dates, long format). Must
        contain ``eom_col``, ``id_col``, ``ret_col``, ``ctff_test_col``,
        and the feature columns listed in ``features``.
    features : pd.DataFrame
        Single-column DataFrame with column ``"features"`` listing the
        402 CTF characteristic names.
    models : list[BasePredictor] | None
        Predictor instances to train and evaluate.  Defaults to
        ``[LassoPredictor(), RidgePredictor()]`` (inherited from
        ``expanding_window_oos``).
    window_length : int
        Rolling training window size in months.  Default 120 (10 years,
        HJKP 2025 standard).
    rank_normalize : bool
        If True (default), apply ``ctf_rank_normalize()`` before training.
        Set to False if the characteristics are already normalized.
    feature_subset : str
        Feature selection mode passed to ``ctf_select_features()``.
        One of ``"all"`` (default), ``"jkp153"``, ``"available"``.
    test_only : bool
        If True (default), filter predictions to dates where
        ``ctff_test_col == True`` before returning.  Set to False to
        include training-period predictions as well.
    ctff_test_col : str
        Column in ``chars`` that flags test-period dates.  Default
        ``"ctff_test"``.
    eom_col : str
        End-of-month formation date column in ``chars``.  Default ``"eom"``.
    id_col : str
        Firm identifier column in ``chars``.  Default ``"id"``.
    ret_col : str
        1-month-ahead excess return column in ``chars``.  Default
        ``"ret_exc_lead1m"``.
    **kwargs : Any
        Additional keyword arguments forwarded to ``expanding_window_oos()``.
        Use to override ``window_type``, ``retrain_freq``, ``n_cv_folds``,
        etc.

    Returns
    -------
    OOSResult
        ``predictions``: DataFrame with columns
        ``[date, permno, ret, pred_<Model>, ...]``.  If ``test_only=True``,
        only test-period rows are included; call
        ``oos.to_ctf_weights(signal_col=..., mode=...)`` directly.

        ``models_history``: dict mapping each retrain date to fitted models.

    Raises
    ------
    ValueError
        If ``feature_subset`` is not one of the accepted values.

    References
    ----------
    Hellum, Jensen, Kelly, and Pedersen (2025). The Power of the Common Task
        Framework. SSRN 5242901.
    """
    if feature_subset not in _VALID_SUBSETS:
        raise ValueError(
            f"feature_subset must be one of {set(_VALID_SUBSETS)}, "
            f"got {feature_subset!r}"
        )

    # --- Step 1: Feature selection ---
    chars_sel, selected_features = ctf_select_features(
        chars, features, subset=feature_subset, eom_col=eom_col
    )

    # --- Step 2: Rank normalization ---
    if rank_normalize:
        chars_sel = ctf_rank_normalize(
            chars_sel, features, eom_col=eom_col
        )
        # Normalization is done; tell expanding_window_oos to skip its own
        rank_normalize_chars = False
    else:
        rank_normalize_chars = False

    # --- Step 3: Build the long-format panel ---
    # expanding_window_oos expects: date_col, entity_col, ret_col, char_cols
    # CTF uses: eom → date, id → permno, ret_exc_lead1m → ret
    keep_cols = [eom_col, id_col, ret_col] + selected_features
    # Only keep columns that are present (ret_col may be absent in some slices)
    available = [c for c in keep_cols if c in chars_sel.columns]
    panel = chars_sel[available].copy()

    # Rename to the standard names expected by expanding_window_oos defaults
    panel = panel.rename(
        columns={eom_col: "date", id_col: "permno", ret_col: "ret"}
    )

    # Drop rows with missing return (can't supervise without a target)
    panel = panel.dropna(subset=["ret"]).reset_index(drop=True)

    # Record which dates are test dates (extracted from original chars before
    # the column rename drops ctff_test_col from the panel)
    if test_only and ctff_test_col in chars.columns:
        _test_dates: set[Any] = set(
            chars.loc[chars[ctff_test_col].astype(bool), eom_col].unique()
        )
    else:
        _test_dates = set()

    # --- Step 4: Guard: no features selected → return empty OOSResult ---
    if not selected_features:
        return OOSResult(predictions=pd.DataFrame(), models_history={})

    # --- Step 5: Run the OOS loop ---
    oos = expanding_window_oos(
        data=panel,
        ret_col="ret",
        char_cols=selected_features,
        models=models,
        date_col="date",
        entity_col="permno",
        train_min_periods=window_length,
        window_type="rolling",
        train_window=window_length,
        rank_normalize_chars=rank_normalize_chars,
        **kwargs,
    )

    # --- Step 6: Filter to test dates if requested ---
    if test_only and _test_dates:
        test_preds = oos.predictions[
            oos.predictions["date"].isin(_test_dates)
        ].reset_index(drop=True)
        return OOSResult(
            predictions=test_preds,
            models_history=oos.models_history,
        )

    return oos
