"""Prediction models for empirical asset pricing.

Provides a unified BasePredictor interface wrapping linear models (OLS, Lasso,
Ridge, Elastic Net), tree ensembles (Random Forest, Gradient Boosted Trees,
LightGBM), a neural network (GKX 2020 architecture), and Instrumented PCA (KPS 2019).

GKX (2020) baseline models: OLSPredictor (OLS3), LassoPredictor, RidgePredictor,
ElasticNetPredictor, RandomForestPredictor (300 trees, depth 6),
GBTPredictor (1000 trees, depth 2, lr=0.01).
LightGBMPredictor: leaf-wise gradient boosting via the LightGBM library.
"""

from __future__ import annotations

import warnings
from abc import ABC, abstractmethod
from typing import Any

import numpy as np
import pandas as pd
from sklearn.ensemble import GradientBoostingRegressor, RandomForestRegressor
from sklearn.linear_model import ElasticNet, Lasso, LinearRegression, Ridge

# Optional dependency: PyTorch
try:
    import torch
    import torch.nn as nn

    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False


# ---------------------------------------------------------------------------
# Abstract base
# ---------------------------------------------------------------------------


class BasePredictor(ABC):
    """Abstract base class for all prediction models in eapctf.predict."""

    def tune(
        self,
        X_train: pd.DataFrame,
        y_train: pd.Series,
        X_val: pd.DataFrame,
        y_val: pd.Series,
    ) -> None:
        """Optional hyperparameter selection using a held-out validation set.

        Override in subclasses that support hyperparameter tuning (e.g.
        CVPredictor).  The default implementation is a no-op so that plain
        predictors can be passed through the OOS loop unchanged.

        Called by ``expanding_window_oos`` at every retrain point when
        ``val_window > 0``.  After ``tune()`` the loop calls ``fit()`` on
        the full training data, so subclasses should only store the selected
        hyperparameters here — actual model fitting happens in ``fit()``.
        """

    @abstractmethod
    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Fit the model on training data.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix (N observations x K characteristics).
        y : pd.Series
            Target variable (returns).
        """

    @abstractmethod
    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Generate predictions for new data.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix with the same columns used during fit.

        Returns
        -------
        pd.Series
            Predicted values, indexed to match X.
        """


# ---------------------------------------------------------------------------
# Linear penalized models
# ---------------------------------------------------------------------------


class LassoPredictor(BasePredictor):
    """L1-penalized linear regression (Lasso).

    Wraps ``sklearn.linear_model.Lasso`` with aligned prediction output.
    """

    def __init__(self, alpha: float = 0.001, **kwargs: Any) -> None:
        self.alpha = alpha
        self.kwargs = kwargs
        # Internal sklearn model; created at fit time
        self._model: Lasso | None = None
        self._feature_names: list[str] = []

    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Fit Lasso on training data."""
        self._feature_names = list(X.columns)
        self._model = Lasso(alpha=self.alpha, max_iter=10_000, **self.kwargs)
        self._model.fit(X.values, y.values)

    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Predict using fitted Lasso, returning entity-indexed Series."""
        if self._model is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        # Align columns to training feature order
        preds = self._model.predict(X[self._feature_names].values)
        return pd.Series(preds, index=X.index, name="pred")


class RidgePredictor(BasePredictor):
    """L2-penalized linear regression (Ridge).

    Wraps ``sklearn.linear_model.Ridge`` with aligned prediction output.
    """

    def __init__(self, alpha: float = 1.0, **kwargs: Any) -> None:
        self.alpha = alpha
        self.kwargs = kwargs
        self._model: Ridge | None = None
        self._feature_names: list[str] = []

    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Fit Ridge on training data."""
        self._feature_names = list(X.columns)
        self._model = Ridge(alpha=self.alpha, **self.kwargs)
        self._model.fit(X.values, y.values)

    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Predict using fitted Ridge, returning entity-indexed Series."""
        if self._model is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        preds = self._model.predict(X[self._feature_names].values)
        return pd.Series(preds, index=X.index, name="pred")


