"""Numerical validation tests against reference implementations.

These tests verify that eap's statistical routines produce results that match
either (a) a known closed-form reference or (b) an authoritative external
library (statsmodels).  A package that passes only internal unit tests can
still silently compute wrong numbers.  These tests close that gap.

Each test documents the reference against which the comparison is made so
that the claim is falsifiable.
"""

from __future__ import annotations

import math

import numpy as np
import pandas as pd
import pytest
import statsmodels.api as sm

from eapctf.crosssection import fama_macbeth
from eapctf.timeseries import grs_test, newey_west_se


# ---------------------------------------------------------------------------
# newey_west_se vs statsmodels HAC
# ---------------------------------------------------------------------------


class TestNeweyWestVsStatsmodels:
    """Compare newey_west_se output against statsmodels OLS HAC standard errors.

    Reference: statsmodels.regression.linear_model.OLS.fit(
        cov_type='HAC', cov_kwds={'maxlags': q, 'use_correction': False}
    )
    The Bartlett kernel implementation and T-normalisation must be identical.
    """

    def _ols_data(self, T: int, K: int, seed: int) -> tuple:
        np.random.seed(seed)
        X = np.column_stack([np.ones(T)] + [np.random.normal(size=T) for _ in range(K - 1)])
        beta = np.random.normal(0, 1, K)
        y = X @ beta + np.random.normal(0, 0.5, T)
        coef = np.linalg.lstsq(X, y, rcond=None)[0]
        resid = y - X @ coef
        return X, y, resid

    @pytest.mark.parametrize("T,K,lags", [
        (120, 2, 4),
        (60,  3, 2),
        (250, 4, 6),
        (100, 2, 0),
    ])
    def test_matches_statsmodels(self, T: int, K: int, lags: int) -> None:
        """newey_west_se must match statsmodels HAC SE to machine precision."""
        X, y, resid = self._ols_data(T, K, seed=T + K + lags)

        our_se = newey_west_se(resid, X, n_lags=lags)

        res_sm = sm.OLS(y, X).fit(
            cov_type="HAC",
            cov_kwds={"maxlags": lags, "use_correction": False},
        )

        np.testing.assert_allclose(
            our_se, res_sm.bse, rtol=1e-10,
            err_msg=f"NW SE mismatch vs statsmodels (T={T}, K={K}, lags={lags})",
        )

    def test_auto_lag_matches_manual(self) -> None:
        """Auto-lag (n_lags=None) must equal the Andrews rule-of-thumb formula."""
        T = 120
        expected_lag = int(math.floor(4.0 * (T / 100.0) ** (2.0 / 9.0)))
        X, _, resid = self._ols_data(T, 2, seed=77)

        se_auto = newey_west_se(resid, X, n_lags=None)
        se_manual = newey_west_se(resid, X, n_lags=expected_lag)

        np.testing.assert_array_equal(se_auto, se_manual)


# ---------------------------------------------------------------------------
# grs_test vs manual closed-form
# ---------------------------------------------------------------------------


