from __future__ import annotations

import datetime as dt

from libs.backtest.daily_gainers import (
    DailyGainerCriteria,
    collect_daily_gainer_signals,
    evaluate_daily_gainer_signals,
)


def _bar(
    open_price: float,
    high: float,
    low: float,
    close: float,
    volume: int,
) -> dict[str, float | int]:
    return {
        "open": open_price,
        "high": high,
        "low": low,
        "close": close,
        "volume": volume,
    }


def test_collect_daily_gainer_signals_ranks_by_return_then_liquidity() -> None:
    start = dt.date(2024, 1, 1)
    dates = [start + dt.timedelta(days=i) for i in range(24)]

    bars_by_symbol: dict[str, dict[dt.date, dict[str, float | int]]] = {
        "AAA": {},
        "BBB": {},
        "CCC": {},
    }
    for idx, date in enumerate(dates[:-1]):
        base_price = 100.0 + idx
        for symbol in bars_by_symbol:
            bars_by_symbol[symbol][date] = _bar(base_price, base_price + 1, base_price - 1, base_price, 1_000_000)

    signal_date = dates[-1]
    bars_by_symbol["AAA"][signal_date] = _bar(120.0, 128.0, 118.0, 126.0, 4_000_000)
    bars_by_symbol["BBB"][signal_date] = _bar(121.0, 130.0, 119.0, 129.0, 3_000_000)
    bars_by_symbol["CCC"][signal_date] = _bar(119.0, 125.0, 112.0, 116.0, 4_000_000)

    criteria = DailyGainerCriteria(
        min_change=0.03,
        max_change=0.35,
        min_volume_ratio=1.2,
        min_close_location=0.55,
        min_avg_dollar_volume=50_000_000.0,
        top_n=2,
    )
    signals_by_date = collect_daily_gainer_signals(
        bars_by_symbol=bars_by_symbol,
        start_date=signal_date,
        end_date=signal_date,
        criteria=criteria,
    )

    signals = signals_by_date[signal_date]
    assert [signal.symbol for signal in signals] == ["BBB", "AAA"]
    assert [signal.rank for signal in signals] == [1, 2]


def test_evaluate_daily_gainer_signals_uses_next_open_entry() -> None:
    signal_date = dt.date(2024, 2, 1)
    entry_date = dt.date(2024, 2, 2)
    day3 = dt.date(2024, 2, 5)
    day5 = dt.date(2024, 2, 7)

    bars_by_symbol = {
        "AAA": {
            dt.date(2024, 1, 2): _bar(100.0, 101.0, 99.0, 100.0, 1_000_000),
            dt.date(2024, 1, 3): _bar(101.0, 102.0, 100.0, 101.0, 1_000_000),
            dt.date(2024, 1, 4): _bar(102.0, 103.0, 101.0, 102.0, 1_000_000),
            dt.date(2024, 1, 5): _bar(103.0, 104.0, 102.0, 103.0, 1_000_000),
            dt.date(2024, 1, 8): _bar(104.0, 105.0, 103.0, 104.0, 1_000_000),
            dt.date(2024, 1, 9): _bar(105.0, 106.0, 104.0, 105.0, 1_000_000),
            dt.date(2024, 1, 10): _bar(106.0, 107.0, 105.0, 106.0, 1_000_000),
            dt.date(2024, 1, 11): _bar(107.0, 108.0, 106.0, 107.0, 1_000_000),
            dt.date(2024, 1, 12): _bar(108.0, 109.0, 107.0, 108.0, 1_000_000),
            dt.date(2024, 1, 16): _bar(109.0, 110.0, 108.0, 109.0, 1_000_000),
            dt.date(2024, 1, 17): _bar(110.0, 111.0, 109.0, 110.0, 1_000_000),
            dt.date(2024, 1, 18): _bar(111.0, 112.0, 110.0, 111.0, 1_000_000),
            dt.date(2024, 1, 19): _bar(112.0, 113.0, 111.0, 112.0, 1_000_000),
            dt.date(2024, 1, 22): _bar(113.0, 114.0, 112.0, 113.0, 1_000_000),
            dt.date(2024, 1, 23): _bar(114.0, 115.0, 113.0, 114.0, 1_000_000),
            dt.date(2024, 1, 24): _bar(115.0, 116.0, 114.0, 115.0, 1_000_000),
            dt.date(2024, 1, 25): _bar(116.0, 117.0, 115.0, 116.0, 1_000_000),
            dt.date(2024, 1, 26): _bar(117.0, 118.0, 116.0, 117.0, 1_000_000),
            dt.date(2024, 1, 29): _bar(118.0, 119.0, 117.0, 118.0, 1_000_000),
            dt.date(2024, 1, 30): _bar(119.0, 120.0, 118.0, 119.0, 1_000_000),
            dt.date(2024, 1, 31): _bar(120.0, 121.0, 119.0, 120.0, 1_000_000),
            signal_date: _bar(123.0, 130.0, 122.0, 128.0, 4_000_000),
            entry_date: _bar(129.0, 132.0, 128.0, 131.0, 3_000_000),
            day3: _bar(130.0, 136.0, 129.0, 135.0, 3_000_000),
            dt.date(2024, 2, 6): _bar(134.0, 138.0, 133.0, 137.0, 3_000_000),
            day5: _bar(137.0, 140.0, 136.0, 139.0, 3_000_000),
        }
    }

    signals_by_date = collect_daily_gainer_signals(
        bars_by_symbol=bars_by_symbol,
        start_date=signal_date,
        end_date=signal_date,
        criteria=DailyGainerCriteria(
            min_change=0.05,
            max_change=0.30,
            min_volume_ratio=1.2,
            min_close_location=0.55,
            min_avg_dollar_volume=50_000_000.0,
            top_n=1,
        ),
    )

    trades = evaluate_daily_gainer_signals(
        signals_by_date=signals_by_date,
        bars_by_symbol=bars_by_symbol,
    )

    assert len(trades) == 1
    trade = trades[0]
    assert trade.entry_date == entry_date
    assert trade.entry_price == 129.0
    assert trade.entry_gap_pct == 129.0 / 128.0 - 1.0
    assert trade.hold_1d_return == 131.0 / 129.0 - 1.0
    assert trade.hold_3d_return == 137.0 / 129.0 - 1.0
    assert trade.hold_5d_return is None