"""Characteristic preparation for ML prediction models.

Implements the standardization pipeline from Hellum, Jensen, Kelly, and
Pedersen (2025, HJKP), which is used as the CTF standard: cross-sectional
rank normalization to [-0.5, 0.5] with zero-fill for missing values (zero
equals the median in this range).  GKX (2020) used [-1, 1]; pass
``rank_scale=1.0`` to recover that convention.
"""

from __future__ import annotations

import pandas as pd

from eapctf.utils.constants import JKP_153
from eapctf.utils.preprocessing import rank_normalize, winsorize


def char_prep(
    data: pd.DataFrame,
    char_cols: list[str],
    use_jkp_whitelist: bool = True,
    winsorize_q: tuple[float, float] | None = None,
    rank_normalize_chars: bool = True,
    rank_scale: float = 0.5,
    fill_method: str = "zero",
    date_col: str = "date",
    entity_col: str = "permno",
) -> tuple[pd.DataFrame, list[str]]:
    """Prepare characteristics for ML models (HJKP 2025 CTF standard).

    Steps:
    1. Filter char_cols to JKP_153 intersection if use_jkp_whitelist=True
    2. (Optional) Winsorize at winsorize_q bounds — skip by default since
       rank normalization is outlier-robust
    3. Rank-normalize to [-rank_scale, +rank_scale] cross-sectionally
       (default [-0.5, 0.5] per HJKP 2025 Section 3.3; use rank_scale=1.0
       for the GKX 2020 [-1, 1] convention)
    4. Fill missing: zero-fill (default — zero equals cross-sectional median
       in [-0.5, 0.5] range, matching HJKP 2025 imputation rule)

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel with date, entity, and characteristic columns.
    char_cols : list[str]
        Candidate characteristic columns.
    use_jkp_whitelist : bool
        If True, restrict to columns that appear in JKP_153.
    winsorize_q : tuple[float, float] | None
        Lower and upper quantile bounds for winsorization before ranking.
        Default None skips winsorization (redundant when rank_normalize_chars=True).
    rank_normalize_chars : bool
        If True, apply cross-sectional rank normalization.
    rank_scale : float
        Half-range of the rank normalization interval.  Default 0.5 produces
        [-0.5, 0.5] (HJKP 2025).  Use 1.0 for GKX 2020 convention [-1, 1].
    fill_method : str
        Missing value strategy: ``"zero"`` (default — equals median under
        [-0.5, 0.5] scaling), ``"median"`` (cross-sectional median fill), or
        ``"drop"`` (drop rows with any NaN).
    date_col : str
        Date column name.
    entity_col : str
        Entity identifier column name.

    Returns
    -------
    tuple[pd.DataFrame, list[str]]
        (prepared_data, used_char_cols) where used_char_cols is the final
        list of characteristic columns after JKP filtering.

    Raises
    ------
    ValueError
        If fill_method is not one of the supported values, or no columns
        survive the JKP whitelist filter.

    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 fill_method up front
    valid_methods = {"zero", "median", "drop"}
    if fill_method not in valid_methods:
        raise ValueError(
            f"fill_method must be one of {valid_methods}, got '{fill_method}'"
        )

    # Step 1: filter to JKP whitelist if requested
    if use_jkp_whitelist:
        jkp_set = set(JKP_153)
        used_cols = [c for c in char_cols if c in jkp_set]
    else:
        used_cols = list(char_cols)

    # Also filter to columns actually present in the data
    used_cols = [c for c in used_cols if c in data.columns]

    if len(used_cols) == 0:
        raise ValueError(
            "No characteristic columns remaining after filtering. "
            "Check that char_cols overlap with data columns"
            + (" and JKP_153" if use_jkp_whitelist else "")
            + "."
        )

    # Work on a copy to avoid mutating the input
    result = data.copy()

    # Step 2: (optional) winsorize before ranking
    if winsorize_q is not None:
        result = winsorize(
            result,
            cols=used_cols,
            q_low=winsorize_q[0],
            q_high=winsorize_q[1],
            by_date=True,
            date_col=date_col,
        )

    # Step 3: cross-sectional rank normalization to [-rank_scale, +rank_scale]
    if rank_normalize_chars:
        result = rank_normalize(
            result,
            char_cols=used_cols,
            entity_col=entity_col,
            date_col=date_col,
            scale=rank_scale,
        )

    # Step 4: fill missing values according to chosen strategy
    if fill_method == "zero":
        # GKX default: replace NaN with zero (neutral signal)
        result[used_cols] = result[used_cols].fillna(0.0)
    elif fill_method == "median":
        # Cross-sectional median fill: each date gets its own median
        for col in used_cols:
            median_vals = result.groupby(date_col)[col].transform("median")
            result[col] = result[col].fillna(median_vals)
        # If entire cross-section is NaN, median is NaN; fall back to zero
        result[used_cols] = result[used_cols].fillna(0.0)
    elif fill_method == "drop":
        # Drop any row that still has NaN in characteristic columns
        result = result.dropna(subset=used_cols).reset_index(drop=True)

    return result, used_cols
