from fxhnt.domain.equity_backtest import (
    universe_asof,
    month_end_rebalance_dates,
    epoch_day,
    pit_universe_by_month,
)


# --- PIT (point-in-time) universe selection — the lookahead-bias fix --------------

def _ym(year: int, month: int) -> int:
    return year * 12 + (month - 1)


def test_pit_universe_excludes_not_yet_liquid_names():
    # THE decisive lookahead test. LATE is illiquid early, hugely liquid late;
    # EARLY is liquid from the start. With top_n=1, an EARLY rebalance must pick
    # EARLY (NOT LATE) even though LATE's PEAK/future dollar-volume is far higher —
    # proving the selection uses only data <= the rebalance date (no lookahead).
    # A LATE rebalance, once LATE's trailing dv overtakes, picks LATE.
    monthly: list[tuple[str, int, float, int, float]] = []
    for m in range(1, 13):  # 2000-01 .. 2000-12
        ym = _ym(2000, m)
        # LATE: tiny early dv (1e3), huge late dv (1e10); EARLY: steady high dv (1e8)
        late_dv = 1e3 if m <= 6 else 1e10
        monthly.append(("LATE", ym, late_dv, 21, 50.0))
        monthly.append(("EARLY", ym, 1e8, 21, 50.0))

    rebal_early = _ym(2000, 3)   # only months 1..3 visible
    rebal_late = _ym(2000, 12)   # months 1..12 visible, LATE's late dv dominates

    early_uni = pit_universe_by_month(
        monthly, [rebal_early], window_months=3, top_n=1,
        min_price=1.0, min_bars=2,
    )
    late_uni = pit_universe_by_month(
        monthly, [rebal_late], window_months=3, top_n=1,
        min_price=1.0, min_bars=2,
    )
    assert early_uni[rebal_early] == ["EARLY"]   # NO lookahead — LATE not yet liquid
    assert late_uni[rebal_late] == ["LATE"]      # LATE now overtakes


def test_pit_universe_min_price_and_min_bars_filters():
    monthly = [
        # PENNY: high dv but avg_close below min_price -> excluded
        ("PENNY", _ym(2000, 1), 1e9, 21, 0.50),
        ("PENNY", _ym(2000, 2), 1e9, 21, 0.50),
        # THIN: high dv but too few in-window bars -> excluded
        ("THIN", _ym(2000, 1), 1e9, 5, 100.0),
        ("THIN", _ym(2000, 2), 1e9, 5, 100.0),
        # GOOD: passes both
        ("GOOD", _ym(2000, 1), 5e7, 21, 100.0),
        ("GOOD", _ym(2000, 2), 5e7, 21, 100.0),
    ]
    uni = pit_universe_by_month(
        monthly, [_ym(2000, 2)], window_months=2, top_n=10,
        min_price=1.0, min_bars=20,
    )
    assert uni[_ym(2000, 2)] == ["GOOD"]


def test_pit_universe_trailing_window_and_tiebreak():
    # Trailing mean over the window: name X has dv 100 then 200 over a 2-month
    # window -> trailing mean 150. TIE: two names with identical trailing dv must
    # order by symbol ascending.
    monthly = [
        ("X", _ym(2000, 1), 100.0, 21, 100.0),
        ("X", _ym(2000, 2), 200.0, 21, 100.0),     # trailing mean 150
        ("ZZZ", _ym(2000, 1), 150.0, 21, 100.0),
        ("ZZZ", _ym(2000, 2), 150.0, 21, 100.0),   # trailing mean 150 -> tie with X
        ("AAA", _ym(2000, 1), 150.0, 21, 100.0),
        ("AAA", _ym(2000, 2), 150.0, 21, 100.0),   # trailing mean 150 -> tie with X
    ]
    uni = pit_universe_by_month(
        monthly, [_ym(2000, 2)], window_months=2, top_n=3,
        min_price=1.0, min_bars=2,
    )
    # all three tie at 150 trailing -> deterministic symbol-ascending order
    assert uni[_ym(2000, 2)] == ["AAA", "X", "ZZZ"]


