import numpy as np
import pandas as pd
import pytest

from quantcore.measurement import estimators as est


@pytest.fixture
def r():
    rng = np.random.default_rng(3)
    x = rng.standard_normal((600, 4)) @ np.array([[1, .5, .2, 0], [0, 1, .3, 0], [0, 0, 1, .4], [0, 0, 0, 1]]) * 0.01
    return pd.DataFrame(x, index=pd.bdate_range("2020-01-01", periods=600), columns=list("ABCD"))


def test_rolling_cov_matches_numpy(r):
    cov = est.rolling_cov(r, 100)
    np.testing.assert_allclose(cov.values, np.cov(r.values[-100:], rowvar=False, ddof=1))
    corr = est.rolling_corr(r, 100)
    np.testing.assert_allclose(np.diag(corr.values), 1.0)
    np.testing.assert_allclose(corr.values, np.corrcoef(r.values[-100:], rowvar=False), atol=1e-12)


def test_ewm_cov_diagonal_matches_pandas(r):
    lam = 0.94
    cov = est.ewm_cov(r, lam, max_obs=len(r))
    pd_var = r.ewm(alpha=1 - lam, adjust=True).var(bias=True).iloc[-1]
    np.testing.assert_allclose(np.diag(cov.values), pd_var.values, rtol=1e-8)
    assert est.ewm_weights(10, lam).sum() == pytest.approx(1.0)


def test_ewm_sequence_recursion_converges_to_direct(r):
    lam = 0.97
    last = None
    for date, cov_prev, _ in est.covariance_sequence(r, "ewm", lam=lam, warmup=100):
        last = (date, cov_prev)
    # Sigma_{T-1} from the recursion vs. the direct (normalised-weight) estimate at T-1
    direct = est.ewm_cov(r.iloc[:-1], lam, max_obs=len(r))
    assert last[0] == r.index[-1]
    np.testing.assert_allclose(last[1], direct.values, rtol=0.05, atol=1e-7)


def test_rolling_sequence_matches_cov(r):
    seq = list(est.covariance_sequence(r, "rolling", window=50))
    date, cov_prev, ret = seq[-1]
    assert date == r.index[-1]
    np.testing.assert_allclose(cov_prev, np.cov(r.values[-51:-1], rowvar=False, ddof=1))
    np.testing.assert_allclose(ret, r.values[-1])


def test_shrinkage_keeps_variances(r):
    cov = est.rolling_cov(r, 200)
    sh = est.shrink(cov, 0.3)
    np.testing.assert_allclose(np.diag(sh.values), np.diag(cov.values))
    assert abs(sh.values[0, 1]) == pytest.approx(0.7 * abs(cov.values[0, 1]))