class TestGRSFormula:
    """Verify GRS statistic against a manually coded closed-form computation.

    Reference: Gibbons, Ross & Shanken (1989), Econometrica 57(5), Eq. 5:
        GRS = ((T-N-K)/N) * (1 + mu_f' Omega^{-1} mu_f)^{-1} * alpha' Sigma^{-1} alpha
    """

    def _manual_grs(
        self,
        R: np.ndarray,
        F: np.ndarray,
    ) -> float:
        T, N = R.shape
        K = F.shape[1]
        ones = np.ones((T, 1))
        X = np.hstack([ones, F])
        XtXi = np.linalg.inv(X.T @ X)
        alphas = np.zeros(N)
        resid_mat = np.zeros((T, N))
        for i in range(N):
            b = XtXi @ (X.T @ R[:, i])
            alphas[i] = b[0]
            resid_mat[:, i] = R[:, i] - X @ b
        Sigma = resid_mat.T @ resid_mat / (T - K - 1)
        mu_f = F.mean(axis=0)
        F_dm = F - mu_f
        Omega = (F_dm.T @ F_dm) / (T - 1)
        sr_sq = float(mu_f @ np.linalg.inv(Omega) @ mu_f)
        return ((T - N - K) / N) * (1.0 / (1.0 + sr_sq)) * float(
            alphas @ np.linalg.inv(Sigma) @ alphas
        )

    @pytest.mark.parametrize("T,N,K,seed", [
        (120, 5, 2, 7),
        (200, 8, 3, 13),
        (100, 4, 1, 42),
    ])
    def test_statistic_matches_formula(self, T: int, N: int, K: int, seed: int) -> None:
        """GRS statistic must match the textbook closed-form to 1e-10."""
        np.random.seed(seed)
        F_arr = np.random.normal(0.005, 0.04, (T, K))
        R_arr = np.random.normal(0.003, 0.05, (T, N))
        factors = pd.DataFrame(F_arr, columns=[f"f{i+1}" for i in range(K)])
        ports = pd.DataFrame(R_arr, columns=[f"p{i+1}" for i in range(N)])

        result = grs_test(ports, factors)
        manual = self._manual_grs(R_arr, F_arr)

        assert abs(result.statistic - manual) < 1e-10, (
            f"GRS mismatch: ours={result.statistic:.12f}, manual={manual:.12f}"
        )

    def test_zero_alpha_gives_small_statistic(self) -> None:
        """When alphas are near zero, GRS stat should be small and p-value large.

        Construction: returns generated as r = 0 + beta'f + noise with no
        intercept, so OLS-estimated alphas will be near zero.
        """
        np.random.seed(0)
        T, N, K = 300, 5, 2
        F_arr = np.random.normal(0.005, 0.04, (T, K))
        betas = np.random.normal(1.0, 0.3, (N, K))
        # Zero-alpha: pure factor model with small residual noise
        R_arr = (F_arr @ betas.T) + np.random.normal(0, 0.002, (T, N))
        factors = pd.DataFrame(F_arr, columns=["f1", "f2"])
        ports = pd.DataFrame(R_arr, columns=[f"p{i+1}" for i in range(N)])

        result = grs_test(ports, factors)

        # With T=300 and tiny noise, estimated alphas should be near zero
        assert result.p_value > 0.05, (
            f"Expected p_value > 0.05 for near-zero alphas, got {result.p_value:.4f}"
        )
        assert np.all(np.abs(result.alphas.values) < 0.005)

    def test_wide_equals_long_format(self) -> None:
        """Wide and long input formats must produce identical statistics."""
        np.random.seed(55)
        T, N, K = 80, 4, 2
        F_arr = np.random.normal(0.005, 0.04, (T, K))
        R_arr = np.random.normal(0.003, 0.05, (T, N))
        factors_w = pd.DataFrame(F_arr, columns=["f1", "f2"])
        ports_w = pd.DataFrame(R_arr, columns=[f"p{i+1}" for i in range(N)])

        result_wide = grs_test(ports_w, factors_w, date_col=None)

        ports_l = ports_w.copy()
        ports_l["date"] = np.arange(T)
        factors_l = factors_w.copy()
        factors_l["date"] = np.arange(T)
        result_long = grs_test(ports_l, factors_l, date_col="date")

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


# ---------------------------------------------------------------------------
# fama_macbeth lambda vs manual period-by-period OLS
# ---------------------------------------------------------------------------