def test_pit_universe_window_only_uses_months_in_range():
    # window_months=2 ending at ym=2000-03 considers months {02, 03} only — a huge
    # dv in 2000-01 (out of window) must NOT count toward the trailing mean, and a
    # name only liquid in 2000-01 with no in-window bars is excluded.
    monthly = [
        ("OLD", _ym(2000, 1), 1e12, 21, 100.0),   # huge but out of window -> excluded
        ("CUR", _ym(2000, 2), 100.0, 21, 100.0),
        ("CUR", _ym(2000, 3), 100.0, 21, 100.0),
    ]
    uni = pit_universe_by_month(
        monthly, [_ym(2000, 3)], window_months=2, top_n=10,
        min_price=1.0, min_bars=2,
    )
    assert uni[_ym(2000, 3)] == ["CUR"]

MEMBERS = [
    ("LIVE", "NASDAQ", "2000-01-01", "2026-06-01"),
    ("DEAD", "NYSE", "1998-01-01", "2005-03-15"),
    ("IPO", "NASDAQ", "2010-06-01", "2026-06-01"),
]


def test_universe_asof_excludes_not_yet_listed_and_delisted():
    dollar_vol = {"LIVE": 1e9, "DEAD": 1e9, "IPO": 1e9}
    assert universe_asof(MEMBERS, dollar_vol, "2008-01-01", n=10) == ["LIVE"]


def test_universe_asof_includes_delisted_while_alive():
    dollar_vol = {"LIVE": 2e9, "DEAD": 1e9}
    assert universe_asof(MEMBERS, dollar_vol, "2003-01-01", n=10) == ["LIVE", "DEAD"]


def test_universe_asof_ranks_by_dollar_volume_desc_and_takes_top_n():
    members = [(s, "NYSE", "2000-01-01", "2026-06-01") for s in ("A", "B", "C")]
    dollar_vol = {"A": 10.0, "B": 30.0, "C": 20.0}
    assert universe_asof(members, dollar_vol, "2020-01-01", n=2) == ["B", "C"]


def test_universe_asof_drops_zero_or_missing_dollar_volume():
    members = [(s, "NYSE", "2000-01-01", "2026-06-01") for s in ("A", "B")]
    assert universe_asof(members, {"A": 5.0, "B": 0.0}, "2020-01-01", n=10) == ["A"]


def test_month_end_rebalance_dates_one_per_month():
    dates = ["2020-01-02", "2020-01-31", "2020-02-03", "2020-02-27", "2020-03-02"]
    assert month_end_rebalance_dates(dates) == ["2020-01-31", "2020-02-27", "2020-03-02"]


def test_epoch_day_roundtrip():
    assert epoch_day("1970-01-01") == 0
    assert epoch_day("1970-01-02") == 1


import pytest

from fxhnt.domain.equity_backtest import daily_returns_for_weights


def test_daily_returns_long_only_matches_weighted_price_change():
    weights = {"A": 0.5, "B": 0.5}
    closes = {
        "A": {0: 100.0, 1: 110.0},
        "B": {0: 50.0, 1: 50.0},
    }
    rets = daily_returns_for_weights(weights, closes, [0, 1], borrow_daily=0.0)
    assert rets == pytest.approx([0.05])


def test_daily_returns_applies_borrow_on_short_notional():
    weights = {"A": -1.0}
    closes = {"A": {0: 100.0, 1: 100.0}}
    rets = daily_returns_for_weights(weights, closes, [0, 1], borrow_daily=0.001)
    assert rets == [-0.001]


def test_daily_returns_skips_day_with_missing_price():
    weights = {"A": 1.0}
    closes = {"A": {0: 100.0, 2: 121.0}}   # day 1 missing
    rets = daily_returns_for_weights(weights, closes, [0, 1, 2], borrow_daily=0.0)
    assert rets == []


def test_daily_returns_holds_flat_on_implausible_ratio():
    # FIX B: a present-but-implausible single-day ratio (1.0 -> 1e6 = +99,999,900%)
    # is a data error (e.g. test ticker / bad tick / unhandled corporate action).
    # The blown name must be DROPPED from that day's contribution (hold-flat), so
    # the day's book return is 0.0 (NOT 999999.0). GOOD has a flat legit day.
    weights = {"BAD": 1.0}
    closes = {"BAD": {0: 1.0, 1: 1.0, 2: 1e6}}
    rets = daily_returns_for_weights(weights, closes, [0, 1, 2], borrow_daily=0.0)
    # day 0->1 is flat (0.0); day 1->2 blows up -> BAD dropped -> empty -> step skipped
    assert rets == pytest.approx([0.0])
    assert all(abs(r) <= 10.0 for r in rets)


