"""CTF conformal prediction wrapper with Adaptive Conformal Inference (ACI).

Implements MAPIE CV+ conformal intervals over the CTF preprocessing pipeline.
At each prediction date the conformal interval width is used to penalise
uncertain predictions before constructing portfolio weights.

References
----------
Gibbs, I. & Candes, E. (2021). Adaptive conformal inference under
    distribution shift. NeurIPS 2021.
Taquet, V. et al. (2022). MAPIE: an open-source library for
    distribution-free uncertainty quantification. JMLR 2023.
"""

from __future__ import annotations

import copy
import warnings
from typing import Any

import numpy as np
import pandas as pd
from scipy.stats import rankdata
from sklearn.base import BaseEstimator, RegressorMixin
from sklearn.model_selection import TimeSeriesSplit

# ---------------------------------------------------------------------------
# Optional dependency: MAPIE
# MAPIE v1.x (installed) uses CrossConformalRegressor instead of
# the deprecated MapieRegressor from v0.8.
# ---------------------------------------------------------------------------
try:
    from mapie.regression import CrossConformalRegressor

    MAPIE_AVAILABLE = True
except ImportError:  # pragma: no cover
    CrossConformalRegressor = None  # noqa: F811
    MAPIE_AVAILABLE = False

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

_VALID_SUBSETS = frozenset({"jkp153", "all", "available"})
_VALID_WEIGHT_MODES = frozenset({"zscore", "rank_weighted", "long_short"})
_VALID_REFIT_FREQS = frozenset({"yearly", "monthly"})


# ---------------------------------------------------------------------------
# Adaptive Conformal Inference
# ---------------------------------------------------------------------------


class AdaptiveConformalInference:
    """Online calibration of the conformal miscoverage level via ACI.

    Tracks a time-varying miscoverage target ``alpha_t``, updating it each
    period based on whether realized values fell within the predicted
    conformal intervals.  The update rule follows Gibbs & Candes (2021)::

        err_t   = 1 - coverage_rate_t          # fraction not covered
        alpha_t = clip(alpha_{t-1} + gamma * (alpha - err_t), 0.01, 0.99)

    When ``err_t > alpha`` (more misses than target) ``alpha_t`` decreases,
    which widens intervals at the next prediction step.  When ``err_t < alpha``
    (better coverage than needed) ``alpha_t`` increases, narrowing intervals.

    Parameters
    ----------
    alpha : float
        Target miscoverage rate. Default 0.1 (90 % coverage).
    gamma : float
        Step size for the ACI update. Smaller values smooth out fluctuations.
        Default 0.05.
    """

    def __init__(self, alpha: float = 0.1, gamma: float = 0.05) -> None:
        self.alpha = alpha
        self.gamma = gamma
        self.alpha_t: float = alpha

    def update(
        self,
        y_realized: np.ndarray,
        lower: np.ndarray,
        upper: np.ndarray,
    ) -> float:
        """Update ``alpha_t`` using cross-sectional coverage for one period.

        Parameters
        ----------
        y_realized : np.ndarray, shape (n,)
            Realized returns for the cross-section at the prior date.
        lower : np.ndarray, shape (n,)
            Lower conformal bounds predicted for that prior date.
        upper : np.ndarray, shape (n,)
            Upper conformal bounds predicted for that prior date.

        Returns
        -------
        float
            Updated ``alpha_t`` (also stored in ``self.alpha_t``).
        """
        covered = (lower <= y_realized) & (y_realized <= upper)
        # avg_err: fraction of stocks NOT covered
        avg_err = 1.0 - float(covered.mean())
        self.alpha_t = float(
            np.clip(
                self.alpha_t + self.gamma * (self.alpha - avg_err),
                0.01,
                0.99,
            )
        )
        return self.alpha_t


# ---------------------------------------------------------------------------
# Sklearn-compatible wrapper for MAPIE
# ---------------------------------------------------------------------------


