"""Tests for eapctf.timeseries: GRS test, alpha regression, Newey-West SE."""

from __future__ import annotations

import math

import numpy as np
import pandas as pd
import pytest

from eapctf.timeseries import (
    AlphaResult,
    GRSResult,
    SpanningResult,
    grs_test,
    newey_west_se,
    rolling_beta,
    spanning_test,
    time_series_alpha,
)


# ---------------------------------------------------------------------------
# newey_west_se
# ---------------------------------------------------------------------------


class TestNeweyWestSE:
    def test_auto_lag_formula(self) -> None:
        """Auto lag should equal floor(4 * (T/100)^(2/9))."""
        np.random.seed(0)
        T = 120
        expected_lag = int(math.floor(4.0 * (T / 100.0) ** (2.0 / 9.0)))

        # Verify the formula matches what the function uses
        X = np.column_stack([np.ones(T), np.random.normal(size=T)])
        y = np.random.normal(size=T)
        coef = np.linalg.lstsq(X, y, rcond=None)[0]
        resid = y - X @ coef

        # Call with n_lags=None (auto)
        se_auto = newey_west_se(resid, X, n_lags=None)
        # Call with manually computed lag
        se_manual = newey_west_se(resid, X, n_lags=expected_lag)

        np.testing.assert_array_almost_equal(se_auto, se_manual, decimal=10)

    def test_se_positive(self) -> None:
        """Standard errors should be non-negative."""
        np.random.seed(1)
        T = 60
        X = np.column_stack([np.ones(T), np.random.normal(size=T)])
        resid = np.random.normal(size=T)
        se = newey_west_se(resid, X, n_lags=3)
        assert (se >= 0).all()

    def test_shape(self) -> None:
        """SE shape should match the number of regressors."""
        T, K = 80, 3
        X = np.random.normal(size=(T, K))
        resid = np.random.normal(size=T)
        se = newey_west_se(resid, X, n_lags=2)
        assert se.shape == (K,)


# ---------------------------------------------------------------------------
# time_series_alpha
# ---------------------------------------------------------------------------


class TestTimeSeriesAlpha:
    def _make_factor_data(self, T: int = 120) -> tuple[pd.Series, pd.DataFrame]:
        np.random.seed(10)
        idx = pd.RangeIndex(T)
        mkt = pd.Series(np.random.normal(0.005, 0.04, T), index=idx, name="mkt")
        factors = pd.DataFrame({"mkt": mkt})
        port = 0.003 + 0.8 * mkt + np.random.normal(0, 0.02, T)
        return pd.Series(port.values, index=idx), factors

    def test_series_returns_alpha_result(self) -> None:
        """time_series_alpha with a Series should return a single AlphaResult."""
        port, factors = self._make_factor_data()
        result = time_series_alpha(port, factors)
        assert isinstance(result, AlphaResult)
        assert isinstance(result.alpha, float)
        assert isinstance(result.r_squared, float)
        assert 0.0 <= result.r_squared <= 1.0

    def test_dataframe_returns_list(self) -> None:
        """time_series_alpha with a DataFrame should return a list of AlphaResults."""
        T = 120
        idx = pd.RangeIndex(T)
        np.random.seed(11)
        factors = pd.DataFrame({"mkt": np.random.normal(0.005, 0.04, T)}, index=idx)
        ports = pd.DataFrame(
            {"p1": np.random.normal(0.003, 0.05, T), "p2": np.random.normal(0.001, 0.04, T)},
            index=idx,
        )
        results = time_series_alpha(ports, factors)
        assert isinstance(results, list)
        assert len(results) == 2
        for r in results:
            assert isinstance(r, AlphaResult)

    def test_known_alpha(self) -> None:
        """With a strong intercept, alpha should be close to the true value."""
        np.random.seed(999)
        T = 500
        mkt = np.random.normal(0.005, 0.04, T)
        true_alpha = 0.01
        port = true_alpha + 0.9 * mkt + np.random.normal(0, 0.005, T)
        idx = pd.RangeIndex(T)
        result = time_series_alpha(
            pd.Series(port, index=idx),
            pd.DataFrame({"mkt": mkt}, index=idx),
        )
        assert abs(result.alpha - true_alpha) < 0.003


# ---------------------------------------------------------------------------
# grs_test
# ---------------------------------------------------------------------------