class TestFMBLambdaFormula:
    """Verify FMB lambda = arithmetic mean of per-period OLS cross-section slopes.

    Reference: Fama & MacBeth (1973), JPE 81(3), 607-636.
    The FM estimator is the time-series average of OLS lambdas from
    T separate cross-sectional regressions.
    """

    def _make_panel(self, T: int, N: int, true_lambda: float, seed: int) -> pd.DataFrame:
        np.random.seed(seed)
        rows = []
        for date in range(T):
            chars = np.random.lognormal(-0.3, 0.5, N)
            rets = true_lambda * chars + np.random.normal(0, 0.05, N)
            for i, p in enumerate(range(1001, 1001 + N)):
                rows.append({"date": date, "permno": p, "ret": rets[i], "bm": chars[i]})
        return pd.DataFrame(rows)

    @pytest.mark.parametrize("T,N,true_lambda,seed", [
        (30, 40, 0.003, 99),
        (50, 60, -0.001, 17),
        (20, 30, 0.005, 5),
    ])
    def test_lambda_equals_manual_period_average(
        self, T: int, N: int, true_lambda: float, seed: int
    ) -> None:
        """FMB lambda must equal arithmetic mean of per-period OLS slopes."""
        data = self._make_panel(T, N, true_lambda, seed)
        result = fama_macbeth(data, ret_col="ret", char_cols=["bm"])

        # Manual: per-period OLS, collect slope for 'bm', average
        manual_lambdas = []
        for date in range(T):
            cs = data[data["date"] == date]
            y = cs["ret"].values
            X = np.column_stack([np.ones(len(y)), cs["bm"].values])
            coef = np.linalg.lstsq(X, y, rcond=None)[0]
            manual_lambdas.append(coef[1])
        manual_mean = float(np.mean(manual_lambdas))

        our_coef = float(result.lambdas.loc["bm", "coef"])
        assert abs(our_coef - manual_mean) < 1e-10, (
            f"FMB lambda mismatch: ours={our_coef:.12f}, manual={manual_mean:.12f}"
        )

    def test_lambda_near_true_value(self) -> None:
        """With large N and T, estimated lambda should be close to the DGP value.

        Uses T=200, N=200 to reduce sampling variance enough for a tight tolerance.
        At small T/N the signal-to-noise ratio (lambda_true / noise_std = 0.08)
        makes the per-period OLS slope noisy; a large sample is required.
        """
        true_lambda = 0.004
        data = self._make_panel(T=200, N=200, true_lambda=true_lambda, seed=42)
        result = fama_macbeth(data, ret_col="ret", char_cols=["bm"])
        our_coef = float(result.lambdas.loc["bm", "coef"])
        assert abs(our_coef - true_lambda) < 0.001, (
            f"FMB lambda far from DGP: ours={our_coef:.5f}, true={true_lambda}"
        )

    def test_nw_se_bounds(self) -> None:
        """NW SE for FM lambda must be positive and less than naive std/sqrt(T)."""
        data = self._make_panel(T=60, N=50, true_lambda=0.003, seed=7)
        result = fama_macbeth(data, ret_col="ret", char_cols=["bm"])

        se_nw = float(result.lambdas.loc["bm", "se_nw"])
        assert se_nw > 0

        # Naive SE = std(lambda_ts) / sqrt(T)
        lambda_ts = result.lambda_ts["bm"].values
        naive_se = float(np.std(lambda_ts, ddof=1) / np.sqrt(len(lambda_ts)))
        # NW SE with lags=0 equals naive SE; with lags>0 it can differ
        # but must remain in a plausible range
        assert se_nw < 5 * naive_se, (
            f"NW SE implausibly large: se_nw={se_nw:.6f}, naive={naive_se:.6f}"
        )


# ---------------------------------------------------------------------------
# pred_HistMean expanding window accuracy
# ---------------------------------------------------------------------------


class TestHistMeanExpandingWindow:
    """Verify that pred_HistMean stores the exact expanding-window training mean.

    Reference: Campbell & Thompson (2008), RFS 21(4), 1509-1531.
    The historical-mean benchmark is the expanding mean of all past
    realized returns available at forecast formation time.
    """

    def test_hist_mean_matches_expanding_training_mean(self) -> None:
        """pred_HistMean at each test date must equal mean(ret[t < test_date])."""
        from eapctf.predict import OLSPredictor, expanding_window_oos

        np.random.seed(42)
        T = 60
        permnos = list(range(1001, 1011))
        rows = []
        for date in range(T):
            for p in permnos:
                rows.append({
                    "date": date, "permno": p,
                    "ret": np.random.normal(0.005, 0.02),
                    "bm": np.random.lognormal(),
                })
        data = pd.DataFrame(rows)

        result = expanding_window_oos(
            data, ret_col="ret", char_cols=["bm"],
            models=[OLSPredictor()],
            train_min_periods=20,
            window_type="expanding",
            train_window=None,
        )
        preds = result.predictions

        mismatches = 0
        for test_date in sorted(preds["date"].unique()):
            train_ret = data[data["date"] < test_date]["ret"]
            expected = float(train_ret.mean())
            actual = float(preds[preds["date"] == test_date]["pred_HistMean"].iloc[0])
            if abs(expected - actual) > 1e-10:
                mismatches += 1

        assert mismatches == 0, (
            f"pred_HistMean mismatch in {mismatches}/{preds['date'].nunique()} test dates"
        )