class _SklearnWrapper(BaseEstimator, RegressorMixin):  # type: ignore[misc]
    """Wrap a :class:`~eapctf.predict.models.BasePredictor` for MAPIE.

    MAPIE requires a scikit-learn estimator interface (numpy array I/O,
    ``clone``-able).  This wrapper converts MAPIE's numpy array inputs to the
    pandas DataFrames expected by ``BasePredictor``.

    Parameters
    ----------
    predictor : BasePredictor
        Underlying predictor instance.  Will be deep-copied on ``__sklearn_clone__``.
    """

    def __init__(self, predictor: BasePredictor) -> None:
        self.predictor = predictor

    def fit(self, X: np.ndarray, y: np.ndarray) -> _SklearnWrapper:
        X_df = pd.DataFrame(X.astype(float))
        y_ser = pd.Series(y.astype(float))
        self.predictor.fit(X_df, y_ser)
        self.is_fitted_ = True
        return self

    def predict(self, X: np.ndarray) -> np.ndarray:
        X_df = pd.DataFrame(X.astype(float))
        result: np.ndarray = self.predictor.predict(X_df).to_numpy().astype(float)
        return result

    def __sklearn_clone__(self) -> _SklearnWrapper:
        return _SklearnWrapper(copy.deepcopy(self.predictor))


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _signal_to_weights(
    signal: np.ndarray,
    mode: str,
    n_stocks: int = 100,
) -> np.ndarray:
    """Convert a cross-sectional signal to portfolio weights."""
    n = len(signal)
    if n == 0:
        return np.array([], dtype=float)

    if mode == "zscore":
        mu = signal - signal.mean()
        abs_sum = float(np.abs(mu).sum())
        if abs_sum < 1e-12:
            return np.zeros(n, dtype=float)
        out: np.ndarray = (mu / abs_sum).astype(float)
        return out

    elif mode == "rank_weighted":
        ranks = rankdata(signal).astype(float) - (n + 1) / 2.0
        abs_sum = float(np.abs(ranks).sum())
        if abs_sum < 1e-12:
            return np.zeros(n, dtype=float)
        out = (ranks / abs_sum).astype(float)
        return out

    elif mode == "long_short":
        k = min(n_stocks, n // 2)
        w = np.zeros(n, dtype=float)
        idx_sorted = np.argsort(signal)
        if k > 0:
            w[idx_sorted[-k:]] = 1.0 / k    # long top-k
            w[idx_sorted[:k]] = -1.0 / k    # short bottom-k
        return w

    else:
        raise ValueError(
            f"weight_mode must be one of {_VALID_WEIGHT_MODES}, got {mode!r}"
        )


def _parse_pis(y_pis: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Extract lower and upper bounds from MAPIE's y_pis output.

    MAPIE v0.8+ returns shape (n, 2, n_alpha).  Earlier versions returned
    shape (n, 2).  This helper handles both conventions.
    """
    if y_pis.ndim == 3:
        lower = y_pis[:, 0, 0].astype(float)
        upper = y_pis[:, 1, 0].astype(float)
    else:
        lower = y_pis[:, 0].astype(float)
        upper = y_pis[:, 1].astype(float)
    # Ensure lower ≤ upper (defensive against any edge-case inversion)
    lo = np.minimum(lower, upper)
    hi = np.maximum(lower, upper)
    return lo, hi


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def ctf_conformal_oos(
    chars: pd.DataFrame,
    features: pd.DataFrame,
    model: BasePredictor,
    window_length: int = 120,
    rank_normalize: bool = True,
    feature_subset: str = "all",
    alpha: float = 0.1,
    gamma: float = 0.05,
    lambda_t: float = 1.0,
    refit_freq: str = "yearly",
    weight_mode: str = "zscore",
    n_conformal_folds: int = 5,
    test_only: bool = True,
    ctff_test_col: str = "ctff_test",
    eom_col: str = "eom",
    id_col: str = "id",
    ret_col: str = "ret_exc_lead1m",
) -> pd.DataFrame:
    """CTF pipeline with MAPIE CV+ conformal prediction and ACI recalibration.

    Builds on the standard CTF preprocessing pipeline (feature selection →
    rank normalization → panel construction) and adds a conformal prediction
    layer on top.  The workflow at each prediction date is:

    1. A MAPIE CV+ regressor (``method="plus"``,
       ``cv=TimeSeriesSplit(n_splits=n_conformal_folds)``) is refitted once
       per calendar year (or every month if ``refit_freq="monthly"``).
    2. The conformal interval width is used to compute an
       *uncertainty-adjusted* signal::

           mu_tilde_i = point_pred_i - lambda_t * (upper_i - lower_i) / 2

    3. The ACI alpha level ``alpha_t`` is updated each period based on
       whether realized returns fell inside the prior period's intervals::

           err_t     = 1 - coverage_rate_t
           alpha_t   = clip(alpha_{t-1} + gamma * (alpha - err_t), 0.01, 0.99)

    The ``alpha_t`` used for prediction at date ``d`` incorporates only
    information up to ``d - 1`` (no look-ahead).

    Parameters
    ----------
    chars : pd.DataFrame
        Full CTF characteristics panel (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 (column ``"features"``) listing CTF
        characteristic names.
    model : BasePredictor
        Predictor instance used as the base learner inside MAPIE.  A deep
        copy is created at each refit, so the original is not modified.
    window_length : int
        Rolling training window in months.  Default 120.
    rank_normalize : bool
        If True (default), apply ``ctf_rank_normalize()`` before training.
    feature_subset : str
        Feature selection mode: ``"all"`` (default), ``"jkp153"``, or
        ``"available"``.
    alpha : float
        Initial (and target) miscoverage rate for conformal intervals.
        Default 0.1 (90 % coverage).
    gamma : float
        ACI step size.  Default 0.05.
    lambda_t : float
        Uncertainty penalty scale.  Higher values penalise high-uncertainty
        predictions more aggressively.  Default 1.0.
    refit_freq : str
        MAPIE refit frequency: ``"yearly"`` (once per calendar year,
        default) or ``"monthly"`` (every period).
    weight_mode : str
        Portfolio weight construction from ``mu_tilde``: ``"zscore"``
        (default), ``"rank_weighted"``, or ``"long_short"``.
    n_conformal_folds : int
        Number of folds for ``TimeSeriesSplit`` inside MAPIE CV+. Default 5.
    test_only : bool
        If True (default), return only rows where ``ctff_test_col == True``.
    ctff_test_col : str
        Column flagging test-period dates.  Default ``"ctff_test"``.
    eom_col : str
        End-of-month formation date column.  Default ``"eom"``.
    id_col : str
        Firm identifier column.  Default ``"id"``.
    ret_col : str
        1-month-ahead excess return column.  Default ``"ret_exc_lead1m"``.

    Returns
    -------
    pd.DataFrame
        Columns: ``id``, ``eom``, ``w``, ``w_baseline``,
        ``interval_width``, ``alpha_t``.

        - ``w``: uncertainty-adjusted weights from ``mu_tilde``.
          For ``"zscore"`` and ``"rank_weighted"`` modes,
          ``sum(abs(w)) ≈ 1`` and ``sum(w) ≈ 0`` per date.
        - ``w_baseline``: weights from point prediction only
          (identical to what ``ctf_expanding_window_oos`` would produce
          with the same model and ``mode``).
        - ``interval_width``: conformal interval width ``upper - lower``.
        - ``alpha_t``: ACI miscoverage target at the time of prediction.

    Raises
    ------
    ImportError
        If ``mapie`` is not installed.
        Install with ``pip install "eapctf[conformal]"`` or
        ``pip install "mapie>=0.8.0"``.
    ValueError
        If ``feature_subset``, ``weight_mode``, or ``refit_freq`` are
        outside the accepted value sets.

    References
    ----------
    Gibbs, I. & Candes, E. (2021). Adaptive conformal inference under
        distribution shift. NeurIPS 2021.
    Taquet, V. et al. (2022). MAPIE: an open-source library for
        distribution-free uncertainty quantification. JMLR 2023.
    Hellum, Jensen, Kelly, and Pedersen (2025). The Power of the Common
        Task Framework. SSRN 5242901.
    """
    if not MAPIE_AVAILABLE:
        raise ImportError(
            "mapie is required for ctf_conformal_oos. "
            'Install with: pip install "eapctf[conformal]" or pip install mapie>=0.8.0'
        )
    if feature_subset not in _VALID_SUBSETS:
        raise ValueError(
            f"feature_subset must be one of {set(_VALID_SUBSETS)}, "
            f"got {feature_subset!r}"
        )
    if weight_mode not in _VALID_WEIGHT_MODES:
        raise ValueError(
            f"weight_mode must be one of {set(_VALID_WEIGHT_MODES)}, "
            f"got {weight_mode!r}"
        )
    if refit_freq not in _VALID_REFIT_FREQS:
        raise ValueError(
            f"refit_freq must be one of {set(_VALID_REFIT_FREQS)}, "
            f"got {refit_freq!r}"
        )

    _EMPTY_COLS = ["id", "eom", "w", "w_baseline", "interval_width", "alpha_t"]

    # ------------------------------------------------------------------
    # Preprocessing — identical to ctf_expanding_window_oos
    # ------------------------------------------------------------------
    chars_sel, selected_features = ctf_select_features(
        chars, features, subset=feature_subset, eom_col=eom_col
    )
    if rank_normalize:
        chars_sel = ctf_rank_normalize(chars_sel, features, eom_col=eom_col)

    keep_cols = [eom_col, id_col, ret_col] + selected_features
    available = [c for c in keep_cols if c in chars_sel.columns]
    panel = chars_sel[available].copy()
    panel = panel.rename(
        columns={eom_col: "date", id_col: "permno", ret_col: "ret"}
    )
    panel = panel.dropna(subset=["ret"]).reset_index(drop=True)

    # Extract test dates from original chars before column rename
    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()

    if not selected_features:
        return pd.DataFrame(columns=_EMPTY_COLS)

    feat_cols = selected_features
    all_dates = sorted(panel["date"].unique())

    # ------------------------------------------------------------------
    # Main prediction loop
    # ------------------------------------------------------------------
    aci = AdaptiveConformalInference(alpha=alpha, gamma=gamma)
    mapie_model: Any = None
    prev_year: int | None = None

    # State carried from previous iteration for ACI update
    prev_y_realized: np.ndarray | None = None
    prev_lower: np.ndarray | None = None
    prev_upper: np.ndarray | None = None

    rows: list[dict[str, Any]] = []

    for d in all_dates:
        # Require at least window_length prior months of training data
        train_dates = [dd for dd in all_dates if dd < d]
        if len(train_dates) < window_length:
            continue

        # ACI update: incorporate previous period's coverage before predicting
        if (
            prev_y_realized is not None
            and prev_lower is not None
            and prev_upper is not None
        ):
            aci.update(prev_y_realized, prev_lower, prev_upper)

        alpha_now = float(np.clip(aci.alpha_t, 0.01, 0.99))

        # Determine whether to refit MAPIE
        cur_year = pd.Timestamp(d).year
        should_refit = (
            mapie_model is None
            or refit_freq == "monthly"
            or (refit_freq == "yearly" and cur_year != prev_year)
        )

        if should_refit:
            # Rolling training window: last window_length months before d
            train_window_dates = train_dates[-window_length:]
            train_mask = panel["date"].isin(train_window_dates)
            X_tr = panel.loc[train_mask, feat_cols].to_numpy(dtype=float)
            y_tr = panel.loc[train_mask, "ret"].to_numpy(dtype=float)

            # Need enough rows for CV splits
            if len(X_tr) < n_conformal_folds * 2:
                prev_year = cur_year
                continue

            wrapper = _SklearnWrapper(copy.deepcopy(model))
            cv = TimeSeriesSplit(n_splits=n_conformal_folds)
            # MAPIE v1.x: CrossConformalRegressor with method="plus" implements
            # jackknife+ / CV+ conformal prediction.
            # Confidence level is set at 1 - alpha (initial); alpha_t is applied
            # dynamically at predict time via the internal _mapie_regressor.predict.
            mapie = CrossConformalRegressor(
                estimator=wrapper,
                confidence_level=1.0 - alpha,
                method="plus",
                cv=cv,
            )
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                mapie.fit_conformalize(X_tr, y_tr)
            mapie_model = mapie
            prev_year = cur_year

        # Predict for current cross-section
        date_mask = panel["date"] == d
        X_te = panel.loc[date_mask, feat_cols].to_numpy(dtype=float)
        permnos = panel.loc[date_mask, "permno"].to_numpy()
        y_realized_arr = panel.loc[date_mask, "ret"].to_numpy(dtype=float)

        if len(X_te) == 0:
            continue

        # Use internal _mapie_regressor to access dynamic alpha_t at predict time.
        # This bypasses CrossConformalRegressor's fixed confidence_level and allows
        # ACI to widen/narrow intervals adaptively without refitting.
        y_pred_arr, y_pis = mapie_model._mapie_regressor.predict(
            X_te, alpha=alpha_now
        )
        lower_arr, upper_arr = _parse_pis(y_pis)
        iw_arr = upper_arr - lower_arr

        # Uncertainty-adjusted signal: penalise wide intervals
        mu_tilde = y_pred_arr - lambda_t * (iw_arr / 2.0)
        mu_baseline = y_pred_arr.copy()

        # Portfolio weights
        w_arr = _signal_to_weights(mu_tilde, weight_mode)
        wb_arr = _signal_to_weights(mu_baseline, weight_mode)

        # Carry state for ACI update at next iteration
        prev_y_realized = y_realized_arr
        prev_lower = lower_arr
        prev_upper = upper_arr

        for permno, wi, wb, iw in zip(permnos, w_arr, wb_arr, iw_arr):
            rows.append(
                {
                    "id": permno,
                    "eom": d,
                    "w": float(wi),
                    "w_baseline": float(wb),
                    "interval_width": float(iw),
                    "alpha_t": float(alpha_now),
                }
            )

    if not rows:
        return pd.DataFrame(columns=_EMPTY_COLS)

    result = pd.DataFrame(rows)

    # Filter to test dates if requested
    if test_only and test_dates:
        result = result[result["eom"].isin(test_dates)].reset_index(drop=True)

    return result
