"""Tests for factor aggregation, regression and the MAX/IVOL correlation."""
import numpy as np
import pandas as pd

from src.factors import (
    daily_factors_to_monthly,
    factor_regression,
    max_ivol_correlation,
)


def test_daily_to_monthly_compounds():
    idx = pd.bdate_range("2021-01-01", "2021-01-31")
    daily = pd.DataFrame({"MKT_RF": [0.001] * len(idx)}, index=idx)
    monthly = daily_factors_to_monthly(daily)
    expected = (1.001) ** len(idx) - 1.0
    assert np.isclose(monthly.iloc[0, 0], expected)


def test_factor_regression_recovers_alpha_and_beta():
    rng = np.random.default_rng(0)
    n = 240
    idx = pd.period_range("2000-01", periods=n, freq="M").to_timestamp("M")
    mkt = pd.Series(rng.normal(0.005, 0.04, n), index=idx)
    true_alpha, true_beta = 0.004, 0.6
    y = true_alpha + true_beta * mkt + pd.Series(rng.normal(0, 0.001, n), index=idx)

    factors = pd.DataFrame({"MKT_RF": mkt})
    out = factor_regression(y, factors, factor_cols=["MKT_RF"])
    assert np.isclose(out["alpha_monthly"], true_alpha, atol=5e-4)
    assert np.isclose(out["table"].loc["MKT_RF", "coef"], true_beta, atol=2e-2)
    assert out["alpha_tstat"] > 2  # strong alpha by construction


def test_max_ivol_correlation_positive_when_linked():
    rng = np.random.default_rng(1)
    months = pd.period_range("2020-01", periods=24, freq="M").to_timestamp("M")
    cols = [f"S{i}" for i in range(20)]
    ivol = pd.DataFrame(rng.uniform(0.01, 0.05, (24, 20)), index=months, columns=cols)
    # MAX strongly increasing in IVOL plus light noise -> high positive corr.
    mx = 2.0 * ivol + rng.normal(0, 0.001, ivol.shape)
    corr = max_ivol_correlation(mx, ivol)
    assert corr["pooled_pearson"] > 0.8
    assert corr["avg_cross_sectional"] > 0.8
