"""Out-of-sample prediction framework for empirical asset pricing.

Implements the expanding-window OOS prediction loop following GKX (2020),
with mandatory temporal leakage checks at each rebalance.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

import pandas as pd

from eapctf.predict.models import BasePredictor, LassoPredictor, RidgePredictor
from eapctf.predict.prep import char_prep
from eapctf.utils.validation import check_no_lookahead


@dataclass
class OOSResult:
    """Container for out-of-sample prediction results.

    Attributes
    ----------
    predictions : pd.DataFrame
        DataFrame with columns [date, entity, ret, pred_<model1>, pred_<model2>, ...].
        ``ret`` is the realized return; ``pred_*`` columns are model predictions.
        All metrics (OOS R², Sharpe of prediction-sorted portfolios, etc.) should
        be computed from these columns after the fact.
    models_history : dict
        Maps each retrain date to a dict of {model_label: fitted_model}.
        Contains one entry per period where retraining occurred.
    """

    predictions: pd.DataFrame
    models_history: dict[Any, dict[str, BasePredictor]] = field(default_factory=dict)

    def to_ctf_weights(
        self,
        signal_col: str,
        mode: str = "long_short",
        n_stocks: int = 100,
        id_col: str = "permno",
        date_col: str = "date",
    ) -> pd.DataFrame:
        """Convert OOS predictions to CTF portfolio weights (id, eom, w).

        Parameters
        ----------
        signal_col : str
            Column in ``predictions`` to use as the ranking signal
            (e.g. ``"pred_LassoPredictor"``).
        mode : str
            ``"long_short"`` (default): top ``n_stocks`` → long (+1/n),
            bottom ``n_stocks`` → short (-1/n); weights sum to 0.
            ``"long_only"``: top ``n_stocks`` → equal weight; weights sum to 1.
            ``"zscore"``: z-score weights, normalized so sum(abs(w)) = 1; mean-zero.
            ``"rank_weighted"``: rank-centered weights, normalized so sum(abs(w)) = 1; mean-zero.
        n_stocks : int
            Number of stocks in each leg for ``"long_short"`` and ``"long_only"``.
            Ignored for ``"zscore"`` and ``"rank_weighted"`` (all stocks included).
            Default 100.
        id_col : str
            Column name for security identifier.  Default ``"permno"``.
        date_col : str
            Column name for date.  Default ``"date"``.

        Returns
        -------
        pd.DataFrame
            Columns: ``id`` (int), ``eom`` (date string YYYY-MM-DD), ``w`` (float).
            Guaranteed no duplicate (id, eom) pairs (CTF Rule 12).
        """
        _VALID_MODES = {"long_short", "long_only", "zscore", "rank_weighted"}
        if mode not in _VALID_MODES:
            raise ValueError(f"mode must be one of {_VALID_MODES}, got '{mode}'")

        preds = self.predictions
        if preds.empty:
            return pd.DataFrame(columns=["id", "eom", "w"])

        records: list[dict[str, Any]] = []
        for date_val, group in preds.groupby(date_col):
            eom_str = str(pd.Timestamp(date_val).date())  # type: ignore[arg-type]
            signal = group[signal_col]
            ids = group[id_col].astype(int).tolist()

            if mode == "zscore":
                # z-score weights: (pred - mean) / std, then normalize so sum(abs) = 1
                mu = float(signal.mean())
                sigma = float(signal.std(ddof=1))
                if sigma == 0 or signal.isna().all():
                    continue
                z = (signal - mu) / sigma
                abs_sum = float(z.abs().sum())
                if abs_sum == 0:
                    continue
                ws = (z / abs_sum).tolist()
                for stock_id, w in zip(ids, ws):
                    records.append({"id": stock_id, "eom": eom_str, "w": w})

            elif mode == "rank_weighted":
                # rank-weighted long-short: (rank - mean_rank) / sum(abs), mean-zero
                r = signal.rank(method="average", ascending=True)
                centered = r - float(r.mean())
                abs_sum = float(centered.abs().sum())
                if abs_sum == 0:
                    continue
                ws = (centered / abs_sum).tolist()
                for stock_id, w in zip(ids, ws):
                    records.append({"id": stock_id, "eom": eom_str, "w": w})

            else:
                ranked = signal.rank(method="first", ascending=False)
                if mode == "long_short":
                    top_mask = ranked <= n_stocks
                    # Exclude top-leg stocks from bottom leg to avoid duplicate
                    # (id, eom) pairs when group size < 2 * n_stocks (CTF Rule 12).
                    bot_mask = (ranked > (len(group) - n_stocks)) & (~top_mask)

                    n_long = int(top_mask.sum())
                    n_short = int(bot_mask.sum())
                    if n_long == 0:
                        continue
                    for idx in group.index[top_mask]:
                        records.append({
                            "id": int(group.loc[idx, id_col]),
                            "eom": eom_str,
                            "w": 1.0 / n_long,
                        })
                    if n_short > 0:
                        for idx in group.index[bot_mask]:
                            records.append({
                                "id": int(group.loc[idx, id_col]),
                                "eom": eom_str,
                                "w": -1.0 / n_short,
                            })
                else:  # long_only
                    top_mask = ranked <= n_stocks
                    n_long = int(top_mask.sum())
                    if n_long == 0:
                        continue
                    for idx in group.index[top_mask]:
                        records.append({
                            "id": int(group.loc[idx, id_col]),
                            "eom": eom_str,
                            "w": 1.0 / n_long,
                        })

        result = pd.DataFrame(records, columns=["id", "eom", "w"])
        result["id"] = result["id"].astype(int)
        return result


def _model_label(model: BasePredictor) -> str:
    """Generate a unique string label for a model instance."""
    return type(model).__name__


def expanding_window_oos(
    data: pd.DataFrame,
    ret_col: str = "ret",
    char_cols: list[str] | None = None,
    models: list[BasePredictor] | None = None,
    date_col: str = "date",
    entity_col: str = "permno",
    train_start: Any | None = None,
    train_min_periods: int = 120,
    window_type: str = "rolling",
    train_window: int = 120,
    retrain_freq: str = "monthly",
    val_window: int = 0,
    n_cv_folds: int = 5,
    retune_freq: str | None = None,
    rank_normalize_chars: bool = True,
    rank_scale: float = 0.5,
    use_jkp_whitelist: bool = False,
) -> OOSResult:
    """Rolling-window out-of-sample prediction loop (HJKP 2025 CTF standard).

    For each test date t:
        - train_data = most recent ``train_window`` periods before t
        - MANDATORY: check_no_lookahead(train_data, test_data_t)
        - Prepare chars: cross-sectional rank → [-0.5, 0.5], zero-fill missing
        - Refit models monthly; tune hyperparameters via 5-fold temporal CV

    Defaults follow the Common Task Framework methodology of Hellum, Jensen,
    Kelly, and Pedersen (2025, HJKP), Section 3.3:
        - 10-year (120-month) rolling training window
        - Monthly retraining
        - 5-fold temporal cross-validation for hyperparameter selection
        - Rank normalization to [-0.5, 0.5]; missing values → 0 (= median)

    To replicate GKX (2020) instead, use:
        ``window_type="expanding", train_window=None, rank_scale=1.0,
          val_window=12, retrain_freq="annual"``

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel with date, entity, return, and characteristic cols.
    ret_col : str
        Name of the return column.
    char_cols : list[str] | None
        Characteristic columns. If None, auto-detect by excluding
        date_col, entity_col, and ret_col.
    models : list[BasePredictor] | None
        List of predictor instances. Defaults to [LassoPredictor(), RidgePredictor()].
    date_col : str
        Date column name.
    entity_col : str
        Entity identifier column name.
    train_start : Any | None
        Earliest date to include in training. If None, use all available.
    train_min_periods : int
        Minimum number of unique training dates before first prediction.
        Default 120 (10 years, matching train_window).
    window_type : str
        ``"rolling"`` (default, HJKP 2025): fixed-size sliding window of
        ``train_window`` most-recent periods.  ``"expanding"`` (GKX 2020):
        training set grows by one period each step.
    train_window : int
        Size of the rolling training window in periods.  Default 120 (10
        years × 12 months, per HJKP 2025).  Ignored when
        ``window_type="expanding"``.
    retrain_freq : str
        ``"monthly"`` (default, HJKP 2025) retrains every period;
        ``"annual"`` retrains once per year (GKX 2020).
    val_window : int
        Legacy rolling validation window for hyperparameter tuning.  Set to
        0 (default) when using ``n_cv_folds``-fold CV instead.  When > 0,
        the most-recent ``val_window`` training periods are held out.
    n_cv_folds : int
        Number of temporal folds for cross-validation hyperparameter tuning.
        Default 5 (HJKP 2025): 10-year window split into 5 folds of ~2 years,
        each taking turns as validation.  Models implementing ``tune_cv()``
        receive this argument.  Set to 0 to disable CV-based tuning.
    retune_freq : str | None
        How often to re-select hyperparameters via ``tune()``.
        ``None`` (default): retune at every retrain point.
        ``"annual"``: retune once per year regardless of ``retrain_freq``.
        ``"never"``: never call ``tune()`` even if ``val_window > 0``.
    rank_normalize_chars : bool
        If True, apply rank normalization during char_prep.
    rank_scale : float
        Half-range of the rank normalization interval.  Default 0.5 →
        [-0.5, 0.5] (HJKP 2025).  Use 1.0 for GKX 2020 convention.
    use_jkp_whitelist : bool
        If True, restrict characteristics to the JKP-approved whitelist.

    Returns
    -------
    OOSResult
        ``predictions``: DataFrame with columns [date, entity, ret, pred_<model>, ...]
        where ``ret`` is the realized return at test time. Compute OOS R², portfolio
        sorts, or any other metric from these columns after the fact.

        ``models_history``: dict mapping each retrain date to the fitted models at
        that point, {date: {model_label: fitted_model}}.

    Raises
    ------
    eapctf.utils.validation.TemporalLeakageError
        If any training window overlaps with the test period.
    ValueError
        If parameters are invalid (see individual checks below).

    References
    ----------
    Hellum, Jensen, Kelly, and Pedersen (2025). The Power of the Common Task
        Framework. SSRN 5242901. Section 3.3.
    Gu, S., Kelly, B., & Xiu, D. (2020). Empirical asset pricing via machine
        learning. Review of Financial Studies, 33(5), 2223-2273.
    """
    # Validate parameters
    if retrain_freq not in {"monthly", "annual"}:
        raise ValueError(
            f"retrain_freq must be 'monthly' or 'annual', got '{retrain_freq}'"
        )
    if window_type not in {"expanding", "rolling"}:
        raise ValueError(
            f"window_type must be 'expanding' or 'rolling', got '{window_type}'"
        )
    if window_type == "rolling" and train_window is None:
        raise ValueError("train_window must be specified when window_type='rolling'")
    if retune_freq not in {None, "annual", "never"}:
        raise ValueError(
            f"retune_freq must be None, 'annual', or 'never', got '{retune_freq}'"
        )

    # Default models: Lasso + Ridge
    if models is None:
        models = [LassoPredictor(), RidgePredictor()]

    # Auto-detect characteristic columns if not specified
    if char_cols is None:
        exclude = {date_col, entity_col, ret_col}
        char_cols = [c for c in data.columns if c not in exclude]

    # Sort data by date for proper temporal ordering
    data_sorted = data.sort_values(date_col).reset_index(drop=True)

    # Filter by train_start if provided
    if train_start is not None:
        data_sorted = data_sorted[data_sorted[date_col] >= train_start].reset_index(drop=True)

    # Get unique sorted dates
    unique_dates = sorted(data_sorted[date_col].unique())
    n_dates = len(unique_dates)

    # Deduplicate model labels if multiple instances of same class
    raw_labels = [_model_label(m) for m in models]
    label_counts: dict[str, int] = {}
    model_labels: list[str] = []
    for lbl in raw_labels:
        count = label_counts.get(lbl, 0)
        label_counts[lbl] = count + 1
        model_labels.append(f"{lbl}_{count}" if count > 0 else lbl)

    # Collect predictions and models across all test periods
    all_predictions: list[pd.DataFrame] = []
    models_history: dict[Any, dict[str, BasePredictor]] = {}

    # Track state for annual retrain/retune
    last_train_year: int | None = None
    last_retune_year: int | None = None

    for t_idx in range(train_min_periods, n_dates):
        test_date = unique_dates[t_idx]

        # Extract year for annual-mode decisions
        try:
            current_year = int(pd.Timestamp(test_date).year)
        except Exception:
            current_year = int(test_date) // 12

        # Determine if we should retrain at this period
        if retrain_freq == "annual":
            if last_train_year is not None and current_year == last_train_year:
                need_retrain = False
            else:
                need_retrain = True
                last_train_year = current_year
        else:
            need_retrain = True

        # Determine if we should retune hyperparameters at this period
        if retune_freq == "never" or val_window == 0:
            need_retune = False
        elif retune_freq == "annual":
            if last_retune_year is not None and current_year == last_retune_year:
                need_retune = False
            else:
                need_retune = need_retrain
                if need_retune:
                    last_retune_year = current_year
        else:
            # None: retune at every retrain point
            need_retune = need_retrain

        # Build training set: all dates strictly before test_date
        all_train_dates = unique_dates[:t_idx]
        if window_type == "rolling" and train_window is not None:
            all_train_dates = all_train_dates[-train_window:]

        train_mask = data_sorted[date_col].isin(set(all_train_dates))
        test_mask = data_sorted[date_col] == test_date
        train_data = data_sorted[train_mask]
        test_data = data_sorted[test_mask]

        if len(test_data) == 0:
            continue

        # MANDATORY temporal leakage check
        check_no_lookahead(train_data, test_data, date_col=date_col)

        # Prepare characteristics on training data
        train_prepped, used_cols = char_prep(
            train_data,
            char_cols=char_cols,
            use_jkp_whitelist=use_jkp_whitelist,
            rank_normalize_chars=rank_normalize_chars,
            rank_scale=rank_scale,
            date_col=date_col,
            entity_col=entity_col,
        )

        # Prepare test data with same column set
        test_prepped, _ = char_prep(
            test_data,
            char_cols=used_cols,
            use_jkp_whitelist=use_jkp_whitelist,
            rank_normalize_chars=rank_normalize_chars,
            rank_scale=rank_scale,
            date_col=date_col,
            entity_col=entity_col,
        )

        X_train = train_prepped[used_cols]
        y_train = train_prepped[ret_col]
        X_test = test_prepped[used_cols]

        # At retune points, select hyperparameters via CV or held-out window
        if need_retune:
            train_dates_series = train_prepped[date_col]
            if n_cv_folds > 0:
                # HJKP 2025: walk-forward temporal CV over n_cv_folds folds
                for model in models:
                    if hasattr(model, "tune_cv") and callable(model.tune_cv):
                        model.tune_cv(
                            X_train, y_train,
                            dates=train_dates_series,
                            n_folds=n_cv_folds,
                        )
            elif val_window > 0:
                # GKX 2020 fallback: single held-out rolling validation window
                train_dates_sorted = sorted(train_dates_series.unique())
                if len(train_dates_sorted) > val_window:
                    val_dates = set(train_dates_sorted[-val_window:])
                    val_mask = train_prepped[date_col].isin(val_dates)
                    X_tune_train = X_train[~val_mask]
                    y_tune_train = y_train[~val_mask]
                    X_tune_val = X_train[val_mask]
                    y_tune_val = y_train[val_mask]
                    for model in models:
                        if hasattr(model, "tune") and callable(model.tune):
                            model.tune(X_tune_train, y_tune_train, X_tune_val, y_tune_val)

        # Fit and predict; store fitted models at each retrain point
        period_preds = test_prepped[[date_col, entity_col]].copy()
        period_preds[ret_col] = test_prepped[ret_col].values  # realized return

        # Campbell & Thompson (2008) historical-mean benchmark: expanding mean of
        # training returns, constant across all stocks in this test period.
        period_preds["pred_HistMean"] = float(y_train.mean())

        period_models: dict[str, BasePredictor] = {}
        for model, lbl in zip(models, model_labels):
            if need_retrain:
                model.fit(X_train, y_train)
            period_preds[f"pred_{lbl}"] = model.predict(X_test).values
            period_models[lbl] = model

        all_predictions.append(period_preds)
        models_history[test_date] = period_models

    # Return empty result if not enough data for any predictions
    if len(all_predictions) == 0:
        return OOSResult(predictions=pd.DataFrame(), models_history={})

    predictions_df = pd.concat(all_predictions, ignore_index=True)
    return OOSResult(predictions=predictions_df, models_history=models_history)
