"""Tests for eapctf.crosssection: FMB, CS regression, multiple testing."""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from eapctf.crosssection import (
    CSRegressionResult,
    FMBResult,
    cs_regression,
    fama_macbeth,
    multiple_testing_correction,
)


# ---------------------------------------------------------------------------
# cs_regression
# ---------------------------------------------------------------------------


class TestCSRegression:
    def test_happy_path(self, small_panel: pd.DataFrame) -> None:
        """cs_regression should return a CSRegressionResult with correct fields."""
        result = cs_regression(small_panel, ret_col="ret", char_cols=["bm", "op"])
        assert isinstance(result, CSRegressionResult)
        # Fields per the updated dataclass
        assert hasattr(result, "coefs")
        assert hasattr(result, "se")
        assert hasattr(result, "t_stats")
        assert hasattr(result, "r_squared")
        assert hasattr(result, "n_obs")

    def test_coef_index_matches_char_cols(self, small_panel: pd.DataFrame) -> None:
        """coefs index should contain the characteristic column names."""
        result = cs_regression(small_panel, ret_col="ret", char_cols=["bm", "op", "inv"])
        assert list(result.coefs.index) == ["bm", "op", "inv"]

    def test_t_stats_sign(self, small_panel: pd.DataFrame) -> None:
        """t-stats should have same sign as coefficients (when se > 0)."""
        result = cs_regression(small_panel, ret_col="ret", char_cols=["bm"])
        coef_sign = np.sign(result.coefs["bm"])
        t_sign = np.sign(result.t_stats["bm"])
        assert coef_sign == t_sign

    def test_double_cluster(self, small_panel: pd.DataFrame) -> None:
        """double_cluster=True should also run without error."""
        result = cs_regression(
            small_panel, ret_col="ret", char_cols=["bm"], double_cluster=True
        )
        assert isinstance(result, CSRegressionResult)
        assert result.n_obs > 0

    def test_no_char_cols_raises(self, small_panel: pd.DataFrame) -> None:
        """Providing no char_cols or controls should raise ValueError."""
        with pytest.raises(ValueError):
            cs_regression(small_panel, ret_col="ret", char_cols=None, controls=None)

    def test_n_obs_correct(self, small_panel: pd.DataFrame) -> None:
        """n_obs should equal the number of non-missing rows."""
        panel = small_panel.copy()
        panel.loc[panel.index[:10], "bm"] = np.nan
        result = cs_regression(panel, ret_col="ret", char_cols=["bm"])
        # Only rows with non-NaN bm and ret should be included
        expected_n = panel.dropna(subset=["ret", "bm"]).shape[0]
        assert result.n_obs == expected_n

    def test_missing_values_handled(self, small_panel: pd.DataFrame) -> None:
        """cs_regression should work even with NaN values in characteristics."""
        panel = small_panel.copy()
        panel.loc[panel["date"] == 1, "bm"] = np.nan
        result = cs_regression(panel, ret_col="ret", char_cols=["bm"])
        assert isinstance(result, CSRegressionResult)

    def test_single_date(self) -> None:
        """cs_regression should work with a single date panel."""
        df = pd.DataFrame(
            {
                "permno": list(range(1001, 1031)),
                "date": [1] * 30,
                "ret": np.random.normal(size=30),
                "bm": np.random.lognormal(size=30),
            }
        )
        result = cs_regression(df, ret_col="ret", char_cols=["bm"])
        assert isinstance(result, CSRegressionResult)


# ---------------------------------------------------------------------------
# fama_macbeth
# ---------------------------------------------------------------------------