class ElasticNetPredictor(BasePredictor):
    """Elastic Net (combined L1 + L2 penalty).

    Wraps ``sklearn.linear_model.ElasticNet`` with aligned prediction output.
    """

    def __init__(
        self,
        alpha: float = 0.001,
        l1_ratio: float = 0.5,
        **kwargs: Any,
    ) -> None:
        self.alpha = alpha
        self.l1_ratio = l1_ratio
        self.kwargs = kwargs
        self._model: ElasticNet | None = None
        self._feature_names: list[str] = []

    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Fit Elastic Net on training data."""
        self._feature_names = list(X.columns)
        self._model = ElasticNet(
            alpha=self.alpha,
            l1_ratio=self.l1_ratio,
            max_iter=10_000,
            **self.kwargs,
        )
        self._model.fit(X.values, y.values)

    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Predict using fitted Elastic Net, returning entity-indexed Series."""
        if self._model is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        preds = self._model.predict(X[self._feature_names].values)
        return pd.Series(preds, index=X.index, name="pred")


class OLSPredictor(BasePredictor):
    """Ordinary least-squares linear regression.

    Baseline predictor from GKX (2020) Sec. 3 (OLS3).
    Wraps ``sklearn.linear_model.LinearRegression`` with aligned output.

    References
    ----------
    Gu, S., Kelly, B., & Xiu, D. (2020). Empirical asset pricing via machine
        learning. RFS, 33(5), 2223-2273. Sec. 3 (OLS3 baseline).
    """

    def __init__(self, **kwargs: Any) -> None:
        self.kwargs = kwargs
        self._model: LinearRegression | None = None
        self._feature_names: list[str] = []

    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Fit OLS on training data."""
        self._feature_names = list(X.columns)
        self._model = LinearRegression(**self.kwargs)
        self._model.fit(X.values, y.values)

    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Predict using fitted OLS, returning entity-indexed Series."""
        if self._model is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        preds = self._model.predict(X[self._feature_names].values)
        return pd.Series(preds, index=X.index, name="pred")


# ---------------------------------------------------------------------------
# Tree-based models
# ---------------------------------------------------------------------------


class RandomForestPredictor(BasePredictor):
    """Random Forest regression.

    Wraps ``sklearn.ensemble.RandomForestRegressor`` with aligned output.
    """

    def __init__(
        self,
        n_estimators: int = 300,
        max_depth: int | None = 6,
        **kwargs: Any,
    ) -> None:
        self.n_estimators = n_estimators
        self.max_depth = max_depth
        self.kwargs = kwargs
        self._model: RandomForestRegressor | None = None
        self._feature_names: list[str] = []

    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Fit Random Forest on training data."""
        self._feature_names = list(X.columns)
        self._model = RandomForestRegressor(
            n_estimators=self.n_estimators,
            max_depth=self.max_depth,
            n_jobs=-1,
            **self.kwargs,
        )
        self._model.fit(X.values, y.values)

    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Predict using fitted Random Forest, returning entity-indexed Series."""
        if self._model is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        preds = self._model.predict(X[self._feature_names].values)
        return pd.Series(preds, index=X.index, name="pred")


class GBTPredictor(BasePredictor):
    """Gradient Boosted Trees via sklearn GradientBoostingRegressor.

    Defaults match GKX (2020) Sec. 3: 1000 trees, max_depth=2 (stumps),
    learning_rate=0.01, using level-wise depth-limited tree growth.

    References
    ----------
    Gu, S., Kelly, B., & Xiu, D. (2020). Empirical asset pricing via machine
        learning. RFS, 33(5), 2223-2273. Sec. 3 (GBR/GBRT specification).
    """

    def __init__(
        self,
        n_estimators: int = 1000,
        max_depth: int = 2,
        learning_rate: float = 0.01,
        **kwargs: Any,
    ) -> None:
        self.n_estimators = n_estimators
        self.max_depth = max_depth
        self.learning_rate = learning_rate
        self.kwargs = kwargs
        self._model: GradientBoostingRegressor | None = None
        self._feature_names: list[str] = []

    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Fit GradientBoostingRegressor on training data."""
        self._feature_names = list(X.columns)
        self._model = GradientBoostingRegressor(
            n_estimators=self.n_estimators,
            max_depth=self.max_depth,
            learning_rate=self.learning_rate,
            **self.kwargs,
        )
        self._model.fit(X.values, y.values)

    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Predict using fitted GBT, returning entity-indexed Series."""
        if self._model is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        preds: np.ndarray = self._model.predict(X[self._feature_names].values)
        return pd.Series(preds, index=X.index, name="pred")


# ---------------------------------------------------------------------------
# Neural network (GKX 2020 architecture)
# ---------------------------------------------------------------------------


