"""Tests for eapctf.sdf: GMM, HJ distance, pricing errors."""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from eapctf.sdf import (
    GMMResult,
    HJResult,
    PricingErrorResult,
    gmm_estimate,
    hj_bounds,
    hj_distance,
    pricing_errors,
)


# ---------------------------------------------------------------------------
# Fixtures: small asset universe for speed
# ---------------------------------------------------------------------------


def _make_sdf_data(
    T: int = 60, N: int = 5, K: int = 2, seed: int = 42
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Create T x N returns and T x K factors."""
    np.random.seed(seed)
    factor_mat = np.random.normal(0.005, 0.04, (T, K))
    factors = pd.DataFrame(factor_mat, columns=[f"f{i+1}" for i in range(K)])

    # Returns: generated from linear factor model
    betas = np.random.normal(1.0, 0.3, (N, K))
    noise = np.random.normal(0, 0.02, (T, N))
    returns_mat = (factor_mat @ betas.T) + noise + 1.0  # gross returns (centered near 1)
    returns = pd.DataFrame(returns_mat, columns=[f"a{i+1}" for i in range(N)])

    return returns, factors


# ---------------------------------------------------------------------------
# gmm_estimate
# ---------------------------------------------------------------------------


class TestGMMEstimate:
    def test_happy_path_two_step(self) -> None:
        """Two-step GMM should return GMMResult with correct shapes."""
        returns, factors = _make_sdf_data()
        result = gmm_estimate(returns, factors, two_step=True)
        assert isinstance(result, GMMResult)
        K = factors.shape[1]
        assert result.b.shape == (K,)
        assert result.b_se.shape == (K,)
        assert len(result.pricing_errors) == returns.shape[1]

    def test_happy_path_one_step(self) -> None:
        """One-step GMM should also work."""
        returns, factors = _make_sdf_data()
        result = gmm_estimate(returns, factors, two_step=False)
        assert isinstance(result, GMMResult)

    def test_j_statistic_non_negative(self) -> None:
        """J-test statistic should be non-negative."""
        returns, factors = _make_sdf_data()
        result = gmm_estimate(returns, factors)
        assert result.j_statistic >= 0.0

    def test_j_p_value_in_range(self) -> None:
        """J-test p-value should be in [0, 1]."""
        returns, factors = _make_sdf_data()
        result = gmm_estimate(returns, factors)
        assert 0.0 <= result.j_p_value <= 1.0

    def test_pricing_errors_series_index(self) -> None:
        """pricing_errors should be a Series indexed by asset names."""
        returns, factors = _make_sdf_data()
        result = gmm_estimate(returns, factors)
        assert list(result.pricing_errors.index) == list(returns.columns)

    def test_se_non_negative(self) -> None:
        """Standard errors for SDF loadings should be non-negative."""
        returns, factors = _make_sdf_data()
        result = gmm_estimate(returns, factors)
        assert (result.b_se >= 0).all()

    def test_single_factor(self) -> None:
        """GMM with a single factor should work."""
        T, N = 60, 5
        np.random.seed(1)
        returns = pd.DataFrame(np.random.normal(1.0, 0.05, (T, N)))
        factors = pd.DataFrame({"mkt": np.random.normal(0.005, 0.04, T)})
        result = gmm_estimate(returns, factors)
        assert isinstance(result, GMMResult)

    def test_row_order_invariance(self) -> None:
        """gmm_estimate must produce identical b regardless of row order in returns."""
        returns, factors = _make_sdf_data(seed=77)
        result_orig = gmm_estimate(returns, factors)

        # Shuffle rows of returns — should give the same b after internal alignment
        returns_shuffled = returns.sample(frac=1, random_state=123)
        result_shuffled = gmm_estimate(returns_shuffled, factors)

        np.testing.assert_allclose(
            result_orig.b, result_shuffled.b, rtol=1e-10,
            err_msg="GMM b changed when return rows were shuffled — alignment bug",
        )

    def test_partial_overlap_drops_missing_periods(self) -> None:
        """When factors cover extra periods not in returns, those are dropped."""
        returns, factors = _make_sdf_data(T=60, seed=99)
        # Append 10 extra rows to factors (indices 60-69)
        extra = pd.DataFrame(
            np.random.normal(size=(10, factors.shape[1])),
            index=range(60, 70),
            columns=factors.columns,
        )
        factors_extended = pd.concat([factors, extra])
        result = gmm_estimate(returns, factors_extended)
        # Should use only 60 common rows
        assert result.b.shape == (factors.shape[1],)


# ---------------------------------------------------------------------------
# hj_distance
# ---------------------------------------------------------------------------


class TestHJDistance:
    def test_happy_path(self) -> None:
        """hj_distance should return an HJResult with distance and p_value."""
        T, N = 60, 5
        np.random.seed(2)
        returns = pd.DataFrame(np.random.normal(1.0, 0.05, (T, N)))
        # A "bad" SDF proxy (random)
        sdf_proxy = pd.Series(np.random.normal(1.0, 0.3, T))
        result = hj_distance(returns, sdf_proxy)
        assert isinstance(result, HJResult)
        assert result.distance >= 0.0

    def test_p_value_in_range(self) -> None:
        """p-value should be in [0, 1]."""
        T, N = 60, 5
        np.random.seed(3)
        returns = pd.DataFrame(np.random.normal(1.0, 0.05, (T, N)))
        sdf_proxy = pd.Series(np.ones(T))
        result = hj_distance(returns, sdf_proxy)
        assert 0.0 <= result.p_value <= 1.0

    def test_distance_se_non_negative(self) -> None:
        """distance_se should be non-negative."""
        T, N = 60, 5
        np.random.seed(4)
        returns = pd.DataFrame(np.random.normal(1.0, 0.05, (T, N)))
        sdf_proxy = pd.Series(np.random.normal(1.0, 0.2, T))
        result = hj_distance(returns, sdf_proxy)
        assert result.distance_se >= 0.0


# ---------------------------------------------------------------------------
# hj_bounds
# ---------------------------------------------------------------------------


class TestHJBounds:
    def test_happy_path(self) -> None:
        """hj_bounds should return an HJResult with lower_bound array."""
        T, N = 60, 5
        np.random.seed(5)
        returns = pd.DataFrame(np.random.normal(1.0, 0.05, (T, N)))
        result = hj_bounds(returns, mean_range=(0.95, 1.05), n_grid=20)
        assert isinstance(result, HJResult)
        assert len(result.lower_bound) == 20
        assert len(result.mean_range) == 20

    def test_lower_bound_non_negative(self) -> None:
        """HJ lower bound (std(m)) should be non-negative everywhere."""
        T, N = 60, 5
        np.random.seed(6)
        returns = pd.DataFrame(np.random.normal(1.0, 0.05, (T, N)))
        result = hj_bounds(returns, n_grid=30)
        assert (result.lower_bound >= 0.0).all()


# ---------------------------------------------------------------------------
# pricing_errors
# ---------------------------------------------------------------------------


class TestPricingErrors:
    def test_happy_path(self) -> None:
        """pricing_errors should return a PricingErrorResult."""
        T, N = 60, 5
        np.random.seed(7)
        returns = pd.DataFrame(np.random.normal(1.0, 0.05, (T, N)))
        sdf_proxy = pd.Series(np.random.normal(1.0, 0.2, T))
        result = pricing_errors(returns, sdf_proxy)
        assert isinstance(result, PricingErrorResult)

    def test_mape_rmse_non_negative(self) -> None:
        """MAPE and RMSE should be non-negative."""
        T, N = 60, 5
        np.random.seed(9)
        returns = pd.DataFrame(np.random.normal(1.0, 0.05, (T, N)))
        sdf_proxy = pd.Series(np.random.normal(1.0, 0.2, T))
        result = pricing_errors(returns, sdf_proxy)
        assert result.mape >= 0.0
        assert result.rmse >= 0.0

    def test_pricing_errors_shape(self) -> None:
        """Pricing errors should be indexed by asset names."""
        T, N = 60, 5
        np.random.seed(8)
        rets = pd.DataFrame(np.random.normal(1.0, 0.05, (T, N)), columns=[f"a{i}" for i in range(N)])
        sdf_proxy = pd.Series(np.random.normal(1.0, 0.2, T))
        result = pricing_errors(rets, sdf_proxy)
        assert hasattr(result, "errors")
        assert len(result.errors) == N
        assert list(result.errors.index) == [f"a{i}" for i in range(N)]