def test_daily_returns_keeps_legit_large_move():
    # A legit +50% day is WITHIN [-0.9, 10.0] and must contribute normally.
    weights = {"A": 1.0}
    closes = {"A": {0: 100.0, 1: 150.0}}
    rets = daily_returns_for_weights(weights, closes, [0, 1], borrow_daily=0.0)
    assert rets == pytest.approx([0.5])


def test_daily_returns_drops_below_min_day_ret():
    # A -95% day is below the -90% floor -> dropped (data error), step skipped.
    weights = {"A": 1.0}
    closes = {"A": {0: 100.0, 1: 5.0}}
    rets = daily_returns_for_weights(weights, closes, [0, 1], borrow_daily=0.0)
    assert rets == []


def test_daily_returns_bad_name_dropped_good_name_kept_same_day():
    # On a day where BAD blows up but GOOD has a legit move, BAD is dropped and
    # GOOD still contributes (the day is NOT skipped because GOOD remains).
    weights = {"BAD": 0.5, "GOOD": 0.5}
    closes = {
        "BAD": {0: 1.0, 1: 1e6},
        "GOOD": {0: 100.0, 1: 150.0},
    }
    rets = daily_returns_for_weights(weights, closes, [0, 1], borrow_daily=0.0)
    # only GOOD contributes: 0.5 * 0.5 = 0.25
    assert rets == pytest.approx([0.25])


def test_delisting_books_terminal_loss_for_long():
    # FIX T9: a long held name whose adjclose series TERMINATES at its delist day
    # books a one-time terminal delisting return (Shumway -0.30) instead of being
    # silently held flat. DEAD has data at day0,day1 only; delist_day says day1 is
    # its delisting day. Step (0,1) is a normal +10%; step (1,2) has price@1 not@2
    # and delist_day==1 -> book 1.0 * -0.30 = -0.30.
    weights = {"DEAD": 1.0}
    closes = {"DEAD": {0: 100.0, 1: 110.0}}
    rets = daily_returns_for_weights(
        weights, closes, [0, 1, 2], borrow_daily=0.0,
        delist_day={"DEAD": 1}, delisting_return=-0.30,
    )
    assert rets == pytest.approx([0.10, -0.30])


def test_delisting_books_gain_for_short():
    # A SHORT (weight<0) profits when a name craters to its delisting return:
    # -1.0 * -0.30 = +0.30 booked on the terminal step.
    weights = {"DEAD": -1.0}
    closes = {"DEAD": {0: 100.0, 1: 110.0}}
    rets = daily_returns_for_weights(
        weights, closes, [0, 1, 2], borrow_daily=0.0,
        delist_day={"DEAD": 1}, delisting_return=-0.30,
    )
    # step (0,1): short of a +10% move = -0.10; step (1,2): terminal +0.30
    assert rets[-1] == pytest.approx(0.30)


def test_no_delisting_when_not_flagged_holds_flat():
    # With no delist_day flag, the terminating series is held flat (existing
    # behavior): step (1,2) has no priced DEAD and no terminal -> step dropped.
    weights = {"DEAD": 1.0}
    closes = {"DEAD": {0: 100.0, 1: 110.0}}
    rets = daily_returns_for_weights(
        weights, closes, [0, 1, 2], borrow_daily=0.0,
        delist_day=None, delisting_return=-0.30,
    )
    assert rets == pytest.approx([0.10])


def test_delisting_return_default_zero_is_noop():
    # delisting_return=0.0 -> terminal contribution 0.0; the terminal step emits
    # 0.0 (a delisting flagged with a zero return is the biased hold-flat baseline
    # value, not a dropped step). The first step is the normal +10%.
    weights = {"DEAD": 1.0}
    closes = {"DEAD": {0: 100.0, 1: 110.0}}
    rets = daily_returns_for_weights(
        weights, closes, [0, 1, 2], borrow_daily=0.0,
        delist_day={"DEAD": 1}, delisting_return=0.0,
    )
    assert rets == pytest.approx([0.10, 0.0])