class NeuralNetPredictor(BasePredictor):
    """3-layer MLP via PyTorch with BatchNorm and Dropout.

    Architecture (following GKX 2020):
        Linear(n_input, hidden_sizes[0]) -> BatchNorm -> ReLU -> Dropout
        Linear(hidden_sizes[0], hidden_sizes[1]) -> BatchNorm -> ReLU -> Dropout
        Linear(hidden_sizes[1], 1)

    No internal StandardScaler — rank normalization is done upstream in
    ``char_prep``. Early stopping on validation loss is applied when
    ``val_fraction > 0``.

    Requires ``torch`` to be installed. Raises ``ImportError`` at fit time
    if the dependency is missing.

    References
    ----------
    Gu, S., Kelly, B., & Xiu, D. (2020). Empirical asset pricing via machine
        learning. Review of Financial Studies, 33(5), 2223-2273.
    """

    def __init__(
        self,
        hidden_sizes: tuple[int, ...] = (32, 32),
        dropout: float = 0.5,
        l1_lambda: float = 1e-3,
        lr: float = 1e-3,
        max_epochs: int = 100,
        patience: int = 10,
        val_fraction: float = 0.1,
        **kwargs: Any,
    ) -> None:
        self.hidden_sizes = hidden_sizes
        self.dropout = dropout
        self.l1_lambda = l1_lambda
        self.lr = lr
        self.max_epochs = max_epochs
        self.patience = patience
        self.val_fraction = val_fraction
        self.kwargs = kwargs
        self._model: Any = None
        self._feature_names: list[str] = []

    def _build_network(self, n_input: int) -> Any:
        """Construct the MLP architecture.

        Returns a ``torch.nn.Sequential`` module.
        """
        # Build layers following the GKX spec:
        # Input -> Hidden1 -> BN -> ReLU -> Dropout -> Hidden2 -> BN -> ReLU -> Dropout -> Output
        layers: list[Any] = []

        # First hidden layer
        layers.append(nn.Linear(n_input, self.hidden_sizes[0]))
        layers.append(nn.BatchNorm1d(self.hidden_sizes[0]))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(self.dropout))

        # Second hidden layer
        layers.append(nn.Linear(self.hidden_sizes[0], self.hidden_sizes[1]))
        layers.append(nn.BatchNorm1d(self.hidden_sizes[1]))
        layers.append(nn.ReLU())
        layers.append(nn.Dropout(self.dropout))

        # Output layer (single node for regression)
        layers.append(nn.Linear(self.hidden_sizes[1], 1))

        return nn.Sequential(*layers)

    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Fit neural network with early stopping on validation loss."""
        if not TORCH_AVAILABLE:
            raise ImportError(
                "torch is required for NeuralNetPredictor. "
                "Install it with: pip install torch"
            )
        self._feature_names = list(X.columns)
        n_input = X.shape[1]

        # Convert to numpy arrays for tensor creation
        X_arr = np.asarray(X.values, dtype=np.float32)
        y_arr = np.asarray(y.values, dtype=np.float32).reshape(-1, 1)

        # Train/validation split (temporal: use last val_fraction rows)
        n_total = len(X_arr)
        n_val = max(1, int(n_total * self.val_fraction))
        n_train = n_total - n_val

        X_train_t = torch.tensor(X_arr[:n_train])
        y_train_t = torch.tensor(y_arr[:n_train])
        X_val_t = torch.tensor(X_arr[n_train:])
        y_val_t = torch.tensor(y_arr[n_train:])

        # Build model and optimizer
        network = self._build_network(n_input)
        optimizer = torch.optim.Adam(network.parameters(), lr=self.lr)
        loss_fn = nn.MSELoss()
        # Training loop with early stopping
        best_val_loss = float("inf")
        patience_counter = 0
        best_state: dict[str, Any] = {}

        for _epoch in range(self.max_epochs):
            # Training step
            network.train()
            optimizer.zero_grad()
            pred_train = network(X_train_t)
            loss = loss_fn(pred_train, y_train_t)
            # L1 weight penalty (GKX 2020): sum of absolute weight values
            if self.l1_lambda > 0.0:
                l1_reg = sum(
                    p.abs().sum()
                    for p in network.parameters()
                    if p.requires_grad
                )
                loss = loss + self.l1_lambda * l1_reg
            loss.backward()
            optimizer.step()

            # Validation step
            network.eval()
            with torch.no_grad():
                pred_val = network(X_val_t)
                val_loss = loss_fn(pred_val, y_val_t).item()

            # Early stopping check
            if val_loss < best_val_loss:
                best_val_loss = val_loss
                patience_counter = 0
                best_state = {k: v.clone() for k, v in network.state_dict().items()}
            else:
                patience_counter += 1
                if patience_counter >= self.patience:
                    break

        # Restore best weights
        if best_state:
            network.load_state_dict(best_state)

        self._model = network

    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Predict using fitted neural network, returning entity-indexed Series."""
        if self._model is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        if not TORCH_AVAILABLE:
            raise ImportError(
                "torch is required for NeuralNetPredictor prediction."
            )
        X_t = torch.tensor(
            np.asarray(X[self._feature_names].values, dtype=np.float32)
        )
        self._model.eval()
        with torch.no_grad():
            preds = self._model(X_t).numpy().flatten()
        return pd.Series(preds, index=X.index, name="pred")


