"""Preprocessing utilities for cross-sectional characteristic transformations.

Implements rank normalization, winsorization, standardization, and portfolio
classification following standard empirical asset pricing conventions.
"""

from __future__ import annotations

import numpy as np
import pandas as pd


def rank_normalize(
    data: pd.DataFrame,
    char_cols: list[str],
    entity_col: str = "permno",
    date_col: str = "date",
    scale: float = 0.5,
) -> pd.DataFrame:
    """Cross-sectional rank normalization to [-scale, +scale].

    Each characteristic is ranked cross-sectionally within each period,
    then mapped to [-scale, +scale] via the transformation:

        x_ranked = 2 * scale * rank / (N + 1) - scale

    where rank is 1-based and N is the number of non-missing observations
    in the cross-section for that period.

    Default scale=0.5 produces the [-0.5, 0.5] range used by
    Hellum, Jensen, Kelly, and Pedersen (2025, "The Power of the Common Task
    Framework"), Section 3.3.  With scale=1.0 the formula recovers the
    [-1, 1] convention from Gu, Kelly, and Xiu (2020, GKX) Appendix A.1.

    When missing values are subsequently zero-filled (the default in
    char_prep), zero equals the cross-sectional median under [-0.5, 0.5]
    scaling, matching the HJKP imputation rule.

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel with date, entity, and characteristic columns.
    char_cols : list[str]
        Columns to rank-normalize.
    entity_col : str
        Entity identifier column (unused directly but preserved in output).
    date_col : str
        Date column for cross-sectional grouping.
    scale : float
        Half-range of the output interval.  Default 0.5 (HJKP 2025).
        Use 1.0 for GKX 2020 convention.

    Returns
    -------
    pd.DataFrame
        Copy of input with char_cols replaced by their rank-normalized values.

    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. Appendix A.1.
    """
    result = data.copy()

    # Process each characteristic column independently
    for col in char_cols:
        # Rank within each cross-section (date group), average ties
        # rank() returns 1-based ranks; NaN inputs remain NaN
        ranked = result.groupby(date_col)[col].rank(method="average")

        # Count non-missing obs per date for normalization denominator
        count = result.groupby(date_col)[col].transform("count")

        # Map ranks to [-scale, +scale]: 2*scale * rank/(N+1) - scale
        # When rank=1 this gives ~ -scale, when rank=N this gives ~ +scale
        normalized = 2.0 * scale * ranked / (count + 1.0) - scale

        result[col] = normalized

    return result


def winsorize(
    data: pd.DataFrame,
    cols: list[str],
    q_low: float = 0.01,
    q_high: float = 0.99,
    by_date: bool = True,
    date_col: str = "date",
) -> pd.DataFrame:
    """Winsorize columns at specified percentiles.

    Values below the q_low quantile are set to the q_low cutoff,
    and values above q_high are set to the q_high cutoff.

    Parameters
    ----------
    data : pd.DataFrame
        Input DataFrame (long-format panel if by_date=True).
    cols : list[str]
        Columns to winsorize.
    q_low : float
        Lower quantile threshold (default 1st percentile).
    q_high : float
        Upper quantile threshold (default 99th percentile).
    by_date : bool
        If True, compute quantile cutoffs within each date cross-section.
        If False, compute cutoffs over the full sample.

    Returns
    -------
    pd.DataFrame
        Copy of input with winsorized columns.
    """
    result = data.copy()

    if by_date:
        # Winsorize within each date group to avoid time-series contamination
        for col in cols:
            # Compute lower and upper bounds per date
            lower = result.groupby(date_col)[col].transform(
                lambda x: x.quantile(q_low)
            )
            upper = result.groupby(date_col)[col].transform(
                lambda x: x.quantile(q_high)
            )
            result[col] = result[col].clip(lower=lower, upper=upper)
    else:
        # Winsorize over the full pooled sample
        for col in cols:
            lower_val = result[col].quantile(q_low)
            upper_val = result[col].quantile(q_high)
            result[col] = result[col].clip(lower=lower_val, upper=upper_val)

    return result


