"""Factor portfolio construction.

General engine for characteristic-based factor portfolios, with specific
implementations for Fama-French 3 and 5 factor models, Hou-Xue-Zhang
q-factor model, Stambaugh-Yuan mispricing factors, and momentum.

References
----------
Fama, E. F., & French, K. R. (1993). Common risk factors in the returns
    on stocks and bonds. Journal of Financial Economics, 33(1), 3-56.
Fama, E. F., & French, K. R. (2015). A five-factor asset pricing model.
    Journal of Financial Economics, 116(1), 1-22.
Hou, K., Xue, C., & Zhang, L. (2015). Digesting anomalies: An investment
    approach. Review of Financial Studies, 28(3), 650-705.
Stambaugh, R. F., & Yuan, Y. (2017). Mispricing factors.
    Review of Financial Studies, 30(4), 1270-1315.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np
import pandas as pd


@dataclass
class FactorResult:
    """Container for factor construction results.

    Attributes
    ----------
    factors : pd.DataFrame
        Date-indexed DataFrame with one column per factor.
    portfolios : dict[str, pd.DataFrame]
        Underlying portfolio returns used to construct the factors.
    """

    factors: pd.DataFrame
    portfolios: dict[str, pd.DataFrame]


# ---------------------------------------------------------------------------
# Internal helpers -- inlined to avoid importing eapctf.utils during development
# ---------------------------------------------------------------------------


def _classify_by_breakpoints(
    values: pd.Series[Any],
    breakpoints: list[float],
) -> pd.Series[Any]:
    """Assign portfolio labels (1-based) given a list of breakpoint values.

    Values at or below the first breakpoint get label 1, between bp[0] and
    bp[1] get label 2, and so on. Values above the last breakpoint get the
    highest label.
    """
    labels = pd.Series(np.nan, index=values.index)
    n_bins = len(breakpoints) + 1
    for idx in values.index:
        val = values[idx]
        if pd.isna(val):
            continue
        assigned = n_bins  # default: above all breakpoints
        for i, bp in enumerate(breakpoints):
            if val <= bp:
                assigned = i + 1
                break
        labels[idx] = assigned
    return labels


def _classify_with_splits(
    data: pd.DataFrame,
    char_col: str,
    splits: tuple[float, ...],
    date_col: str,
    ascending: bool = True,
    nyse_breakpoints_flag: bool = True,
    exchcd_col: str = "exchcd",
) -> pd.Series[Any]:
    """Classify stocks into groups using percentage splits per date.

    For example, splits=(30, 40, 30) creates 3 groups at the 30th and
    70th percentiles.

    Parameters
    ----------
    ascending : bool
        If True, low-char stocks get label 1 (e.g., small stocks = 1).
        If False, high-char stocks get label 1.
    """
    # Validate splits sum to 100 (within floating-point tolerance)
    splits_sum = sum(splits)
    if abs(splits_sum - 100.0) > 1e-6:
        raise ValueError(
            f"splits must sum to 100, got {splits_sum}. "
            f"Example: (30, 40, 30) for a 30/40/30 split."
        )

    # Convert splits to cumulative percentile breakpoints
    # e.g., (30, 40, 30) -> [0.3, 0.7]
    cumulative = np.cumsum(splits[:-1]) / np.sum(splits)
    n_groups = len(splits)

    all_labels: list[pd.Series[Any]] = []
    for _dt, group in data.groupby(date_col, sort=False):
        if nyse_breakpoints_flag and exchcd_col in group.columns:
            # Compute breakpoints from NYSE stocks only
            nyse_vals = group.loc[group[exchcd_col] == 1, char_col].dropna()
        else:
            nyse_vals = group[char_col].dropna()

        if len(nyse_vals) == 0:
            lbl: pd.Series[Any] = pd.Series(np.nan, index=group.index)
            all_labels.append(lbl)
            continue

        # Compute breakpoint values at the cumulative percentiles
        bp_values = [float(np.quantile(nyse_vals, q)) for q in cumulative]

        # Assign labels to ALL stocks using these breakpoints
        lbl = _classify_by_breakpoints(group[char_col], bp_values)

        if not ascending:
            # Reverse labels so that high-char = label 1
            def _reverse_label(x: float, ng: int = n_groups) -> float:
                return float(ng - x + 1) if pd.notna(x) else float(np.nan)
            lbl = lbl.map(_reverse_label)

        all_labels.append(lbl)

    if all_labels:
        result: pd.Series[Any] = pd.concat(all_labels)
        return result
    return pd.Series(dtype=float)


def _compute_vw_portfolio_returns(
    data: pd.DataFrame,
    port_col: str,
    ret_col: str,
    date_col: str,
    mktcap_col: str,
    weight_col: str | None = None,
) -> pd.DataFrame:
    """Compute value-weighted returns per (date, portfolio) group.

    Returns a date x portfolio wide DataFrame.
    """
    w_col = weight_col if weight_col is not None else mktcap_col
    df = data.dropna(subset=[port_col, ret_col, w_col]).copy()

    df["_wret"] = df[ret_col] * df[w_col]
    grouped = df.groupby([date_col, port_col], sort=True)
    port_ret_series = grouped["_wret"].sum() / grouped[w_col].sum()
    port_ret_df = port_ret_series.reset_index()
    port_ret_df.columns = pd.Index([date_col, port_col, "ret"])
    wide = pd.DataFrame(
        port_ret_df.pivot(index=date_col, columns=port_col, values="ret"),
    )
    return wide


def char_factor(
    data: pd.DataFrame,
    char_col: str,
    ret_col: str = "ret",
    date_col: str = "date",
    entity_col: str = "permno",
    mktcap_col: str = "mktcap",
    exchcd_col: str = "exchcd",
    char_split: tuple[float, ...] = (30, 40, 30),
    nyse_breakpoints: bool = True,
    ascending: bool = True,
    size_col: str | None = None,
    weight_col: str | None = None,
    primary_col: str | None = None,
) -> FactorResult:
    """General characteristic-based factor portfolio.

    Constructs a long-short factor portfolio from any firm characteristic.
    If ``size_col`` is provided, performs a 2x3 double sort (size x char)
    and returns both the characteristic factor (HML-style) and a size
    factor (SMB-style).

    Parameters
    ----------
    data : pd.DataFrame
        Panel data.
    char_col : str
        Characteristic column to sort on.
    ret_col : str, default "ret"
        Return column.
    date_col : str, default "date"
        Date column.
    entity_col : str, default "permno"
        Entity identifier.
    mktcap_col : str, default "mktcap"
        Market cap column for value-weighting.
    exchcd_col : str, default "exchcd"
        Exchange code column.
    char_split : tuple[float, ...], default (30, 40, 30)
        Percentage split for the characteristic sort (e.g., 30/40/30 tercile).
    nyse_breakpoints : bool, default True
        Use NYSE-only breakpoints.
    ascending : bool, default True
        If True, low-char stocks are in group 1 (standard for book-to-market).
    size_col : str or None, default None
        If provided, perform a 2x3 double sort (size median x char tercile).
    weight_col : str or None, default None
        Alternative weighting column (defaults to mktcap_col).
    primary_col : str or None, default None
        Column flagging primary securities. Rows where this column is False
        or 0 are dropped before portfolio construction, excluding ADRs,
        preferred shares, and other non-primary listings (pyanomaly
        convention).

    Returns
    -------
    FactorResult
        factors: DataFrame with the factor return series.
        portfolios: dict mapping portfolio names to return DataFrames.

    References
    ----------
    Fama, E. F., & French, K. R. (1993). Common risk factors in the returns
        on stocks and bonds. Journal of Financial Economics, 33(1), 3-56.
        Section II (factor construction procedure).
    """
    df = data.copy()
    # Apply primary securities filter (pyanomaly convention)
    if primary_col is not None:
        df = df[df[primary_col].astype(bool)].copy()
    n_char_groups = len(char_split)

    # Classify stocks by characteristic (tercile by default)
    df["_char_group"] = _classify_with_splits(
        df, char_col, char_split, date_col,
        ascending=ascending,
        nyse_breakpoints_flag=nyse_breakpoints,
        exchcd_col=exchcd_col,
    )

    if size_col is not None:
        # 2x3 double sort: size (median) x char (tercile)
        # Classify by size: median split -> 2 groups
        df["_size_group"] = _classify_with_splits(
            df, size_col, (50, 50), date_col,
            ascending=True,  # small = 1, big = 2
            nyse_breakpoints_flag=nyse_breakpoints,
            exchcd_col=exchcd_col,
        )

        # Drop stocks without valid group assignments
        valid = df["_char_group"].notna() & df["_size_group"].notna()
        df = df[valid].copy()

        # Compute value-weighted returns for each of the 6 portfolios (2 size x 3 char)
        df["_port_label"] = (
            df["_size_group"].astype(int).astype(str)
            + "_"
            + df["_char_group"].astype(int).astype(str)
        )
        wide = _compute_vw_portfolio_returns(
            df, "_port_label", ret_col, date_col, mktcap_col, weight_col,
        )

        # Label portfolios: size_1 = small, size_2 = big
        # char_1 = low, char_2 = mid, char_3 = high (if ascending=True)
        # HML-style factor: (small_high + big_high)/2 - (small_low + big_low)/2
        _zero: Any = pd.Series(0.0, index=wide.index)
        small_high: pd.Series[Any] = wide.get("1_3", _zero)
        big_high: pd.Series[Any] = wide.get("2_3", _zero)
        small_low: pd.Series[Any] = wide.get("1_1", _zero)
        big_low: pd.Series[Any] = wide.get("2_1", _zero)

        char_factor_ret = (small_high + big_high) / 2 - (small_low + big_low) / 2

        # SMB-style factor: avg(small) - avg(big)
        small_cols = [c for c in wide.columns if str(c).startswith("1_")]
        big_cols = [c for c in wide.columns if str(c).startswith("2_")]
        smb_ret = wide[small_cols].mean(axis=1) - wide[big_cols].mean(axis=1)

        # Build factors DataFrame
        factors = pd.DataFrame(
            {
                f"{char_col}_factor": char_factor_ret,
                "smb": smb_ret,
            },
            index=wide.index,
        )
        factors.index.name = date_col

        # Build portfolios dict
        portfolios: dict[str, pd.DataFrame] = {}
        char_labels_map: dict[str, str] = {"1": "low", "2": "mid", "3": "high"}
        for col in wide.columns:
            col_str = str(col)
            size_label = "small" if col_str.startswith("1_") else "big"
            char_idx = col_str.split("_")[1]
            char_label = char_labels_map.get(char_idx, char_idx)
            port_name = f"{size_label}_{char_label}"
            portfolios[port_name] = wide[[col]].rename(columns={col: "ret"})

    else:
        # Simple univariate sort (no size conditioning)
        valid = df["_char_group"].notna()
        df = df[valid].copy()

        wide = _compute_vw_portfolio_returns(
            df, "_char_group", ret_col, date_col, mktcap_col, weight_col,
        )

        # Long-short: high minus low
        high_label = float(n_char_groups)
        low_label = 1.0
        _zero_univ: Any = pd.Series(0.0, index=wide.index)
        high_ret: pd.Series[Any] = wide.get(high_label, _zero_univ)
        low_ret: pd.Series[Any] = wide.get(low_label, _zero_univ)
        char_factor_ret = high_ret - low_ret

        factors = pd.DataFrame(
            {f"{char_col}_factor": char_factor_ret},
            index=wide.index,
        )
        factors.index.name = date_col

        # Build portfolios dict
        portfolios = {}
        label_names: dict[float, str] = {1.0: "low", 2.0: "mid", 3.0: "high"}
        for col in wide.columns:
            name = label_names.get(float(col), f"group_{int(col)}")
            portfolios[name] = wide[[col]].rename(columns={col: "ret"})

    return FactorResult(factors=factors, portfolios=portfolios)


def _compute_market_return(
    data: pd.DataFrame,
    ret_col: str,
    date_col: str,
    mktcap_col: str,
    rf_col: str,
    primary_col: str | None = None,
) -> pd.Series[Any]:
    """Compute value-weighted market return minus risk-free rate."""
    df = data.dropna(subset=[ret_col, mktcap_col]).copy()
    if primary_col is not None:
        df = df[df[primary_col].astype(bool)]
    df["_wret"] = df[ret_col] * df[mktcap_col]
    grouped = df.groupby(date_col, sort=True)
    mkt_ret = grouped["_wret"].sum() / grouped[mktcap_col].sum()

    # Get risk-free rate (one per date)
    rf = df.groupby(date_col, sort=True)[rf_col].first()

    mkt_rf = mkt_ret - rf
    mkt_rf.name = "mkt_rf"
    return mkt_rf


def ff3_factors(
    data: pd.DataFrame,
    ret_col: str = "ret",
    date_col: str = "date",
    entity_col: str = "permno",
    mktcap_col: str = "mktcap",
    bm_col: str = "bm",
    rf_col: str = "rf",
    primary_col: str | None = None,
) -> FactorResult:
    """Fama-French 3-factor model: MKT-RF, SMB, HML.

    Constructs the three factors following the standard 2x3 sorting
    procedure on size and book-to-market.

    Parameters
    ----------
    data : pd.DataFrame
        Panel data with return, market cap, book-to-market, and risk-free rate.
    ret_col : str, default "ret"
        Return column.
    date_col : str, default "date"
        Date column.
    entity_col : str, default "permno"
        Entity identifier.
    mktcap_col : str, default "mktcap"
        Market capitalization column.
    bm_col : str, default "bm"
        Book-to-market ratio column.
    rf_col : str, default "rf"
        Risk-free rate column.
    primary_col : str or None, default None
        Column flagging primary securities; passed to char_factor and market
        return computation to exclude non-primary listings.

    Returns
    -------
    FactorResult
        factors DataFrame has columns [mkt_rf, smb, hml].

    References
    ----------
    Fama, E. F., & French, K. R. (1993). Common risk factors in the returns
        on stocks and bonds. Journal of Financial Economics, 33(1), 3-56.
    """
    # HML via 2x3 sort on size x BM (ascending=True: low BM = 1, high BM = 3)
    hml_result = char_factor(
        data, bm_col,
        ret_col=ret_col, date_col=date_col, entity_col=entity_col,
        mktcap_col=mktcap_col, size_col=mktcap_col,
        ascending=True,
        primary_col=primary_col,
    )

    # Extract HML and SMB from the double-sort result
    hml = hml_result.factors[f"{bm_col}_factor"]
    smb = hml_result.factors["smb"]

    # MKT-RF: value-weighted market return minus risk-free rate
    mkt_rf = _compute_market_return(
        data, ret_col, date_col, mktcap_col, rf_col, primary_col=primary_col
    )

    # Align all factors on a common date index
    factors = pd.DataFrame(
        {
            "mkt_rf": mkt_rf,
            "smb": smb,
            "hml": hml,
        },
    )
    factors.index.name = date_col

    # Combine all underlying portfolios
    all_portfolios = dict(hml_result.portfolios)

    return FactorResult(factors=factors, portfolios=all_portfolios)


def ff5_factors(
    data: pd.DataFrame,
    ret_col: str = "ret",
    date_col: str = "date",
    entity_col: str = "permno",
    mktcap_col: str = "mktcap",
    bm_col: str = "bm",
    op_col: str = "op",
    inv_col: str = "inv",
    rf_col: str = "rf",
    primary_col: str | None = None,
) -> FactorResult:
    """Fama-French 5-factor model: MKT-RF, SMB, HML, RMW, CMA.

    Extends the 3-factor model with profitability (RMW) and investment (CMA).

    Parameters
    ----------
    data : pd.DataFrame
        Panel data.
    ret_col, date_col, entity_col, mktcap_col : str
        Standard column names.
    bm_col : str, default "bm"
        Book-to-market ratio column.
    op_col : str, default "op"
        Operating profitability column.
    inv_col : str, default "inv"
        Investment (asset growth) column.
    rf_col : str, default "rf"
        Risk-free rate column.
    primary_col : str or None, default None
        Column flagging primary securities; passed to all char_factor and
        market return calls to exclude non-primary listings.

    Returns
    -------
    FactorResult
        factors DataFrame has columns [mkt_rf, smb, hml, rmw, cma].

    References
    ----------
    Fama, E. F., & French, K. R. (2015). A five-factor asset pricing model.
        Journal of Financial Economics, 116(1), 1-22.
    """
    # HML: 2x3 on size x BM
    hml_result = char_factor(
        data, bm_col,
        ret_col=ret_col, date_col=date_col, entity_col=entity_col,
        mktcap_col=mktcap_col, size_col=mktcap_col,
        ascending=True,
        primary_col=primary_col,
    )

    # RMW: 2x3 on size x operating profitability (high profit = robust)
    rmw_result = char_factor(
        data, op_col,
        ret_col=ret_col, date_col=date_col, entity_col=entity_col,
        mktcap_col=mktcap_col, size_col=mktcap_col,
        ascending=True,
        primary_col=primary_col,
    )

    # CMA: 2x3 on size x investment (low investment = conservative)
    # ascending=False flips so that group 3 = conservative, group 1 = aggressive
    # HML-style = group_3 - group_1 = conservative - aggressive = CMA
    cma_result = char_factor(
        data, inv_col,
        ret_col=ret_col, date_col=date_col, entity_col=entity_col,
        mktcap_col=mktcap_col, size_col=mktcap_col,
        ascending=False,
        primary_col=primary_col,
    )

    # Combine SMB from all three sorts (FF2015 averages SMB across sorts)
    smb = (
        hml_result.factors["smb"]
        + rmw_result.factors["smb"]
        + cma_result.factors["smb"]
    ) / 3

    hml = hml_result.factors[f"{bm_col}_factor"]
    rmw = rmw_result.factors[f"{op_col}_factor"]
    cma = cma_result.factors[f"{inv_col}_factor"]

    # MKT-RF
    mkt_rf = _compute_market_return(
        data, ret_col, date_col, mktcap_col, rf_col, primary_col=primary_col
    )

    factors = pd.DataFrame(
        {
            "mkt_rf": mkt_rf,
            "smb": smb,
            "hml": hml,
            "rmw": rmw,
            "cma": cma,
        },
    )
    factors.index.name = date_col

    # Merge all portfolios
    all_portfolios: dict[str, pd.DataFrame] = {}
    for prefix, result in [("hml", hml_result), ("rmw", rmw_result), ("cma", cma_result)]:
        for name, pf in result.portfolios.items():
            all_portfolios[f"{prefix}_{name}"] = pf

    return FactorResult(factors=factors, portfolios=all_portfolios)


def hxz4_factors(
    data: pd.DataFrame,
    ret_col: str = "ret",
    date_col: str = "date",
    entity_col: str = "permno",
    mktcap_col: str = "mktcap",
    inv_col: str = "inv",
    roe_col: str = "roe",
    rf_col: str = "rf",
    primary_col: str | None = None,
) -> FactorResult:
    """Hou-Xue-Zhang q-factor model: MKT-RF, ME, IA, ROE.

    Parameters
    ----------
    data : pd.DataFrame
        Panel data.
    ret_col, date_col, entity_col, mktcap_col : str
        Standard column names.
    inv_col : str, default "inv"
        Investment-to-assets column.
    roe_col : str, default "roe"
        Return on equity column.
    rf_col : str, default "rf"
        Risk-free rate column.
    primary_col : str or None, default None
        Column flagging primary securities; forwarded to all char_factor and
        market return calls.

    Returns
    -------
    FactorResult
        factors DataFrame has columns [mkt_rf, me, ia, roe].

    References
    ----------
    Hou, K., Xue, C., & Zhang, L. (2015). Digesting anomalies: An investment
        approach. Review of Financial Studies, 28(3), 650-705.
    """
    # ME (size factor): 2x3 on size x size
    me_result = char_factor(
        data, mktcap_col,
        ret_col=ret_col, date_col=date_col, entity_col=entity_col,
        mktcap_col=mktcap_col, size_col=mktcap_col,
        ascending=True,
        primary_col=primary_col,
    )

    # IA (investment factor): 2x3 on size x investment
    # Low investment minus high investment -> ascending=False
    ia_result = char_factor(
        data, inv_col,
        ret_col=ret_col, date_col=date_col, entity_col=entity_col,
        mktcap_col=mktcap_col, size_col=mktcap_col,
        ascending=False,
        primary_col=primary_col,
    )

    # ROE factor: 2x3 on size x ROE (high ROE minus low ROE)
    roe_result = char_factor(
        data, roe_col,
        ret_col=ret_col, date_col=date_col, entity_col=entity_col,
        mktcap_col=mktcap_col, size_col=mktcap_col,
        ascending=True,
        primary_col=primary_col,
    )

    # MKT-RF
    mkt_rf = _compute_market_return(
        data, ret_col, date_col, mktcap_col, rf_col, primary_col=primary_col
    )

    me = me_result.factors["smb"]
    ia = ia_result.factors[f"{inv_col}_factor"]
    roe = roe_result.factors[f"{roe_col}_factor"]

    factors = pd.DataFrame(
        {
            "mkt_rf": mkt_rf,
            "me": me,
            "ia": ia,
            "roe": roe,
        },
    )
    factors.index.name = date_col

    all_portfolios: dict[str, pd.DataFrame] = {}
    for prefix, result in [("me", me_result), ("ia", ia_result), ("roe", roe_result)]:
        for name, pf in result.portfolios.items():
            all_portfolios[f"{prefix}_{name}"] = pf

    return FactorResult(factors=factors, portfolios=all_portfolios)


def sy4_factors(
    data: pd.DataFrame,
    ret_col: str = "ret",
    date_col: str = "date",
    entity_col: str = "permno",
    mktcap_col: str = "mktcap",
    bm_col: str = "bm",
    mgmt_anomaly_cols: list[str] | None = None,
    perf_anomaly_cols: list[str] | None = None,
    rf_col: str = "rf",
    primary_col: str | None = None,
) -> FactorResult:
    """Stambaugh-Yuan 4-factor model: MKT-RF, SMB, MGMT, PERF.

    MGMT and PERF are constructed by averaging cross-sectional ranks
    of component anomalies into composite scores, then sorting.

    Parameters
    ----------
    data : pd.DataFrame
        Panel data.
    ret_col, date_col, entity_col, mktcap_col : str
        Standard column names.
    bm_col : str, default "bm"
        Book-to-market column (not used directly, but available for SMB).
    mgmt_anomaly_cols : list[str] or None
        Columns for management-related anomalies (e.g., net issuance,
        composite equity issuance, accruals, net operating assets, asset growth).
    perf_anomaly_cols : list[str] or None
        Columns for performance-related anomalies (e.g., financial distress,
        O-score, momentum, gross profitability, ROA).
    rf_col : str, default "rf"
        Risk-free rate column.
    primary_col : str or None, default None
        Column flagging primary securities. Applied before composite ranking
        so ADRs/preferred shares are excluded from rank normalization too.

    Returns
    -------
    FactorResult
        factors DataFrame has columns [mkt_rf, smb, mgmt, perf].

    References
    ----------
    Stambaugh, R. F., & Yuan, Y. (2017). Mispricing factors.
        Review of Financial Studies, 30(4), 1270-1315.
    """
    if mgmt_anomaly_cols is None:
        mgmt_anomaly_cols = []
    if perf_anomaly_cols is None:
        perf_anomaly_cols = []

    df = data.copy()
    # Apply primary securities filter before composite ranking so non-primary
    # stocks do not distort the cross-sectional rank normalization
    if primary_col is not None:
        df = df[df[primary_col].astype(bool)].copy()

    # Build MGMT composite: average cross-sectional percentile rank
    if mgmt_anomaly_cols:
        mgmt_rank_cols: list[str] = []
        for col in mgmt_anomaly_cols:
            rank_col = f"_rank_{col}"
            df[rank_col] = df.groupby(date_col, sort=False)[col].rank(pct=True)
            mgmt_rank_cols.append(rank_col)
        df["_mgmt_composite"] = df[mgmt_rank_cols].mean(axis=1)
    else:
        df["_mgmt_composite"] = np.nan

    # Build PERF composite: average cross-sectional percentile rank
    if perf_anomaly_cols:
        perf_rank_cols: list[str] = []
        for col in perf_anomaly_cols:
            rank_col = f"_rank_{col}"
            df[rank_col] = df.groupby(date_col, sort=False)[col].rank(pct=True)
            perf_rank_cols.append(rank_col)
        df["_perf_composite"] = df[perf_rank_cols].mean(axis=1)
    else:
        df["_perf_composite"] = np.nan

    # MGMT factor: 2x3 on size x mgmt_composite
    mgmt_result = char_factor(
        df, "_mgmt_composite",
        ret_col=ret_col, date_col=date_col, entity_col=entity_col,
        mktcap_col=mktcap_col, size_col=mktcap_col,
        ascending=True,
    )

    # PERF factor: 2x3 on size x perf_composite
    perf_result = char_factor(
        df, "_perf_composite",
        ret_col=ret_col, date_col=date_col, entity_col=entity_col,
        mktcap_col=mktcap_col, size_col=mktcap_col,
        ascending=True,
    )

    # SMB: average of SMB from MGMT and PERF sorts
    smb = (mgmt_result.factors["smb"] + perf_result.factors["smb"]) / 2

    mgmt = mgmt_result.factors["_mgmt_composite_factor"]
    perf = perf_result.factors["_perf_composite_factor"]

    # MKT-RF
    mkt_rf = _compute_market_return(
        data, ret_col, date_col, mktcap_col, rf_col, primary_col=primary_col
    )

    factors = pd.DataFrame(
        {
            "mkt_rf": mkt_rf,
            "smb": smb,
            "mgmt": mgmt,
            "perf": perf,
        },
    )
    factors.index.name = date_col

    all_portfolios: dict[str, pd.DataFrame] = {}
    for prefix, result in [("mgmt", mgmt_result), ("perf", perf_result)]:
        for name, pf in result.portfolios.items():
            all_portfolios[f"{prefix}_{name}"] = pf

    return FactorResult(factors=factors, portfolios=all_portfolios)


def mom_factor(
    data: pd.DataFrame,
    ret_col: str = "ret",
    date_col: str = "date",
    entity_col: str = "permno",
    mktcap_col: str = "mktcap",
    exchcd_col: str = "exchcd",
    formation_start: int = 12,
    formation_end: int = 2,
    primary_col: str | None = None,
) -> FactorResult:
    """Momentum factor UMD.

    Computes cumulative returns over the formation window (months t-lag to
    t-formation_end), then forms a long-short factor portfolio using a 2x3
    (size x momentum) double sort with NYSE breakpoints.

    Formation window: ``shift(formation_start)`` stores the cumulative log
    return through t-formation_start, so the momentum signal captures returns
    from month t-(formation_start-1) through t-formation_end.
    With defaults (formation_start=12, formation_end=2) this gives 10 months
    (t-11 through t-2), the standard Jegadeesh-Titman skip-1-month window.

    Parameters
    ----------
    data : pd.DataFrame
        Panel data with returns.
    ret_col : str, default "ret"
        Return column.
    date_col : str, default "date"
        Date column.
    entity_col : str, default "permno"
        Entity identifier.
    mktcap_col : str, default "mktcap"
        Market capitalization column.
    exchcd_col : str, default "exchcd"
        Exchange code column forwarded to NYSE breakpoint computation.
    formation_start : int, default 12
        Lag at the start of the formation window. Returns from
        t-(formation_start-1) onward are included in the signal.
    formation_end : int, default 2
        Lag at the end of the formation window (inclusive).
    primary_col : str or None, default None
        Column flagging primary securities. Applied before momentum signal
        construction so non-primary stocks are excluded throughout.

    Returns
    -------
    FactorResult
        factors DataFrame has a single column ["mom"].

    References
    ----------
    Jegadeesh, N., & Titman, S. (1993). Returns to buying winners and
        selling losers: Implications for stock market efficiency.
        Journal of Finance, 48(1), 65-91.
    """
    df = data.copy()
    # Apply primary securities filter before signal construction
    if primary_col is not None:
        df = df[df[primary_col].astype(bool)].copy()
    df = df.sort_values([entity_col, date_col])

    # Compute cumulative return over the formation window for each stock
    # Use log returns for additive rolling, then convert back
    df["_log1r"] = np.log1p(df[ret_col])

    # Within each entity, compute the cumulative sum of log(1+r)
    df["_cumlog"] = df.groupby(entity_col, sort=False)["_log1r"].cumsum()

    # Momentum signal at time t:
    # _cumlog_start = cumulative log return through t-formation_start
    # _cumlog_end   = cumulative log return through t-formation_end
    # _mom = exp(cumlog_end - cumlog_start) - 1
    #      = product of gross returns from t-(formation_start-1) to t-formation_end
    df["_cumlog_start"] = df.groupby(entity_col, sort=False)["_cumlog"].shift(
        formation_start,
    )
    df["_cumlog_end"] = df.groupby(entity_col, sort=False)["_cumlog"].shift(
        formation_end,
    )

    # Momentum = exp(cumlog_end - cumlog_start) - 1
    df["_mom"] = np.exp(df["_cumlog_end"] - df["_cumlog_start"]) - 1

    # Drop rows without valid momentum signal
    df_valid = df.dropna(subset=["_mom"]).copy()

    # Use char_factor with 2x3 sort on size x momentum; forward exchcd_col for
    # NYSE breakpoints so the caller's column naming is respected
    mom_result = char_factor(
        df_valid, "_mom",
        ret_col=ret_col, date_col=date_col, entity_col=entity_col,
        mktcap_col=mktcap_col, exchcd_col=exchcd_col, size_col=mktcap_col,
        ascending=True,
        primary_col=primary_col,
    )

    # Rename factor column
    factors = mom_result.factors.rename(columns={"_mom_factor": "mom"})
    # Keep only the mom column (drop smb from this sort)
    if "smb" in factors.columns:
        factors = factors[["mom"]]

    return FactorResult(factors=factors, portfolios=mom_result.portfolios)
