"""Correctness tests for the signal layer."""
import numpy as np
import pandas as pd

from src.signals import (
    daily_returns,
    max_signal,
    monthly_returns,
    month_valid_counts,
)


def _one_month_returns(values, ticker="AAA"):
    """Build a single-ticker daily-return frame for Jan 2021 business days."""
    idx = pd.bdate_range("2021-01-01", periods=len(values))
    return pd.DataFrame({ticker: values}, index=idx)


def test_max_signal_is_mean_of_top_n():
    vals = [0.01, 0.05, 0.02, 0.10, -0.03, 0.07, 0.00, 0.04,
            0.06, -0.01, 0.03, 0.08, 0.02, 0.01, 0.09, 0.05]
    rets = _one_month_returns(vals)
    # MAX(1) = single largest daily return.
    m1 = max_signal(rets, n_high=1, min_days=10).iloc[0, 0]
    assert np.isclose(m1, max(vals))
    # MAX(3) = mean of the three largest (0.10, 0.09, 0.08).
    m3 = max_signal(rets, n_high=3, min_days=10).iloc[0, 0]
    assert np.isclose(m3, np.mean([0.10, 0.09, 0.08]))


def test_max_signal_respects_min_days():
    rets = _one_month_returns([0.02, 0.03, 0.05])  # only 3 days
    out = max_signal(rets, n_high=1, min_days=15)
    assert np.isnan(out.iloc[0, 0])


def test_max_signal_index_is_month_end():
    rets = _one_month_returns([0.01] * 20)
    out = max_signal(rets, n_high=1, min_days=10)
    ts = out.index[0]
    assert ts.month == 1 and ts.year == 2021
    # last day of January
    assert ts == pd.Timestamp("2021-01-31")


def test_monthly_returns_compound_correctly():
    vals = [0.10, -0.05, 0.02]
    rets = _one_month_returns(vals)
    out = monthly_returns(rets).iloc[0, 0]
    expected = (1.10 * 0.95 * 1.02) - 1.0
    assert np.isclose(out, expected)


def test_daily_returns_recovers_known_path():
    prices = pd.DataFrame(
        {"AAA": [100.0, 110.0, 99.0]},
        index=pd.bdate_range("2021-01-01", periods=3),
    )
    r = daily_returns(prices)["AAA"].dropna()
    assert np.isclose(r.iloc[0], 0.10)
    assert np.isclose(r.iloc[1], -0.10)


def test_month_valid_counts_ignores_nan():
    df = _one_month_returns([0.01, np.nan, 0.02, np.nan, 0.03])
    counts = month_valid_counts(df).iloc[0, 0]
    assert counts == 3