def standardize(
    data: pd.DataFrame,
    cols: list[str],
    by_date: bool = True,
    date_col: str = "date",
) -> pd.DataFrame:
    """Demean and scale to unit variance.

    Parameters
    ----------
    data : pd.DataFrame
        Input DataFrame (long-format panel if by_date=True).
    cols : list[str]
        Columns to standardize.
    by_date : bool
        If True, compute mean and std within each date cross-section.
        If False, compute over the full sample.

    Returns
    -------
    pd.DataFrame
        Copy of input with standardized columns.
    """
    result = data.copy()

    if by_date:
        for col in cols:
            # Cross-sectional demeaning and scaling
            mean = result.groupby(date_col)[col].transform("mean")
            std = result.groupby(date_col)[col].transform("std")
            # Avoid division by zero when std is zero (constant cross-section)
            result[col] = (result[col] - mean) / std.replace(0.0, np.nan)
    else:
        for col in cols:
            mean_val = result[col].mean()
            std_val = result[col].std()
            if std_val == 0.0:
                result[col] = np.nan
            else:
                result[col] = (result[col] - mean_val) / std_val

    return result


def classify(
    data: pd.DataFrame,
    col: str,
    n_portfolios: int | None = None,
    splits: list[float] | None = None,
    ascending: bool = True,
    by_nyse: bool = True,
    exchcd_col: str = "exchcd",
    date_col: str = "date",
) -> pd.Series:
    """Assign integer portfolio labels based on quantile breakpoints.

    Computes breakpoints either from equal-spaced quantiles (n_portfolios)
    or from user-specified quantile splits, then assigns each observation
    to a portfolio bucket. When by_nyse=True, breakpoints are computed
    on NYSE stocks only (exchcd==1) but applied to all stocks.

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel with at least col, date_col columns.
    col : str
        Characteristic column used for sorting.
    n_portfolios : int | None
        Number of equal-spaced portfolios. Mutually exclusive with splits.
    splits : list[float] | None
        Custom quantile breakpoints (e.g. [0.3, 0.7] for terciles).
        Mutually exclusive with n_portfolios.
    ascending : bool
        If True, portfolio 0 has the lowest values. If False, reversed.
    by_nyse : bool
        If True, compute breakpoints on NYSE stocks (exchcd==1) only,
        then apply to all stocks.
    exchcd_col : str
        Exchange code column (CRSP convention: 1 = NYSE).
    date_col : str
        Date column for cross-sectional grouping.

    Returns
    -------
    pd.Series
        Integer portfolio labels (0-based), same index as input.

    Raises
    ------
    ValueError
        If neither n_portfolios nor splits is provided, or both are.

    References
    ----------
    Equivalent to pyanomaly.datatools.classify() with NYSE breakpoints.
    """
    # Validate inputs: exactly one of n_portfolios or splits must be given
    if (n_portfolios is None) == (splits is None):
        raise ValueError("Provide exactly one of n_portfolios or splits.")

    # Build quantile boundaries from n_portfolios if not user-specified
    if n_portfolios is not None:
        boundaries = [i / n_portfolios for i in range(1, n_portfolios)]
    else:
        assert splits is not None  # for type checker
        boundaries = sorted(splits)

    # Initialize result series with NaN (int-compatible after fillna)
    result = pd.Series(np.nan, index=data.index, dtype="float64")

    # Process each date cross-section independently
    for date_val, group in data.groupby(date_col):
        # Determine which subset to use for computing breakpoints
        if by_nyse and exchcd_col in group.columns:
            bp_data = group.loc[group[exchcd_col] == 1, col].dropna()
        else:
            bp_data = group[col].dropna()

        # Skip dates with no valid breakpoint data
        if len(bp_data) == 0:
            continue

        # Compute breakpoint values at specified quantiles
        bp_values = [bp_data.quantile(q) for q in boundaries]

        # Assign portfolio labels using np.searchsorted
        # searchsorted returns the bin index for each value
        vals = group[col]
        bp_array = np.array(sorted(bp_values), dtype=np.float64)
        labels = np.searchsorted(bp_array, vals.to_numpy(dtype=np.float64), side="right")

        # Reverse labeling if descending sort is requested
        if not ascending:
            labels = len(boundaries) - labels

        # Assign labels; keep NaN where original values are missing
        label_series = pd.Series(labels, index=group.index, dtype="float64")
        label_series[vals.isna()] = np.nan
        result.loc[group.index] = label_series

    return result
