"""Univariate and bivariate portfolio sorts.

Implements single- and double-sort portfolio construction following
standard empirical asset pricing methodology.

References
----------
Bali, T. G., Engle, R. F., & Murray, S. (2016). Empirical asset pricing:
    The cross section of stock returns. Wiley. Ch. 5-7.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Literal

import numpy as np
import pandas as pd

from eapctf.sorting.breakpoints import nyse_breakpoints


@dataclass
class SortResult:
    """Container for portfolio sort results.

    Attributes
    ----------
    portfolio_returns : pd.DataFrame
        Date-indexed DataFrame with columns [port_1, ..., port_N, long_short].
    portfolio_stats : pd.DataFrame
        Summary statistics per portfolio: mean, std, SR, alpha, t-alpha.
    breakpoints : pd.DataFrame
        Date-indexed breakpoint values used in the sort.
    """

    portfolio_returns: pd.DataFrame
    portfolio_stats: pd.DataFrame
    breakpoints: pd.DataFrame


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


def _classify_column(
    data: pd.DataFrame,
    char_col: str,
    n_portfolios: int,
    date_col: str,
    bp_df: pd.DataFrame | None = None,
) -> pd.Series[Any]:
    """Assign integer portfolio labels 1..n_portfolios per date.

    If ``bp_df`` is provided, those breakpoints are used (e.g., NYSE-only).
    Otherwise, unconditional within-date quantiles are computed.
    """
    if bp_df is not None:
        # Merge breakpoints onto every row via date
        bp_cols = [c for c in bp_df.columns if c.startswith("q_")]
        merged = data[[date_col, char_col]].merge(
            bp_df, on=date_col, how="left",
        )

        # Vectorized bin assignment using numpy searchsorted per row
        char_vals = merged[char_col].to_numpy(dtype=float)
        bp_matrix = merged[bp_cols].to_numpy(dtype=float)
        result_arr = np.full(len(char_vals), -1, dtype=int)
        for i in range(len(char_vals)):
            v = char_vals[i]
            if np.isnan(v):
                result_arr[i] = -1
                continue
            # searchsorted finds insertion point; +1 for 1-based labels.
            # side="right" puts exact breakpoint values in the upper portfolio,
            # matching the Fama-French convention (value >= breakpoint → upper bin).
            bps = bp_matrix[i]
            if np.any(np.isnan(bps)):
                # No valid NYSE stocks for this date — skip row
                result_arr[i] = -1
                continue
            result_arr[i] = int(np.searchsorted(bps, v, side="right")) + 1

        return pd.Series(result_arr, index=data.index)

    # Fallback: equal-frequency quantile bins per date
    all_labels: list[pd.Series[Any]] = []
    for _dt, group in data.groupby(date_col, sort=False):
        vals = group[char_col]
        try:
            bins = pd.qcut(vals, n_portfolios, labels=False, duplicates="drop") + 1
        except ValueError:
            # Tie-breaking fallback: rank first so qcut always finds distinct quantiles
            bins = pd.qcut(vals.rank(method="first"), n_portfolios, labels=False) + 1
        all_labels.append(bins)

    if all_labels:
        labels: pd.Series[Any] = pd.concat(all_labels)
    else:
        labels = pd.Series(dtype="Int64")
    return labels.astype("Int64")


def _compute_portfolio_returns(
    data: pd.DataFrame,
    port_col: str,
    ret_col: str,
    date_col: str,
    mktcap_col: str,
    weighting: str,
    n_portfolios: int,
    exchcd_col: str | None = None,
) -> pd.DataFrame:
    """Compute value- or equal-weighted returns for each portfolio per date."""
    # Drop rows without a valid portfolio assignment or missing return
    df = data.dropna(subset=[port_col, ret_col]).copy()
    df = df[df[port_col] > 0]

    if weighting in ("vw", "vw_cap"):
        if weighting == "vw_cap" and exchcd_col is not None and exchcd_col in df.columns:
            # Cap mktcap at NYSE 80th percentile per date (JKP/pyanomaly vw_cap
            # convention). Winsorizes large stocks to reduce size concentration.
            nyse_p80 = (
                df[df[exchcd_col] == 1]
                .groupby(date_col, sort=False)[mktcap_col]
                .quantile(0.80)
                .rename("_nyse_p80")
            )
            df = df.merge(nyse_p80, on=date_col, how="left")
            w_col = "_capped_mktcap"
            df[w_col] = df[[mktcap_col, "_nyse_p80"]].min(axis=1)
        else:
            w_col = mktcap_col
        # Value-weighted returns: sum(ret * w) / sum(w) per group.
        # Only include firms with valid returns in the denominator weight so
        # that NaN-return firms do not dilute the weighted average.
        df["_weighted_ret"] = df[ret_col] * df[w_col]
        grouped = df.groupby([date_col, port_col], sort=True)
        port_ret_series = grouped["_weighted_ret"].sum() / grouped[w_col].sum()
    else:
        # Equal-weighted returns
        port_ret_series = df.groupby([date_col, port_col], sort=True)[ret_col].mean()

    # Pivot to wide format: date x portfolio
    port_ret_df = port_ret_series.reset_index()
    port_ret_df.columns = pd.Index([date_col, port_col, "ret"])
    wide = port_ret_df.pivot(index=date_col, columns=port_col, values="ret")

    # Rename columns to port_1, port_2, ...
    col_map = {p: f"port_{int(p)}" for p in sorted(wide.columns)}
    wide = wide.rename(columns=col_map)

    # Ensure all portfolio columns are present (fill missing with NaN)
    for i in range(1, n_portfolios + 1):
        col_name = f"port_{i}"
        if col_name not in wide.columns:
            wide[col_name] = np.nan

    # Reorder columns
    ordered_cols = [f"port_{i}" for i in range(1, n_portfolios + 1)]
    wide = wide[ordered_cols]

    return wide


def _compute_portfolio_stats(port_returns: pd.DataFrame) -> pd.DataFrame:
    """Compute summary statistics for each portfolio column."""
    stats_list: list[dict[str, object]] = []
    for col in port_returns.columns:
        ret_series = port_returns[col].dropna()
        mean_ret = float(ret_series.mean())
        std_ret = float(ret_series.std())
        # Annualized Sharpe ratio (monthly data, no risk-free subtraction here)
        sr = (mean_ret / std_ret * np.sqrt(12)) if std_ret > 0 else 0.0
        # t-statistic for the mean
        n = len(ret_series)
        t_stat = (mean_ret / (std_ret / np.sqrt(n))) if std_ret > 0 and n > 1 else 0.0
        stats_list.append(
            {
                "portfolio": col,
                "mean": mean_ret,
                "std": std_ret,
                "sharpe": sr,
                "t_mean": t_stat,
            }
        )
    return pd.DataFrame(stats_list).set_index("portfolio")


def univariate_sort(
    data: pd.DataFrame,
    char_col: str,
    ret_col: str = "ret",
    date_col: str = "date",
    entity_col: str = "permno",
    mktcap_col: str = "mktcap",
    n_portfolios: int = 10,
    breakpoints: Literal["all", "nyse"] = "nyse",
    weighting: Literal["vw", "ew", "vw_cap"] = "vw",
    long_port: int | None = None,
    short_port: int | None = None,
    min_nyse_pct: float | None = 20.0,
) -> SortResult:
    """Single-sort portfolios.

    Sorts stocks into ``n_portfolios`` groups each period based on
    ``char_col``, then computes portfolio returns using value- or
    equal-weighting.

    Parameters
    ----------
    data : pd.DataFrame
        Panel data with columns for date, entity, return, market cap,
        and the sorting characteristic.
    char_col : str
        Characteristic column name to sort on.
    ret_col : str, default "ret"
        Return column name.
    date_col : str, default "date"
        Date column name.
    entity_col : str, default "permno"
        Entity identifier column name.
    mktcap_col : str, default "mktcap"
        Market capitalization column for value-weighting.
    n_portfolios : int, default 10
        Number of portfolios to form.
    breakpoints : {"all", "nyse"}, default "nyse"
        Whether to use all stocks or only NYSE stocks for breakpoints.
    weighting : {"vw", "ew", "vw_cap"}, default "vw"
        "vw": value-weighted by mktcap. "ew": equal-weighted. "vw_cap":
        value-weighted with mktcap capped at NYSE p80 per date (JKP/pyanomaly
        convention). Requires exchcd column for "vw_cap"; falls back to "vw"
        if exchcd is absent.
    long_port : int or None, default None
        Portfolio number for the long leg. Defaults to n_portfolios.
    short_port : int or None, default None
        Portfolio number for the short leg. Defaults to 1.
    min_nyse_pct : float or None, default 20.0
        Exclude stocks whose market cap falls below this NYSE size percentile
        each period (JKP micro-cap filter). Default is 20, which drops the
        bottom 20% of NYSE size — the JKP (2023) standard. Set to None to
        include all stocks. Requires ``exchcd`` and ``mktcap`` columns.

    Returns
    -------
    SortResult
        Contains portfolio_returns, portfolio_stats, and breakpoints.

    References
    ----------
    Bali, T. G., Engle, R. F., & Murray, S. (2016). Empirical asset pricing:
        The cross section of stock returns. Wiley. Ch. 5-6.
    """
    # Set default long/short portfolios
    if long_port is None:
        long_port = n_portfolios
    if short_port is None:
        short_port = 1

    # Work on a copy to avoid mutating user data
    df = data.copy()

    # Optional micro-cap filter: drop stocks below min_nyse_pct NYSE size percentile
    if min_nyse_pct is not None and "exchcd" in df.columns and mktcap_col in df.columns:
        pct_frac = min_nyse_pct / 100.0
        size_cutoffs = (
            df[df["exchcd"] == 1]
            .groupby(date_col)[mktcap_col]
            .quantile(pct_frac)
            .rename("_size_cutoff")
        )
        df = df.merge(size_cutoffs, on=date_col, how="left")
        df = df[df[mktcap_col] >= df["_size_cutoff"]].drop(columns=["_size_cutoff"])

    # Compute or skip NYSE breakpoints
    bp_df: pd.DataFrame | None = None
    exchcd_col = "exchcd"
    if breakpoints == "nyse" and exchcd_col in df.columns:
        bp_df = nyse_breakpoints(
            df, char_col, n_portfolios=n_portfolios,
            date_col=date_col, exchange_col=exchcd_col,
        )

    # Classify each stock into a portfolio
    port_col = "_port"
    df[port_col] = _classify_column(
        df, char_col, n_portfolios, date_col, bp_df=bp_df,
    )

    # Compute portfolio returns (wide format)
    port_returns = _compute_portfolio_returns(
        df, port_col, ret_col, date_col, mktcap_col, weighting, n_portfolios,
        exchcd_col=exchcd_col,
    )

    # Build long-short column
    long_col = f"port_{long_port}"
    short_col = f"port_{short_port}"
    port_returns["long_short"] = port_returns[long_col] - port_returns[short_col]

    # Summary statistics
    port_stats = _compute_portfolio_stats(port_returns)

    # Build breakpoints DataFrame for output
    if bp_df is not None:
        bp_out = bp_df.copy()
    else:
        # Compute breakpoints on all stocks for the output record
        bp_out = _all_stock_breakpoints(df, char_col, n_portfolios, date_col)

    return SortResult(
        portfolio_returns=port_returns,
        portfolio_stats=port_stats,
        breakpoints=bp_out,
    )


def _all_stock_breakpoints(
    data: pd.DataFrame,
    char_col: str,
    n_portfolios: int,
    date_col: str,
) -> pd.DataFrame:
    """Compute quantile breakpoints using all stocks (fallback)."""
    quantile_fracs = np.linspace(0.0, 1.0, n_portfolios + 1)[1:-1]
    labels = [f"q_{i + 1}" for i in range(len(quantile_fracs))]

    rows: list[dict[str, object]] = []
    for dt, group in data.groupby(date_col, sort=True):
        vals = group[char_col].dropna()
        row: dict[str, object] = {date_col: dt}
        if len(vals) == 0:
            for lbl in labels:
                row[lbl] = np.nan
        else:
            q = np.quantile(vals, quantile_fracs)
            for lbl, qval in zip(labels, q):
                row[lbl] = float(qval)
        rows.append(row)

    return pd.DataFrame(rows)


def bivariate_sort(
    data: pd.DataFrame,
    char1_col: str,
    char2_col: str,
    ret_col: str = "ret",
    date_col: str = "date",
    entity_col: str = "permno",
    mktcap_col: str = "mktcap",
    n1: int = 5,
    n2: int = 5,
    sort_type: Literal["independent", "dependent"] = "independent",
    weighting: Literal["vw", "ew", "vw_cap"] = "vw",
    breakpoints: Literal["all", "nyse"] = "nyse",
    min_nyse_pct: float | None = 20.0,
) -> SortResult:
    """Double-sort portfolios (independent or sequential conditional sort).

    Parameters
    ----------
    data : pd.DataFrame
        Panel data.
    char1_col : str
        First sorting characteristic.
    char2_col : str
        Second sorting characteristic.
    ret_col : str, default "ret"
        Return column name.
    date_col : str, default "date"
        Date column name.
    entity_col : str, default "permno"
        Entity identifier column name.
    mktcap_col : str, default "mktcap"
        Market capitalization column.
    n1 : int, default 5
        Number of portfolios for the first sort.
    n2 : int, default 5
        Number of portfolios for the second sort.
    sort_type : {"independent", "dependent"}, default "independent"
        Independent: classify each characteristic separately, form intersections.
        Dependent: sort on char1 first, then char2 within each char1 group.
    weighting : {"vw", "ew", "vw_cap"}, default "vw"
        "vw": value-weighted by mktcap. "ew": equal-weighted. "vw_cap":
        value-weighted with mktcap capped at NYSE p80 per date.
    breakpoints : {"all", "nyse"}, default "nyse"
        Source of breakpoints.
    min_nyse_pct : float or None, default 20.0
        Exclude stocks whose market cap falls below this NYSE size percentile
        each period (JKP micro-cap filter). Default is 20, which drops the
        bottom 20% of NYSE size — the JKP (2023) standard. Set to None to
        include all stocks. Requires ``exchcd`` and ``mktcap`` columns.

    Returns
    -------
    SortResult
        portfolio_returns has columns for each (i, j) intersection portfolio
        named port_i_j.

    References
    ----------
    Bali, T. G., Engle, R. F., & Murray, S. (2016). Empirical asset pricing:
        The cross section of stock returns. Wiley. Ch. 7.
    """
    df = data.copy()

    # Optional micro-cap filter
    if min_nyse_pct is not None and "exchcd" in df.columns and mktcap_col in df.columns:
        pct_frac = min_nyse_pct / 100.0
        size_cutoffs = (
            df[df["exchcd"] == 1]
            .groupby(date_col)[mktcap_col]
            .quantile(pct_frac)
            .rename("_size_cutoff")
        )
        df = df.merge(size_cutoffs, on=date_col, how="left")
        df = df[df[mktcap_col] >= df["_size_cutoff"]].drop(columns=["_size_cutoff"])

    # Determine breakpoint source
    exchcd_col = "exchcd"
    bp1_df: pd.DataFrame | None = None
    bp2_df: pd.DataFrame | None = None
    if breakpoints == "nyse" and exchcd_col in df.columns:
        bp1_df = nyse_breakpoints(
            df, char1_col, n_portfolios=n1,
            date_col=date_col, exchange_col=exchcd_col,
        )
        bp2_df = nyse_breakpoints(
            df, char2_col, n_portfolios=n2,
            date_col=date_col, exchange_col=exchcd_col,
        )

    # First sort: classify on char1
    df["_port1"] = _classify_column(df, char1_col, n1, date_col, bp_df=bp1_df)

    if sort_type == "independent":
        # Independent: classify char2 on the full sample independently
        df["_port2"] = _classify_column(df, char2_col, n2, date_col, bp_df=bp2_df)
    else:
        # Dependent: classify char2 within each char1 group
        all_bins: list[pd.Series[Any]] = []
        for (_dt, _p1), group in df.groupby([date_col, "_port1"], sort=False):
            vals = group[char2_col]
            try:
                bins = pd.qcut(vals, n2, labels=False, duplicates="drop") + 1
            except ValueError:
                # Tie-breaking fallback: rank first so qcut always finds distinct quantiles
                bins = pd.qcut(vals.rank(method="first"), n2, labels=False) + 1
            all_bins.append(bins)

        if all_bins:
            df["_port2"] = pd.concat(all_bins)
        else:
            df["_port2"] = np.nan

    # Form intersection labels: "i_j"
    df["_biport"] = df["_port1"].astype(str) + "_" + df["_port2"].astype(str)

    # Drop rows without valid assignments
    valid = (df["_port1"] > 0) & (df["_port2"] > 0)
    df = df[valid].copy()

    # Compute weighted returns for each intersection portfolio per date
    if weighting in ("vw", "vw_cap"):
        if weighting == "vw_cap" and exchcd_col in df.columns:
            # Cap mktcap at NYSE 80th percentile per date (pyanomaly convention)
            nyse_p80 = (
                df[df[exchcd_col] == 1]
                .groupby(date_col, sort=False)[mktcap_col]
                .quantile(0.80)
                .rename("_nyse_p80")
            )
            df = df.merge(nyse_p80, on=date_col, how="left")
            w_col = "_capped_mktcap"
            df[w_col] = df[[mktcap_col, "_nyse_p80"]].min(axis=1)
        else:
            w_col = mktcap_col
        df["_wret"] = df[ret_col] * df[w_col]
        grouped = df.groupby([date_col, "_biport"], sort=True)
        port_ret_series = grouped["_wret"].sum() / grouped[w_col].sum()
    else:
        port_ret_series = df.groupby([date_col, "_biport"], sort=True)[ret_col].mean()

    port_ret_df = port_ret_series.reset_index()
    port_ret_df.columns = pd.Index([date_col, "biport", "ret"])
    wide = port_ret_df.pivot(index=date_col, columns="biport", values="ret")

    # Rename columns to port_i_j format
    col_rename = {c: f"port_{c}" for c in wide.columns}
    wide = wide.rename(columns=col_rename)

    # Compute stats
    port_stats = _compute_portfolio_stats(wide)

    # Build breakpoint output (char1 breakpoints)
    if bp1_df is not None:
        bp_out = bp1_df.copy()
    else:
        bp_out = _all_stock_breakpoints(df, char1_col, n1, date_col)

    return SortResult(
        portfolio_returns=wide,
        portfolio_stats=port_stats,
        breakpoints=bp_out,
    )