class TestFamaMacBeth:
    def test_happy_path_char_cols(self, small_panel: pd.DataFrame) -> None:
        """fama_macbeth with char_cols should return an FMBResult."""
        result = fama_macbeth(small_panel, ret_col="ret", char_cols=["bm", "op"])
        assert isinstance(result, FMBResult)
        assert isinstance(result.lambdas, pd.DataFrame)
        assert "coef" in result.lambdas.columns
        assert "t_nw" in result.lambdas.columns
        assert "se_nw" in result.lambdas.columns

    def test_lambda_index_matches_chars(self, small_panel: pd.DataFrame) -> None:
        """lambdas index should contain the characteristic column names."""
        result = fama_macbeth(small_panel, ret_col="ret", char_cols=["bm", "op"])
        assert set(result.lambdas.index) == {"bm", "op"}

    def test_r_squared_in_range(self, small_panel: pd.DataFrame) -> None:
        """Average cross-sectional R-squared should be in [0, 1]."""
        result = fama_macbeth(small_panel, ret_col="ret", char_cols=["bm"])
        assert 0.0 <= result.r_squared <= 1.0

    def test_n_periods_positive(self, small_panel: pd.DataFrame) -> None:
        """n_periods should be positive."""
        result = fama_macbeth(small_panel, ret_col="ret", char_cols=["bm", "inv"])
        assert result.n_periods > 0

    def test_lambda_ts_shape(self, small_panel: pd.DataFrame) -> None:
        """lambda_ts should have shape (n_periods, n_char_cols)."""
        result = fama_macbeth(small_panel, ret_col="ret", char_cols=["bm", "op"])
        assert result.lambda_ts.shape[1] == 2

    def test_no_factor_no_char_raises(self, small_panel: pd.DataFrame) -> None:
        """Providing neither factor_cols nor char_cols should raise ValueError."""
        with pytest.raises(ValueError):
            fama_macbeth(small_panel, ret_col="ret")

    def test_missing_values(self, small_panel: pd.DataFrame) -> None:
        """fama_macbeth should handle NaN values in characteristics."""
        panel = small_panel.copy()
        panel.loc[panel["date"] == 1, "bm"] = np.nan
        result = fama_macbeth(panel, ret_col="ret", char_cols=["bm"])
        assert isinstance(result, FMBResult)

    def test_two_pass_factor_cols(self, small_panel: pd.DataFrame) -> None:
        """Two-pass FMB with factor_cols should return risk premia indexed by beta_<factor>."""
        result = fama_macbeth(
            small_panel, ret_col="ret", factor_cols=["bm", "op"],
            shanken_correction=True,
        )
        assert isinstance(result, FMBResult)
        # lambda index should be beta_bm, beta_op
        assert "beta_bm" in result.lambdas.index
        assert "beta_op" in result.lambdas.index
        assert result.shanken_applied is True

    def test_se_nw_column_present(self, small_panel: pd.DataFrame) -> None:
        """lambdas DataFrame must have se_nw and t_nw columns (not se_ols/t_ols)."""
        result = fama_macbeth(small_panel, ret_col="ret", char_cols=["bm"])
        assert "se_nw" in result.lambdas.columns
        assert "t_nw" in result.lambdas.columns
        assert "se_ols" not in result.lambdas.columns
        assert "t_ols" not in result.lambdas.columns

    def test_shanken_se_ge_nw_se(self, small_panel: pd.DataFrame) -> None:
        """Shanken SE must be >= NW SE for beta regressors (c >= 1 by construction)."""
        result = fama_macbeth(
            small_panel, ret_col="ret", factor_cols=["bm"],
            shanken_correction=True,
        )
        assert result.shanken_applied is True
        for _, row in result.lambdas.iterrows():
            assert row["se_shanken"] >= row["se_nw"] - 1e-12


# ---------------------------------------------------------------------------
# multiple_testing_correction
# ---------------------------------------------------------------------------


class TestMultipleTestingCorrection:
    def _make_t_stats(self, n: int = 20) -> pd.Series:
        np.random.seed(42)
        vals = np.random.normal(0, 2.0, n)
        return pd.Series(vals, index=[f"anomaly_{i}" for i in range(n)])

    def test_bhy_method(self) -> None:
        """BHY correction should return a DataFrame with expected columns."""
        t_stats = self._make_t_stats()
        result = multiple_testing_correction(t_stats, method="bhy")
        assert isinstance(result, pd.DataFrame)
        assert "t_stat" in result.columns
        assert "p_value" in result.columns
        assert "significant" in result.columns
        assert len(result) == len(t_stats)

    def test_bonferroni_method(self) -> None:
        """Bonferroni correction should return correct shape."""
        t_stats = self._make_t_stats()
        result = multiple_testing_correction(t_stats, method="bonferroni")
        assert len(result) == len(t_stats)

    def test_holm_method(self) -> None:
        """Holm correction should return correct shape."""
        t_stats = self._make_t_stats()
        result = multiple_testing_correction(t_stats, method="holm")
        assert len(result) == len(t_stats)

    def test_unknown_method_raises(self) -> None:
        """Unknown correction method should raise ValueError."""
        t_stats = self._make_t_stats()
        with pytest.raises(ValueError):
            multiple_testing_correction(t_stats, method="unknown_method")

    def test_large_t_stat_significant(self) -> None:
        """A very large t-statistic should be significant even after correction."""
        t_stats = pd.Series({"anomaly_big": 10.0, "anomaly_small": 0.1})
        result = multiple_testing_correction(t_stats, method="bonferroni", alpha=0.05)
        assert result.loc["anomaly_big", "significant"]

    def test_p_values_in_range(self) -> None:
        """p-values should all be in [0, 1]."""
        t_stats = self._make_t_stats(30)
        result = multiple_testing_correction(t_stats, method="bhy")
        assert (result["p_value"] >= 0.0).all()
        assert (result["p_value"] <= 1.0).all()