# ---------------------------------------------------------------------------
# LightGBM
# ---------------------------------------------------------------------------


class LightGBMPredictor(BasePredictor):
    """LightGBM gradient boosting regressor with optional early stopping.

    Uses leaf-wise tree growth (``num_leaves``) rather than sklearn's
    level-wise (``max_depth``), which is faster and often more accurate
    for tabular data with many features.

    Early stopping is applied on a hold-out validation split when
    ``val_fraction > 0``: training halts once validation loss stops improving
    for ``early_stopping_rounds`` consecutive rounds and the best iteration
    is restored.  Set ``val_fraction=0`` to disable early stopping and train
    for exactly ``n_estimators`` rounds.

    Parameters
    ----------
    n_estimators : int
        Maximum number of boosting rounds. Default 1000.
    num_leaves : int
        Maximum number of leaves per tree. Controls model capacity.
        Default 31 (LightGBM default).
    learning_rate : float
        Shrinkage applied to each tree. Default 0.01.
    min_child_samples : int
        Minimum observations in a leaf. Default 20.
    subsample : float
        Fraction of training rows sampled per tree. Default 0.8.
    colsample_bytree : float
        Fraction of features sampled per tree. Default 0.8.
    reg_alpha : float
        L1 regularisation on leaf weights. Default 0.0.
    reg_lambda : float
        L2 regularisation on leaf weights. Default 0.0.
    val_fraction : float
        Fraction of training data held out for early stopping.
        Must be in (0, 1). Set to 0 to disable. Default 0.1.
    early_stopping_rounds : int
        Patience for early stopping. Default 50.
    verbose : int
        LightGBM verbosity level. Default -1 (silent).
    **kwargs : Any
        Additional keyword arguments forwarded to ``lgb.train``.

    References
    ----------
    Ke, G., et al. (2017). LightGBM: A highly efficient gradient boosting
        decision tree. NeurIPS 30.
    """

    def __init__(
        self,
        n_estimators: int = 1000,
        num_leaves: int = 31,
        learning_rate: float = 0.01,
        min_child_samples: int = 20,
        subsample: float = 0.8,
        colsample_bytree: float = 0.8,
        reg_alpha: float = 0.0,
        reg_lambda: float = 0.0,
        val_fraction: float = 0.1,
        early_stopping_rounds: int = 50,
        verbose: int = -1,
        **kwargs: Any,
    ) -> None:
        self.n_estimators = n_estimators
        self.num_leaves = num_leaves
        self.learning_rate = learning_rate
        self.min_child_samples = min_child_samples
        self.subsample = subsample
        self.colsample_bytree = colsample_bytree
        self.reg_alpha = reg_alpha
        self.reg_lambda = reg_lambda
        self.val_fraction = val_fraction
        self.early_stopping_rounds = early_stopping_rounds
        self.verbose = verbose
        self.kwargs = kwargs
        self._booster: Any = None
        self._feature_names: list[str] = []

    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Fit LightGBM regressor on training data."""
        try:
            import lightgbm as lgb
        except ImportError as exc:
            raise ImportError(
                "lightgbm is required for LightGBMPredictor. "
                "Install it with: pip install lightgbm"
            ) from exc

        self._feature_names = list(X.columns)
        X_arr = X.values.astype(np.float64)
        y_arr = y.values.astype(np.float64)

        params = {
            "objective": "regression",
            "metric": "mse",
            "num_leaves": self.num_leaves,
            "learning_rate": self.learning_rate,
            "min_child_samples": self.min_child_samples,
            "subsample": self.subsample,
            "colsample_bytree": self.colsample_bytree,
            "reg_alpha": self.reg_alpha,
            "reg_lambda": self.reg_lambda,
            "verbose": self.verbose,
            "n_jobs": -1,
        }
        params.update(self.kwargs)

        callbacks: list[Any] = []
        valid_sets: list[Any] = []

        if self.val_fraction > 0.0:
            # Temporal split: last val_fraction rows are validation
            n_val = max(1, int(len(X_arr) * self.val_fraction))
            n_train = len(X_arr) - n_val
            train_ds = lgb.Dataset(X_arr[:n_train], label=y_arr[:n_train])
            val_ds = lgb.Dataset(X_arr[n_train:], label=y_arr[n_train:], reference=train_ds)
            valid_sets = [val_ds]
            callbacks.append(
                lgb.early_stopping(
                    stopping_rounds=self.early_stopping_rounds, verbose=False
                )
            )
            callbacks.append(lgb.log_evaluation(period=-1))
        else:
            train_ds = lgb.Dataset(X_arr, label=y_arr)

        self._booster = lgb.train(
            params,
            train_ds,
            num_boost_round=self.n_estimators,
            valid_sets=valid_sets if valid_sets else None,
            callbacks=callbacks if callbacks else None,
        )

    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Predict using fitted LightGBM model, returning entity-indexed Series."""
        if self._booster is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        preds = self._booster.predict(X[self._feature_names].values.astype(np.float64))
        return pd.Series(np.asarray(preds).ravel(), index=X.index, name="pred")


