"""CTF-specific feature preprocessing utilities.

Rank normalization and feature selection helpers for preparing characteristic
panels before model training within the Common Task Framework (CTF).
"""

from __future__ import annotations

import numpy as np
import pandas as pd
from scipy.stats import rankdata

from eapctf.utils.constants import JKP_153


def ctf_rank_normalize(
    chars: pd.DataFrame,
    features: pd.DataFrame,
    scale: float = 0.5,
    fill_value: float = 0.0,
    eom_col: str = "eom",
) -> pd.DataFrame:
    """Rank-normalize CTF feature columns to [-scale, +scale] via ECDF.

    For each end-of-month cross-section, each feature is ranked using the
    ECDF (average-tie method) and mapped to the interval [-scale, +scale]:

        normalized = (rank / (n + 1) - 0.5) * 2 * scale

    where n is the number of non-missing observations in the cross-section.
    Missing values are filled with ``fill_value`` after normalization, which
    equals the cross-sectional median (0.0) under default scale=0.5, matching
    the HJKP (2025) imputation rule.

    Non-feature columns (e.g., ``id``, ``eom``, ``ret_exc_lead1m``,
    ``ctff_test``) are returned unchanged.

    Parameters
    ----------
    chars : pd.DataFrame
        Full characteristics panel (firms x dates, long format). Must contain
        ``eom_col`` and all feature columns listed in ``features``.
    features : pd.DataFrame
        DataFrame with a column named ``"features"`` (or its first column)
        listing the feature names to normalize.
    scale : float
        Half-range of the output interval. Default 0.5 maps to [-0.5, 0.5]
        as used in HJKP (2025). Use 1.0 for [-1, 1] (GKX 2020 convention).
    fill_value : float
        Value used to fill remaining NaN after normalization. Default 0.0
        corresponds to the cross-sectional median under scale=0.5 (HJKP 2025).
    eom_col : str
        Name of the end-of-month date column. Default ``"eom"``.

    Returns
    -------
    pd.DataFrame
        Copy of ``chars`` with feature columns replaced by normalized values.
        All other columns are preserved without modification.

    References
    ----------
    Hellum, Jensen, Kelly, and Pedersen (2025). The Power of the Common Task
        Framework. SSRN 5242901. Section 3.3.
    """
    # Resolve feature column names from the features DataFrame
    feat_col = "features" if "features" in features.columns else features.columns[0]
    feat_names: list[str] = features[feat_col].tolist()

    # Only process feature columns that actually exist in chars
    feat_present = [f for f in feat_names if f in chars.columns]

    result = chars.copy()

    def _rank_normalize_group(group: pd.Series) -> pd.Series:
        """Rank a single cross-section and map to [-scale, +scale]."""
        mask = group.notna()
        if mask.sum() == 0:
            return group  # all-NaN cross-section, leave unchanged

        out = group.copy()
        n = int(mask.sum())
        # rankdata uses average tie-breaking; slice only non-NaN values
        ranks = rankdata(group[mask].to_numpy(), method="average")
        # ECDF mapping: rank / (n + 1) in (0, 1), then center and scale
        out[mask] = (ranks / (n + 1) - 0.5) * 2.0 * scale
        return out

    # Apply cross-sectional rank normalization for each feature
    for feat in feat_present:
        result[feat] = result.groupby(eom_col, group_keys=False)[feat].apply(
            _rank_normalize_group
        )

    # Fill remaining NaN (stocks with missing feature) with fill_value
    result[feat_present] = result[feat_present].fillna(fill_value)

    return result


def ctf_select_features(
    chars: pd.DataFrame,
    features: pd.DataFrame,
    subset: str = "jkp153",
    eom_col: str = "eom",
    min_coverage: float = 0.5,
) -> tuple[pd.DataFrame, list[str]]:
    """Select a subset of CTF features and return a filtered characteristics panel.

    Three selection modes are supported:

    - ``"jkp153"``: Restrict to the 153 characteristics from Jensen, Kelly, and
      Pedersen (2023, JF). Only columns present in ``chars`` are retained.
    - ``"all"``: Return every feature listed in ``features`` that exists in
      ``chars``.
    - ``"available"``: Retain features whose non-null fraction across all
      observations meets or exceeds ``min_coverage``.

    Non-feature columns (those not listed in ``features``) are always included
    in the returned DataFrame regardless of selection mode.

    Parameters
    ----------
    chars : pd.DataFrame
        Full characteristics panel (firms x dates, long format).
    features : pd.DataFrame
        DataFrame with a column named ``"features"`` (or its first column)
        listing candidate feature names.
    subset : str
        Feature selection mode. One of ``"jkp153"``, ``"all"``,
        ``"available"``. Default ``"jkp153"``.
    eom_col : str
        Name of the end-of-month date column. Retained in output. Default
        ``"eom"``.
    min_coverage : float
        Minimum fraction of non-null observations required to keep a feature
        when ``subset="available"``. Default 0.5 (50% coverage).

    Returns
    -------
    tuple[pd.DataFrame, list[str]]
        A 2-tuple of:

        - Filtered ``chars`` DataFrame containing non-feature columns plus
          the selected feature columns.
        - List of selected feature column names.

    Raises
    ------
    ValueError
        If ``subset`` is not one of ``"jkp153"``, ``"all"``, ``"available"``.

    References
    ----------
    Jensen, T., Kelly, B., & Pedersen, L. H. (2023). Is There a Replication
        Crisis in Finance? Journal of Finance, 78(5), 2465-2518.
    """
    # Resolve feature column names from the features DataFrame
    feat_col = "features" if "features" in features.columns else features.columns[0]
    all_feat_names: list[str] = features[feat_col].tolist()

    # Identify non-feature columns that should always be kept
    feat_set = set(all_feat_names)
    non_feat_cols: list[str] = [c for c in chars.columns if c not in feat_set]

    # Features that actually exist in chars (universe for all modes)
    feat_in_chars: list[str] = [f for f in all_feat_names if f in chars.columns]

    if subset == "jkp153":
        jkp_set = set(JKP_153)
        selected = [f for f in feat_in_chars if f in jkp_set]

    elif subset == "all":
        selected = feat_in_chars

    elif subset == "available":
        n_total = len(chars)
        if n_total == 0:
            selected = []
        else:
            # Compute non-null fraction for each candidate feature
            coverage: np.ndarray = chars[feat_in_chars].notna().mean(axis=0).to_numpy()
            selected = [f for f, cov in zip(feat_in_chars, coverage) if cov >= min_coverage]

    else:
        raise ValueError(
            f"subset must be one of 'jkp153', 'all', 'available'; got {subset!r}"
        )

    # Reconstruct filtered DataFrame: non-feature cols + selected features
    keep_cols = non_feat_cols + selected
    filtered = chars[keep_cols].copy()

    return filtered, selected