class TestGRSTest:
    def _make_wide_data(self, T: int = 120, N: int = 5, K: int = 2) -> tuple:
        np.random.seed(42)
        factors = pd.DataFrame(
            np.random.normal(0.005, 0.04, (T, K)), columns=[f"f{i+1}" for i in range(K)]
        )
        port_returns = pd.DataFrame(
            np.random.normal(0.003, 0.05, (T, N)), columns=[f"p{i+1}" for i in range(N)]
        )
        return port_returns, factors

    def test_happy_path_wide(self) -> None:
        """grs_test with wide-format inputs should return GRSResult."""
        ports, factors = self._make_wide_data()
        result = grs_test(ports, factors)
        assert isinstance(result, GRSResult)
        assert isinstance(result.statistic, float)
        assert 0.0 <= result.p_value <= 1.0
        assert len(result.alphas) == 5

    def test_grs_wide_equals_long(self) -> None:
        """Wide and long input formats should produce identical GRS statistics."""
        ports, factors = self._make_wide_data()

        # Wide format (default)
        result_wide = grs_test(ports, factors, date_col=None)

        # Long format: add a date column and use date_col
        ports_long = ports.copy()
        ports_long["date"] = np.arange(len(ports))
        factors_long = factors.copy()
        factors_long["date"] = np.arange(len(factors))

        result_long = grs_test(ports_long, factors_long, date_col="date")

        assert abs(result_wide.statistic - result_long.statistic) < 1e-10

    def test_single_portfolio(self) -> None:
        """GRS test with a single portfolio (N=1) should work."""
        np.random.seed(5)
        T = 80
        factors = pd.DataFrame({"mkt": np.random.normal(0.005, 0.04, T)})
        single_port = pd.Series(np.random.normal(0.003, 0.05, T), name="p1")
        result = grs_test(single_port, factors)
        assert isinstance(result, GRSResult)
        assert result.p_value >= 0.0

    def test_p_value_in_range(self) -> None:
        """p-value should be in [0, 1]."""
        ports, factors = self._make_wide_data()
        result = grs_test(ports, factors)
        assert 0.0 <= result.p_value <= 1.0

    def test_sharpe_factors_non_negative(self) -> None:
        """sharpe_factors should be non-negative."""
        ports, factors = self._make_wide_data()
        result = grs_test(ports, factors)
        assert result.sharpe_factors >= 0.0

    def test_insufficient_obs_raises(self) -> None:
        """GRS requires T > N + K; otherwise raise ValueError."""
        np.random.seed(9)
        T, N, K = 20, 18, 3   # T - N - K = -1 <= 0
        ports = pd.DataFrame(np.random.normal(size=(T, N)))
        factors = pd.DataFrame(np.random.normal(size=(T, K)))
        with pytest.raises(ValueError, match="T > N \\+ K"):
            grs_test(ports, factors)

    def test_exact_boundary_raises(self) -> None:
        """T - N - K == 0 should also raise ValueError."""
        np.random.seed(10)
        T, N, K = 10, 7, 3   # T - N - K = 0
        ports = pd.DataFrame(np.random.normal(size=(T, N)))
        factors = pd.DataFrame(np.random.normal(size=(T, K)))
        with pytest.raises(ValueError):
            grs_test(ports, factors)


# ---------------------------------------------------------------------------
# rolling_beta
# ---------------------------------------------------------------------------


class TestRollingBeta:
    def test_shape_and_type(self) -> None:
        """rolling_beta should return a Series with the same length."""
        np.random.seed(7)
        T = 100
        idx = pd.RangeIndex(T)
        returns = pd.Series(np.random.normal(0, 0.05, T), index=idx)
        mkt = pd.Series(np.random.normal(0.005, 0.04, T), index=idx)
        betas = rolling_beta(returns, mkt, window=24, min_periods=12)
        assert isinstance(betas, pd.Series)
        # First 11 observations should be NaN (min_periods=12)
        assert betas.iloc[:11].isna().all()

    def test_known_beta(self) -> None:
        """A portfolio that perfectly mirrors the market should have beta=1."""
        np.random.seed(8)
        T = 200
        idx = pd.RangeIndex(T)
        mkt = pd.Series(np.random.normal(0.005, 0.04, T), index=idx)
        # r = 0 + 1.0 * mkt => beta should be 1.0
        returns = mkt.copy()
        betas = rolling_beta(returns, mkt, window=60, min_periods=30)
        valid_betas = betas.dropna()
        np.testing.assert_allclose(valid_betas.values, 1.0, atol=1e-10)


# ---------------------------------------------------------------------------
# spanning_test
# ---------------------------------------------------------------------------


class TestSpanningTest:
    def test_happy_path(self) -> None:
        """spanning_test should return a SpanningResult."""
        np.random.seed(20)
        T = 100
        factors_new = pd.DataFrame({"f_new": np.random.normal(0.005, 0.04, T)})
        factors_base = pd.DataFrame({"mkt": np.random.normal(0.005, 0.04, T)})
        result = spanning_test(factors_new, factors_base)
        assert isinstance(result, SpanningResult)

    def test_p_value_range(self) -> None:
        """p-value from spanning test should be in [0, 1]."""
        np.random.seed(21)
        T = 80
        factors_new = pd.DataFrame(
            {"f1": np.random.normal(0, 0.04, T), "f2": np.random.normal(0, 0.03, T)}
        )
        factors_base = pd.DataFrame({"mkt": np.random.normal(0.005, 0.04, T)})
        result = spanning_test(factors_new, factors_base)
        assert 0.0 <= result.p_value <= 1.0
