"""Tests for eapctf.sorting: univariate/bivariate sorts and factor construction."""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from eapctf.sorting import (
    FactorResult,
    SortResult,
    bivariate_sort,
    char_factor,
    ff3_factors,
    nyse_breakpoints,
    univariate_sort,
)


# ---------------------------------------------------------------------------
# nyse_breakpoints
# ---------------------------------------------------------------------------


class TestNyseBreakpoints:
    def test_output_shape(self, small_panel: pd.DataFrame) -> None:
        """nyse_breakpoints should return a DataFrame with date + q_* columns."""
        bp = nyse_breakpoints(small_panel, char_col="bm", n_portfolios=5)
        assert "date" in bp.columns
        bp_cols = [c for c in bp.columns if c.startswith("q_")]
        # n_portfolios=5 means 4 quantile breakpoints
        assert len(bp_cols) == 4
        assert len(bp) == small_panel["date"].nunique()

    def test_single_date(self) -> None:
        """nyse_breakpoints with a single date should work."""
        df = pd.DataFrame(
            {
                "permno": list(range(1001, 1021)),
                "date": [1] * 20,
                "bm": np.random.lognormal(size=20),
                "exchcd": [1] * 10 + [2] * 10,
            }
        )
        bp = nyse_breakpoints(df, char_col="bm", n_portfolios=3)
        assert len(bp) == 1


# ---------------------------------------------------------------------------
# univariate_sort
# ---------------------------------------------------------------------------


class TestUnivariateSort:
    def test_happy_path(self, small_panel: pd.DataFrame) -> None:
        """univariate_sort should return a SortResult with correct structure."""
        result = univariate_sort(small_panel, char_col="bm", n_portfolios=5)
        assert isinstance(result, SortResult)
        # portfolio_returns: date-indexed DataFrame
        assert "port_1" in result.portfolio_returns.columns
        assert "port_5" in result.portfolio_returns.columns
        assert "long_short" in result.portfolio_returns.columns
        # portfolio_stats: one row per portfolio
        assert len(result.portfolio_stats) == 6  # 5 ports + long_short
        # breakpoints: date-indexed
        assert len(result.breakpoints) == small_panel["date"].nunique()

    def test_equal_weighted(self, small_panel: pd.DataFrame) -> None:
        """equal-weighted sort should return a valid SortResult."""
        result = univariate_sort(small_panel, char_col="bm", n_portfolios=3, weighting="ew")
        assert "port_1" in result.portfolio_returns.columns
        assert "long_short" in result.portfolio_returns.columns

    def test_all_breakpoints(self, small_panel: pd.DataFrame) -> None:
        """breakpoints='all' should also work (no NYSE filter)."""
        result = univariate_sort(
            small_panel, char_col="bm", n_portfolios=3, breakpoints="all"
        )
        assert isinstance(result, SortResult)

    def test_long_short_is_difference(self, small_panel: pd.DataFrame) -> None:
        """long_short column should equal port_N minus port_1."""
        result = univariate_sort(
            small_panel, char_col="bm", n_portfolios=5, weighting="ew", breakpoints="all"
        )
        expected_ls = result.portfolio_returns["port_5"] - result.portfolio_returns["port_1"]
        pd.testing.assert_series_equal(
            result.portfolio_returns["long_short"].dropna(),
            expected_ls.dropna(),
            check_names=False,
        )

    def test_missing_characteristic(self, small_panel: pd.DataFrame) -> None:
        """Panel with NaN in characteristic should still run."""
        panel = small_panel.copy()
        panel.loc[panel["date"] == 1, "bm"] = np.nan
        result = univariate_sort(panel, char_col="bm", n_portfolios=5)
        assert isinstance(result, SortResult)

    def test_single_date(self) -> None:
        """Panel with a single date should not crash."""
        df = pd.DataFrame(
            {
                "permno": list(range(1001, 1021)),
                "date": [1] * 20,
                "ret": np.random.normal(size=20),
                "mktcap": np.random.lognormal(size=20),
                "exchcd": [1] * 10 + [2] * 10,
                "bm": np.random.lognormal(size=20),
            }
        )
        result = univariate_sort(df, char_col="bm", n_portfolios=3, breakpoints="all")
        assert isinstance(result, SortResult)