# ---------------------------------------------------------------------------
# Instrumented PCA (Kelly, Pruitt, Su 2019)
# ---------------------------------------------------------------------------


class IPCAPredictor(BasePredictor):
    """Instrumented PCA -- thin wrapper around bkelly-lab/ipca (Kelly, Pruitt, Su 2019).

    Delegates all estimation to ``ipca.InstrumentedPCA``, which implements
    the correct alternating least squares algorithm with SVD initialization,
    Cholesky+SVD orthogonalization of Gamma, and proper unbalanced-panel support.

    Panel indices (entity id, time id) are required for multi-period IPCA.
    Pass ``date_col`` and ``entity_col`` as constructor parameters OR include
    those columns in X at fit/predict time.  When neither is available the
    predictor falls back to single-period mode (equivalent to reduced-rank OLS)
    and emits a UserWarning.

    Parameters
    ----------
    n_factors : int
        Number of latent IPCA factors L. Default 3.
    max_iter : int
        Maximum ALS iterations. Default 10000 (bkelly-lab default).
    tol : float
        ALS convergence tolerance. Default 1e-6.
    intercept : bool
        Whether to include an intercept (alpha) factor. Default False.
    date_col : str or None
        Name of the date/period column. When present in X at fit/predict time,
        used to build time indices for the panel. Default None.
    entity_col : str or None
        Name of the entity column. When present in X at fit/predict time,
        used to build entity indices for the panel. Default None.

    References
    ----------
    Kelly, B., Pruitt, S., & Su, Y. (2019). Characteristics are covariances:
        A unified model of risk and return. Journal of Financial Economics,
        134(3), 501-524.
    """

    def __init__(
        self,
        n_factors: int = 3,
        max_iter: int = 10000,
        tol: float = 1e-6,
        intercept: bool = False,
        date_col: str | None = None,
        entity_col: str | None = None,
    ) -> None:
        self.n_factors = n_factors
        self.max_iter = max_iter
        self.tol = tol
        self.intercept = intercept
        self.date_col = date_col
        self.entity_col = entity_col
        self._ipca: Any = None
        self._feature_names: list[str] = []
        self._single_period: bool = False

    def _extract_indices(
        self, X: pd.DataFrame
    ) -> tuple[np.ndarray, list[str], bool]:
        """Build (entity, time) index array and characteristic column list.

        Returns
        -------
        indices : np.ndarray, shape (N, 2)  -- [entity_id, time_id] per row
        char_cols : list[str]               -- columns to use as instruments
        single_period : bool                -- True when falling back to T=1
        """
        drop: set[str] = set()
        dc = self.date_col
        ec = self.entity_col
        if dc and dc in X.columns:
            drop.add(dc)
        if ec and ec in X.columns:
            drop.add(ec)

        char_cols = [c for c in X.columns if c not in drop]

        if dc and ec and dc in X.columns and ec in X.columns:
            dates = X[dc].values
            entities = X[ec].values
            date_map = {d: i for i, d in enumerate(sorted(set(dates)))}
            entity_map = {e: i for i, e in enumerate(sorted(set(entities)))}
            time_idx = np.array([date_map[d] for d in dates], dtype=np.int64)
            entity_idx = np.array([entity_map[e] for e in entities], dtype=np.int64)
            indices = np.column_stack([entity_idx, time_idx])
            return indices, char_cols, False

        # Fallback: single cross-section (T=1, each row is a unique entity)
        N = len(X)
        indices = np.column_stack(
            [np.arange(N, dtype=np.int64), np.zeros(N, dtype=np.int64)]
        )
        return indices, char_cols, True

    def fit(
        self,
        X: pd.DataFrame,
        y: pd.Series,
        date_col: str | None = None,
        entity_col: str | None = None,
    ) -> None:
        """Fit IPCA via the official bkelly-lab ALS algorithm.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix. May optionally contain date_col / entity_col
            columns for proper panel indexing. If absent, the predictor
            falls back to single-period mode with a UserWarning.
        y : pd.Series
            Target variable (returns).
        date_col : str or None
            Overrides the constructor date_col for this call.
        entity_col : str or None
            Overrides the constructor entity_col for this call.
        """
        from ipca import InstrumentedPCA  # lazy import to keep ipca optional

        if date_col is not None:
            self.date_col = date_col
        if entity_col is not None:
            self.entity_col = entity_col

        indices, char_cols, single_period = self._extract_indices(X)
        self._feature_names = char_cols
        self._single_period = single_period

        if single_period:
            warnings.warn(
                "IPCAPredictor: date_col and/or entity_col not found in X. "
                "Falling back to single-period IPCA (equivalent to reduced-rank "
                "regression). For proper panel IPCA pass date_col and entity_col "
                "as constructor parameters or include them in X.",
                UserWarning,
                stacklevel=2,
            )

        X_chars = X[char_cols].values.astype(np.float64)
        y_arr = y.values.astype(np.float64)

        self._ipca = InstrumentedPCA(
            n_factors=self.n_factors,
            intercept=self.intercept,
            max_iter=self.max_iter,
            iter_tol=self.tol,
        )
        self._ipca.fit(X_chars, y_arr, indices=indices, data_type="panel")

    def predict(
        self,
        X: pd.DataFrame,
        date_col: str | None = None,
        entity_col: str | None = None,
    ) -> pd.Series:
        """Predict returns using fitted IPCA model.

        Uses time-averaged factors (mean_factor=True), consistent with the
        Kelly-Pruitt-Su (2019) out-of-sample forecasting convention.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix with the same characteristic columns used in fit.
        date_col : str or None
            Overrides the constructor date_col for this call.
        entity_col : str or None
            Overrides the constructor entity_col for this call.
        """
        if self._ipca is None:
            raise RuntimeError("Model not fitted. Call fit() first.")

        # Temporarily apply per-call overrides
        orig_dc, orig_ec = self.date_col, self.entity_col
        if date_col is not None:
            self.date_col = date_col
        if entity_col is not None:
            self.entity_col = entity_col

        indices, char_cols, _ = self._extract_indices(X)
        self.date_col, self.entity_col = orig_dc, orig_ec

        X_chars = X[char_cols].values.astype(np.float64)

        preds = self._ipca.predict(
            X=X_chars,
            indices=indices,
            mean_factor=True,
            data_type="panel",
        )
        return pd.Series(np.asarray(preds).ravel(), index=X.index, name="pred")


# ---------------------------------------------------------------------------
# Cross-validation wrapper
# ---------------------------------------------------------------------------


class CVPredictor(BasePredictor):
    """Hyperparameter-tuning wrapper following GKX (2020) validation methodology.

    Separates hyperparameter selection from model fitting:

    1. ``tune(X_train, y_train, X_val, y_val)`` — evaluate each parameter
       combination by training on X_train / y_train and scoring on the held-out
       X_val / y_val.  Stores the best combination in ``best_params_``.
    2. ``fit(X, y)`` — refit on the full dataset (train + validation) using
       ``best_params_``.  If ``tune()`` has not been called yet, falls back to
       the first entry in ``param_grid``.

    This design mirrors GKX (2020) Sec. 1.1: hyperparameters are selected once
    per retraining cycle on a rolling validation window, then the model is
    refitted on all available past data before generating OOS predictions.
    ``expanding_window_oos`` calls ``tune()`` automatically when
    ``val_window > 0``.

    Parameters
    ----------
    predictor_class : type[BasePredictor]
        The predictor class to tune (e.g. LassoPredictor).
    param_grid : dict[str, list]
        Hyperparameter grid, e.g. ``{"alpha": [0.0001, 0.001, 0.01, 0.1]}``.
    scoring : str
        ``"mse"`` (default) or ``"r2"``.  Lower validation MSE / higher R²
        wins.
    **fixed_kwargs : Any
        Fixed constructor kwargs forwarded regardless of the grid
        (e.g. ``max_iter=10_000``).

    Examples
    --------
    Tune Lasso alpha — GKX style with annual retuning::

        cv_lasso = CVPredictor(
            LassoPredictor,
            param_grid={"alpha": [0.0001, 0.001, 0.01, 0.1]},
        )
        result = expanding_window_oos(
            data, models=[cv_lasso],
            retrain_freq="annual", val_window=12,
        )

    Tune GBRT depth and learning rate::

        cv_gbt = CVPredictor(
            GBTPredictor,
            param_grid={
                "max_depth": [1, 2, 3],
                "learning_rate": [0.01, 0.05, 0.1],
            },
        )
    """

    def __init__(
        self,
        predictor_class: type[BasePredictor],
        param_grid: dict[str, list[Any]],
        scoring: str = "mse",
        **fixed_kwargs: Any,
    ) -> None:
        self.predictor_class = predictor_class
        self.param_grid = param_grid
        self.scoring = scoring
        self.fixed_kwargs = fixed_kwargs
        self.best_params_: dict[str, Any] = {}
        self._fitted_model: BasePredictor | None = None

    def _iter_param_combinations(self) -> list[dict[str, Any]]:
        """Cartesian product of param_grid values."""
        keys = list(self.param_grid.keys())
        values = [self.param_grid[k] for k in keys]
        combos: list[dict[str, Any]] = []

        def _expand(idx: int, current: dict[str, Any]) -> None:
            if idx == len(keys):
                combos.append(dict(current))
                return
            for v in values[idx]:
                current[keys[idx]] = v
                _expand(idx + 1, current)

        _expand(0, {})
        return combos

    def _score(
        self,
        y_true: np.ndarray,
        y_pred: np.ndarray,
    ) -> float:
        """Return validation score (lower = better for all metrics)."""
        if self.scoring == "mse":
            return float(np.mean((y_true - y_pred) ** 2))
        # r2: negate so lower is better
        ss_res = float(np.sum((y_true - y_pred) ** 2))
        ss_tot = float(np.sum((y_true - np.mean(y_true)) ** 2))
        return -(1.0 - ss_res / ss_tot) if ss_tot > 0 else 0.0

    def tune(
        self,
        X_train: pd.DataFrame,
        y_train: pd.Series,
        X_val: pd.DataFrame,
        y_val: pd.Series,
    ) -> None:
        """Select best hyperparameters using a held-out validation set.

        For each parameter combination: fit on (X_train, y_train), evaluate on
        (X_val, y_val).  The combination with the lowest validation score is
        stored in ``best_params_`` and used by the next ``fit()`` call.
        """
        combos = self._iter_param_combinations()
        y_val_arr = np.asarray(y_val.values, dtype=np.float64)

        best_score = np.inf
        best_combo = combos[0]

        for combo in combos:
            model = self.predictor_class(**combo, **self.fixed_kwargs)
            try:
                model.fit(X_train, y_train)
                preds = np.asarray(model.predict(X_val).values, dtype=np.float64)
            except Exception:
                continue
            score = self._score(y_val_arr, preds)
            if score < best_score:
                best_score = score
                best_combo = combo

        self.best_params_ = best_combo

    def tune_cv(
        self,
        X: pd.DataFrame,
        y: pd.Series,
        dates: pd.Series,
        n_folds: int = 5,
    ) -> None:
        """Select best hyperparameters using walk-forward temporal cross-validation.

        Implements the 5-fold temporal CV from Hellum, Jensen, Kelly, and Pedersen
        (2025, HJKP), Section 3.3: the training window is split into ``n_folds``
        consecutive temporal blocks; fold k uses blocks 0..k-1 as training data
        and block k as validation.  Block 0 is always training-only; blocks 1
        through n_folds-1 each serve once as the validation set.

        Parameters
        ----------
        X : pd.DataFrame
            Feature matrix aligned with y and dates.
        y : pd.Series
            Target returns, same index as X.
        dates : pd.Series
            Date values corresponding to rows of X (same index).  Used to split
            the window into temporal folds without lookahead.
        n_folds : int
            Number of temporal folds.  Default 5 (HJKP 2025): a 120-month window
            yields ~5 folds of 24 months each.

        References
        ----------
        Hellum, Jensen, Kelly, and Pedersen (2025). The Power of the Common Task
            Framework. SSRN 5242901. Section 3.3.
        """
        unique_dates = sorted(dates.unique())
        n_dates = len(unique_dates)

        if n_dates < n_folds:
            # Fall back to first param combo — not enough data to form folds
            combos = self._iter_param_combinations()
            self.best_params_ = combos[0] if combos else {}
            return

        # Split dates into n_folds consecutive blocks (last block absorbs remainder)
        base_size = n_dates // n_folds
        fold_edges: list[int] = []
        for f in range(n_folds + 1):
            fold_edges.append(f * base_size)
        fold_edges[-1] = n_dates  # ensure last fold includes all remaining dates

        combos = self._iter_param_combinations()
        # scores_per_combo[i] accumulates val scores across usable folds
        scores_per_combo: list[list[float]] = [[] for _ in combos]

        # Walk-forward: fold k (1-indexed) trains on blocks 0..k-1, validates on block k
        for k in range(1, n_folds):
            train_date_set = set(unique_dates[: fold_edges[k]])
            val_date_set = set(unique_dates[fold_edges[k] : fold_edges[k + 1]])

            train_mask = dates.isin(train_date_set)
            val_mask = dates.isin(val_date_set)

            X_tr = X[train_mask]
            y_tr = y[train_mask]
            X_v = X[val_mask]
            y_v = y[val_mask]

            if len(X_tr) == 0 or len(X_v) == 0:
                continue

            y_v_arr = np.asarray(y_v.values, dtype=np.float64)

            for i, combo in enumerate(combos):
                model = self.predictor_class(**combo, **self.fixed_kwargs)
                try:
                    model.fit(X_tr, y_tr)
                    preds = np.asarray(model.predict(X_v).values, dtype=np.float64)
                except Exception:
                    continue
                scores_per_combo[i].append(self._score(y_v_arr, preds))

        # Select combo with lowest average CV score; fall back to first if no scores
        best_score = np.inf
        best_combo = combos[0]
        for i, combo in enumerate(combos):
            if scores_per_combo[i]:
                avg = float(np.mean(scores_per_combo[i]))
                if avg < best_score:
                    best_score = avg
                    best_combo = combo

        self.best_params_ = best_combo

    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        """Fit on the full dataset using the hyperparameters from the last ``tune()`` call.

        Falls back to the first param_grid combination if ``tune()`` has not
        been called (e.g. for the very first OOS period before val_window is
        available).
        """
        params = self.best_params_ if self.best_params_ else self._iter_param_combinations()[0]
        self._fitted_model = self.predictor_class(**params, **self.fixed_kwargs)
        self._fitted_model.fit(X, y)

    def predict(self, X: pd.DataFrame) -> pd.Series:
        """Predict using the model fitted with the selected hyperparameters."""
        if self._fitted_model is None:
            raise RuntimeError("Model not fitted. Call fit() first.")
        return self._fitted_model.predict(X)