# ---------------------------------------------------------------------------
# bivariate_sort
# ---------------------------------------------------------------------------


class TestBivariateSort:
    def test_happy_path(self, small_panel: pd.DataFrame) -> None:
        """bivariate_sort should return a SortResult."""
        result = bivariate_sort(
            small_panel, char1_col="bm", char2_col="op", n1=3, n2=3, breakpoints="all"
        )
        assert isinstance(result, SortResult)
        # portfolio_returns columns should have intersection format
        port_cols = [c for c in result.portfolio_returns.columns if c.startswith("port_")]
        assert len(port_cols) > 0

    def test_dependent_sort(self, small_panel: pd.DataFrame) -> None:
        """Dependent bivariate sort should also work."""
        result = bivariate_sort(
            small_panel,
            char1_col="bm",
            char2_col="op",
            n1=3,
            n2=3,
            sort_type="dependent",
            breakpoints="all",
        )
        assert isinstance(result, SortResult)

    def test_portfolio_stats_index(self, small_panel: pd.DataFrame) -> None:
        """portfolio_stats should be indexed by portfolio names."""
        result = bivariate_sort(
            small_panel, char1_col="bm", char2_col="op", n1=2, n2=2, breakpoints="all"
        )
        assert len(result.portfolio_stats) > 0


# ---------------------------------------------------------------------------
# primary_col filter and vw_cap weighting
# ---------------------------------------------------------------------------


class TestPrimaryColFilter:
    def test_primary_col_excludes_non_primary(self, small_panel: pd.DataFrame) -> None:
        """char_factor with primary_col should exclude rows where primary is False."""
        panel = small_panel.copy()
        # Mark only the first 25 stocks (NYSE) as primary
        panel["primary"] = panel["permno"] < 1026

        result_all = char_factor(panel, "bm", size_col="mktcap")
        result_primary = char_factor(panel, "bm", size_col="mktcap", primary_col="primary")

        assert isinstance(result_primary, FactorResult)
        # Factor values should differ because half the universe is excluded
        common_idx = result_all.factors.index.intersection(result_primary.factors.index)
        factor_col = "bm_factor"
        diff = (
            result_all.factors.loc[common_idx, factor_col]
            - result_primary.factors.loc[common_idx, factor_col]
        ).abs()
        # At least some periods should differ (non-primary stocks affect weights)
        assert diff.sum() > 0

    def test_ff3_primary_col_runs(self, small_panel: pd.DataFrame) -> None:
        """ff3_factors with primary_col should return a FactorResult with mkt_rf, smb, hml."""
        panel = small_panel.copy()
        panel["primary"] = panel["permno"] < 1026
        result = ff3_factors(panel, primary_col="primary")
        assert isinstance(result, FactorResult)
        assert "mkt_rf" in result.factors.columns
        assert "smb" in result.factors.columns
        assert "hml" in result.factors.columns


class TestVwCapWeighting:
    def test_vw_cap_returns_sort_result(self, small_panel: pd.DataFrame) -> None:
        """univariate_sort with vw_cap weighting should succeed."""
        result = univariate_sort(
            small_panel, char_col="bm", n_portfolios=5, weighting="vw_cap"
        )
        assert isinstance(result, SortResult)
        assert "port_1" in result.portfolio_returns.columns
        assert "long_short" in result.portfolio_returns.columns

    def test_vw_cap_differs_from_vw(self, small_panel: pd.DataFrame) -> None:
        """vw_cap capping should produce different portfolio returns than standard vw."""
        result_vw = univariate_sort(
            small_panel, char_col="bm", n_portfolios=5, weighting="vw"
        )
        result_cap = univariate_sort(
            small_panel, char_col="bm", n_portfolios=5, weighting="vw_cap"
        )
        # At least one portfolio should differ (capping at NYSE p80 changes weights)
        diff = (
            result_vw.portfolio_returns - result_cap.portfolio_returns
        ).abs().sum().sum()
        assert diff > 0