# ---------------------------------------------------------------------------
# Factory function
# ---------------------------------------------------------------------------

# Registry mapping string names to predictor classes
_MODEL_REGISTRY: dict[str, type[BasePredictor]] = {
    "ols": OLSPredictor,
    "lasso": LassoPredictor,
    "ridge": RidgePredictor,
    "elastic_net": ElasticNetPredictor,
    "rf": RandomForestPredictor,
    "gbt": GBTPredictor,
    "lgbm": LightGBMPredictor,
    "nn": NeuralNetPredictor,
    "ipca": IPCAPredictor,
}


def make_predictor(model_type: str, **kwargs: Any) -> BasePredictor:
    """Factory function to create a predictor by name.

    Parameters
    ----------
    model_type : str
        One of: ``ols``, ``lasso``, ``ridge``, ``elastic_net``, ``rf``,
        ``gbt``, ``nn``, ``ipca``.
    **kwargs : Any
        Keyword arguments forwarded to the predictor constructor.

    Returns
    -------
    BasePredictor
        An instance of the requested predictor class.

    Raises
    ------
    ValueError
        If ``model_type`` is not recognized.
    """
    if model_type not in _MODEL_REGISTRY:
        raise ValueError(
            f"Unknown model_type '{model_type}'. "
            f"Choose from: {list(_MODEL_REGISTRY.keys())}"
        )
    return _MODEL_REGISTRY[model_type](**kwargs)
