okx-codex-trader/tests/test_explore_ultrashort.py

import importlib.util
import sys
from pathlib import Path

import pytest

from okx_codex_trader.models import Candle
from okx_codex_trader.sampled_report import SegmentResult


def load_explore_module():
    path = Path(__file__).resolve().parents[1] / "scripts" / "explore_ultrashort.py"
    spec = importlib.util.spec_from_file_location("explore_ultrashort", path)
    assert spec is not None
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def build_candles(count: int) -> list[Candle]:
    return [
        Candle(
            symbol="BTC-USDT-SWAP",
            ts=index * 60_000,
            open=100.0 + index,
            high=101.0 + index,
            low=99.0 + index,
            close=100.0 + index,
            volume=1_000.0 + index,
        )
        for index in range(count)
    ]


def build_result(
    total_return: float,
    trade_count: int,
    win_rate: float,
    max_drawdown: float,
    trade_returns: list[float] | None = None,
    equity_curve: list[dict[str, float | int]] | None = None,
) -> SegmentResult:
    return SegmentResult(
        trade_count=trade_count,
        total_return=total_return,
        win_rate=win_rate,
        max_drawdown=max_drawdown,
        trades=[
            {
                "side": "Long",
                "entry_time": "2026-04-01 00:00",
                "exit_time": "2026-04-01 00:15",
                "entry_price": 100.0,
                "exit_price": 100.0 * (1.0 + trade_return),
                "pnl": trade_return * 10_000.0,
                "return_pct": trade_return * 100.0,
            }
            for trade_return in (trade_returns or [])
        ],
        open_position=None,
        candles=[],
        equity_curve=equity_curve or [],
        entries=[],
        exits=[],
    )


def test_evaluate_candidate_all_windows_uses_complete_non_overlapping_windows():
    module = load_explore_module()
    returns = iter([0.01, -0.02, 0.03])
    calls: list[dict[str, object]] = []

    def run_segment(*, candles, leverage, warmup_bars):
        calls.append(
            {
                "first_ts": candles[0].ts,
                "count": len(candles),
                "leverage": leverage,
                "warmup_bars": warmup_bars,
            }
        )
        return build_result(
            next(returns),
            trade_count=2,
            win_rate=0.5,
            max_drawdown=0.04,
            trade_returns=[0.02, -0.01],
        )

    metrics = module.evaluate_candidate_all_windows(
        candidate=module.Candidate("test-candidate", warmup_bars=2, run=run_segment),
        candles=build_candles(17),
        window_size=5,
        leverage=3,
    )

    assert calls == [
        {"first_ts": 0, "count": 7, "leverage": 3, "warmup_bars": 2},
        {"first_ts": 300_000, "count": 7, "leverage": 3, "warmup_bars": 2},
        {"first_ts": 600_000, "count": 7, "leverage": 3, "warmup_bars": 2},
    ]
    assert metrics["sample_count"] == 3
    assert metrics["avg_return"] == pytest.approx(0.0066666667)
    assert metrics["median_return"] == pytest.approx(0.01)
    assert metrics["positive_window_rate"] == pytest.approx(2 / 3)
    assert metrics["worst_return"] == pytest.approx(-0.02)
    assert metrics["p10_return"] == pytest.approx(-0.014)
    assert metrics["p90_return"] == pytest.approx(0.026)
    assert metrics["best_return"] == pytest.approx(0.03)
    assert metrics["trades"] == 6
    assert metrics["avg_trades_per_window"] == pytest.approx(2.0)
    assert metrics["win_rate"] == pytest.approx(0.5)
    assert metrics["trade_win_rate"] == pytest.approx(0.5)
    assert metrics["avg_trade_return"] == pytest.approx(0.005)
    assert metrics["avg_win_return"] == pytest.approx(0.02)
    assert metrics["avg_loss_return_abs"] == pytest.approx(0.01)
    assert metrics["payoff_ratio"] == pytest.approx(2.0)
    assert metrics["profit_factor"] == pytest.approx(2.0)
    assert metrics["expectancy_per_trade"] == pytest.approx(0.005)
    assert metrics["max_drawdown"] == pytest.approx(0.04)
    assert metrics["return_drawdown_ratio"] == pytest.approx(0.0066666667 / 0.04)
    assert metrics["ci95_low"] < metrics["avg_return"] < metrics["ci95_high"]


def test_evaluate_candidate_window_rows_keeps_window_timestamps():
    module = load_explore_module()

    def run_segment(*, candles, leverage, warmup_bars):
        return build_result(
            0.01,
            trade_count=1,
            win_rate=1.0,
            max_drawdown=0.02,
            trade_returns=[0.01],
        )

    rows = module.evaluate_candidate_window_rows(
        candidate=module.Candidate("test-candidate", warmup_bars=2, run=run_segment),
        candles=build_candles(12),
        window_size=5,
        leverage=3,
    )

    assert rows[0]["window_start_ts"] == 120_000
    assert rows[0]["window_end_ts"] == 360_000
    assert rows[1]["window_start_ts"] == 420_000
    assert rows[1]["window_end_ts"] == 660_000
    assert module.summarize_window_rows(rows)["sample_count"] == 2


def test_sort_robust_results_prioritizes_confidence_interval_lower_bound():
    module = load_explore_module()
    pandas = pytest.importorskip("pandas")
    frame = pandas.DataFrame(
        [
            {"name": "higher-average", "avg_return": 0.10, "ci95_low": -0.01},
            {"name": "supported", "avg_return": 0.03, "ci95_low": 0.01},
            {"name": "tie-breaker", "avg_return": 0.04, "ci95_low": 0.01},
        ]
    )

    sorted_frame = module.sort_robust_results(frame)

    assert sorted_frame["name"].tolist() == ["tie-breaker", "supported", "higher-average"]


def test_add_cost_metrics_and_sort_cost_results():
    module = load_explore_module()
    pandas = pytest.importorskip("pandas")
    frame = pandas.DataFrame(
        [
            {"name": "active", "avg_return": 0.01, "ci95_low": 0.004, "ci95_high": 0.016, "avg_trades_per_window": 10.0},
            {"name": "selective", "avg_return": 0.006, "ci95_low": 0.003, "ci95_high": 0.009, "avg_trades_per_window": 2.0},
        ]
    )

    with_cost = module.add_cost_metrics(frame, 0.0012)
    sorted_frame = module.sort_cost_results(with_cost)

    assert with_cost.loc[0, "net_avg_return"] == pytest.approx(-0.002)
    assert with_cost.loc[0, "net_ci95_low"] == pytest.approx(-0.008)
    assert with_cost.loc[0, "breakeven_roundtrip_cost_on_margin"] == pytest.approx(0.001)
    assert with_cost.loc[1, "net_avg_return"] == pytest.approx(0.0036)
    assert with_cost.loc[1, "net_ci95_low"] == pytest.approx(0.0006)
    assert sorted_frame["name"].tolist() == ["selective", "active"]


def test_summarize_periods_applies_trade_cost_to_period_average():
    module = load_explore_module()
    pandas = pytest.importorskip("pandas")
    frame = pandas.DataFrame(
        [
            {"window_end_ts": 0, "total_return": 0.01, "trade_count": 2, "win_rate": 0.5, "max_drawdown": 0.01},
            {"window_end_ts": 86_400_000, "total_return": -0.02, "trade_count": 1, "win_rate": 0.0, "max_drawdown": 0.03},
        ]
    )

    monthly = module.summarize_periods(frame, "M", 0.001)

    assert monthly.loc[0, "period"] == "1970-01"
    assert monthly.loc[0, "window_count"] == 2
    assert monthly.loc[0, "avg_return"] == pytest.approx(-0.005)
    assert monthly.loc[0, "positive_window_rate"] == pytest.approx(0.5)
    assert monthly.loc[0, "trades"] == 3
    assert monthly.loc[0, "net_avg_return"] == pytest.approx(-0.0065)


def test_summarize_cost_adjusted_trade_equity_periods_compounds_net_trade_returns():
    module = load_explore_module()
    result = build_result(
        0.0,
        trade_count=2,
        win_rate=0.5,
        max_drawdown=0.0,
        trade_returns=[0.01, -0.02],
        equity_curve=[{"ts": 1_775_001_600_000, "equity": 10_000.0, "close": 100.0}],
    )

    monthly = module.summarize_cost_adjusted_trade_equity_periods(result, "M", 0.001)

    assert monthly.loc[0, "period"] == "2026-04"
    assert monthly.loc[0, "trades"] == 2
    assert monthly.loc[0, "end_equity"] == pytest.approx(10_000.0 * 1.009 * 0.979)


def test_add_market_regime_columns_adds_net_return_and_buckets():
    module = load_explore_module()
    rows = [
        {"window_start_ts": 300 * 60_000, "window_end_ts": 359 * 60_000, "total_return": 0.02, "trade_count": 2, "win_rate": 0.5, "max_drawdown": 0.01, "trades": []},
        {"window_start_ts": 360 * 60_000, "window_end_ts": 419 * 60_000, "total_return": -0.01, "trade_count": 1, "win_rate": 0.0, "max_drawdown": 0.02, "trades": []},
        {"window_start_ts": 420 * 60_000, "window_end_ts": 479 * 60_000, "total_return": 0.03, "trade_count": 3, "win_rate": 1.0, "max_drawdown": 0.01, "trades": []},
    ]

    frame = module.add_market_regime_columns(build_candles(500), rows, 0.001)

    assert frame.loc[0, "net_return"] == pytest.approx(0.018)
    assert set(frame["market_return_bucket"].astype(str)) <= {"down", "flat", "up"}
    assert "realized_vol" in frame
    assert "ma240_distance" in frame


def test_annualized_metrics_from_cost_adjusted_equity():
    module = load_explore_module()
    pandas = pytest.importorskip("pandas")
    frame = pandas.DataFrame(
        [
            {"ts": pandas.Timestamp("2025-01-01", tz="UTC"), "equity": 100.0},
            {"ts": pandas.Timestamp("2025-01-02", tz="UTC"), "equity": 120.0},
            {"ts": pandas.Timestamp("2025-01-03", tz="UTC"), "equity": 90.0},
            {"ts": pandas.Timestamp("2026-01-01", tz="UTC"), "equity": 110.0},
        ]
    )

    metrics = module.annualized_metrics_from_equity(frame, 1_735_689_600_000, 1_767_225_600_000)

    assert metrics["net_total_return"] == pytest.approx(0.10)
    assert metrics["net_annualized_return"] == pytest.approx(0.10)
    assert metrics["net_max_drawdown"] == pytest.approx(0.25)
    assert metrics["net_calmar"] == pytest.approx(0.4)
    assert "net_sharpe_daily" in metrics


def test_recent_horizon_metrics_use_equity_at_cutoff():
    module = load_explore_module()
    pandas = pytest.importorskip("pandas")
    frame = pandas.DataFrame(
        [
            {"ts": pandas.Timestamp("2024-01-01 00:00", tz="UTC"), "equity": 100.0},
            {"ts": pandas.Timestamp("2025-06-01 00:00", tz="UTC"), "equity": 120.0},
            {"ts": pandas.Timestamp("2025-12-01 00:00", tz="UTC"), "equity": 90.0},
            {"ts": pandas.Timestamp("2026-04-01 00:00", tz="UTC"), "equity": 110.0},
        ]
    )

    rows = module.recent_horizon_metrics_from_equity(
        frame,
        int(pandas.Timestamp("2026-04-01 00:00", tz="UTC").timestamp() * 1000),
        (("1y", pandas.DateOffset(years=1)), ("3m", pandas.DateOffset(months=3))),
    )

    assert rows.loc[0, "horizon"] == "1y"
    assert rows.loc[0, "horizon_start"] == "2025-04-01 00:00"
    assert rows.loc[0, "net_total_return"] == pytest.approx(110.0 / 100.0 - 1.0)
    assert rows.loc[1, "horizon"] == "3m"
    assert rows.loc[1, "horizon_start"] == "2026-01-01 00:00"
    assert rows.loc[1, "net_total_return"] == pytest.approx(110.0 / 90.0 - 1.0)


def test_build_ma_cross_candidate_names_and_warmup():
    module = load_explore_module()

    candidate = module.build_ma_cross_candidate(20, 80, "long")

    assert candidate.name == "ma-cross-long-f20-s80"
    assert candidate.warmup_bars == 80


def test_build_rsi2_long_guarded_candidate_names_and_warmup():
    module = load_explore_module()

    candidate = module.build_rsi2_long_guarded_candidate(240, 2.0, 55.0, 0.008, 96)

    assert candidate.name == "rsi2-long-guarded-t240-l2.0-x55.0-sl0.008-mh96"
    assert candidate.warmup_bars == 240


def test_build_trend_rsi_bb_long_candidate_names_and_warmup():
    module = load_explore_module()

    candidate = module.build_trend_rsi_bb_long_candidate(240, 20, 2.5, 3.0, 45.0, 0.005)

    assert candidate.name == "trend-rsi-bb-long-t240-b20-m2.5-r3.0-x45.0-sl0.005"
    assert candidate.warmup_bars == 240


def test_build_regime_hybrid_candidate_names_and_warmup():
    module = load_explore_module()

    candidate = module.build_regime_hybrid_candidate(240, 240, 0.015, 2.0, 55.0, 2.5, 0.008)

    assert candidate.name == "regime-hybrid-t240-r240-n0.015-l2.0-x55.0-m2.5-sl0.008"
    assert candidate.warmup_bars == 240


def test_align_pair_candles_keeps_shared_timestamps():
    module = load_explore_module()
    left = build_candles(4)
    right = [
        Candle(symbol="BTC-USDT-SWAP", ts=60_000, open=1.0, high=1.0, low=1.0, close=1.0, volume=1.0),
        Candle(symbol="BTC-USDT-SWAP", ts=180_000, open=3.0, high=3.0, low=3.0, close=3.0, volume=1.0),
    ]

    left_aligned, right_aligned = module.align_pair_candles(left, right)

    assert [candle.ts for candle in left_aligned] == [60_000, 180_000]
    assert [candle.ts for candle in right_aligned] == [60_000, 180_000]


def test_build_eth_btc_rsi_filter_candidate_names_and_warmup():
    module = load_explore_module()

    candidate = module.build_eth_btc_rsi_filter_candidate(50, 3.0, 55.0, 240, 96, 0.01)

    assert candidate.name == "eth-btc-rsi-filter-et50-l3.0-x55.0-bt240-bm96-br0.01"
    assert candidate.warmup_bars == 240


def test_build_eth_btc_shock_filter_candidate_names_and_warmup():
    module = load_explore_module()

    candidate = module.build_eth_btc_shock_filter_candidate(50, 3.0, 55.0, 480, 240, 0.0, 240, 0.006, 0.05)

    assert candidate.name == "eth-btc-shock-filter-et50-l3.0-x55.0-bt480-bm240-br0.0-sw240-sv0.006-sd0.05"
    assert candidate.warmup_bars == 480


def test_build_eth_btc_ratio_pullback_candidate_names_and_warmup():
    module = load_explore_module()

    candidate = module.build_eth_btc_ratio_pullback_candidate(480, 96, 0.01, 48, 2.0, 5.0, 0.008)

    assert candidate.name == "eth-btc-ratio-pullback-bt480-bm96-br0.01-rl48-rs2.0-rr5.0-sl0.008"
    assert candidate.warmup_bars == 480


def test_build_btc_lead_eth_lag_candidate_names_and_warmup():
    module = load_explore_module()

    candidate = module.build_btc_lead_eth_lag_candidate(16, 0.018, 0.012, 32, 0.006, 0.012)

    assert candidate.name == "btc-lead-eth-lag-lb16-br0.018-gap0.012-mh32-sl0.006-tp0.012"
    assert candidate.warmup_bars == 16


def test_btc_lead_eth_lag_enters_after_btc_outperformance():
    module = load_explore_module()
    eth = build_candles(8)
    btc = build_candles(8)
    btc[3] = Candle(symbol="BTC-USDT-SWAP", ts=180_000, open=100.0, high=100.0, low=100.0, close=103.0, volume=1.0)
    btc[4] = Candle(symbol="BTC-USDT-SWAP", ts=240_000, open=103.0, high=103.0, low=103.0, close=103.0, volume=1.0)
    eth[3] = Candle(symbol="ETH-USDT-SWAP", ts=180_000, open=100.0, high=100.0, low=100.0, close=100.0, volume=1.0)
    eth[4] = Candle(symbol="ETH-USDT-SWAP", ts=240_000, open=100.0, high=102.0, low=100.0, close=101.0, volume=1.0)

    result = module.run_btc_lead_eth_lag_segment(
        eth_candles=eth,
        btc_candles=btc,
        leverage=3,
        warmup_bars=3,
        lead_lookback=3,
        btc_return_threshold=0.02,
        lag_gap=0.02,
        max_hold_bars=3,
        stop_loss_pct=0.01,
        take_profit_pct=0.01,
    )

    assert result.trade_count == 1
    assert result.trades[0]["side"] == "Long"
    assert result.total_return > 0.0


def test_get_candles_cached_saves_exhausted_history_and_updates_latest(tmp_path):
    module = load_explore_module()

    class Client:
        def __init__(self):
            self.limits: list[int] = []

        def get_candles(self, symbol, bar, limit):
            self.limits.append(limit)
            if len(self.limits) == 1:
                return build_candles(3)
            return [
                Candle(
                    symbol=symbol,
                    ts=3 * 60_000,
                    open=103.0,
                    high=104.0,
                    low=102.0,
                    close=103.0,
                    volume=1_003.0,
                )
            ]

    client = Client()

    first = module.get_candles_cached(client, "BTC-USDT-SWAP", "15m", 500, tmp_path)
    second = module.get_candles_cached(client, "BTC-USDT-SWAP", "15m", 500, tmp_path)

    assert [candle.ts for candle in first] == [0, 60_000, 120_000]
    assert [candle.ts for candle in second] == [0, 60_000, 120_000, 180_000]
    assert client.limits == [500, 300]
    assert (tmp_path / "BTC-USDT-SWAP" / "15m.csv").exists()
    assert module.load_cached_candles(tmp_path, "BTC-USDT-SWAP", "15m")[1] is True


def test_get_candles_cached_bridges_stale_cache_gap(tmp_path):
    module = load_explore_module()
    module.save_cached_candles(
        tmp_path,
        "BTC-USDT-SWAP",
        "3m",
        [
            Candle(symbol="BTC-USDT-SWAP", ts=0, open=100.0, high=101.0, low=99.0, close=100.0, volume=1.0),
            Candle(symbol="BTC-USDT-SWAP", ts=180_000, open=101.0, high=102.0, low=100.0, close=101.0, volume=1.0),
        ],
        history_exhausted=True,
    )

    class Client:
        def __init__(self):
            self.limits: list[int] = []

        def get_candles(self, symbol, bar, limit):
            self.limits.append(limit)
            if limit == 300:
                return [
                    Candle(symbol=symbol, ts=720_000, open=104.0, high=105.0, low=103.0, close=104.0, volume=1.0),
                    Candle(symbol=symbol, ts=900_000, open=105.0, high=106.0, low=104.0, close=105.0, volume=1.0),
                ]
            return [
                Candle(symbol=symbol, ts=360_000, open=102.0, high=103.0, low=101.0, close=102.0, volume=1.0),
                Candle(symbol=symbol, ts=540_000, open=103.0, high=104.0, low=102.0, close=103.0, volume=1.0),
                Candle(symbol=symbol, ts=720_000, open=104.0, high=105.0, low=103.0, close=104.0, volume=1.0),
                Candle(symbol=symbol, ts=900_000, open=105.0, high=106.0, low=104.0, close=105.0, volume=1.0),
            ]

    candles = module.get_candles_cached(Client(), "BTC-USDT-SWAP", "3m", 10, tmp_path)

    assert [candle.ts for candle in candles] == [0, 180_000, 360_000, 540_000, 720_000, 900_000]


def test_get_candles_cached_repairs_existing_internal_gap(tmp_path):
    module = load_explore_module()
    module.save_cached_candles(
        tmp_path,
        "BTC-USDT-SWAP",
        "3m",
        [
            Candle(symbol="BTC-USDT-SWAP", ts=0, open=100.0, high=101.0, low=99.0, close=100.0, volume=1.0),
            Candle(symbol="BTC-USDT-SWAP", ts=180_000, open=101.0, high=102.0, low=100.0, close=101.0, volume=1.0),
            Candle(symbol="BTC-USDT-SWAP", ts=720_000, open=104.0, high=105.0, low=103.0, close=104.0, volume=1.0),
            Candle(symbol="BTC-USDT-SWAP", ts=900_000, open=105.0, high=106.0, low=104.0, close=105.0, volume=1.0),
        ],
        history_exhausted=True,
    )

    class Client:
        def get_candles(self, symbol, bar, limit):
            if limit == 300:
                return [
                    Candle(symbol=symbol, ts=720_000, open=104.0, high=105.0, low=103.0, close=104.0, volume=1.0),
                    Candle(symbol=symbol, ts=900_000, open=105.0, high=106.0, low=104.0, close=105.0, volume=1.0),
                ]
            return [
                Candle(symbol=symbol, ts=360_000, open=102.0, high=103.0, low=101.0, close=102.0, volume=1.0),
                Candle(symbol=symbol, ts=540_000, open=103.0, high=104.0, low=102.0, close=103.0, volume=1.0),
                Candle(symbol=symbol, ts=720_000, open=104.0, high=105.0, low=103.0, close=104.0, volume=1.0),
                Candle(symbol=symbol, ts=900_000, open=105.0, high=106.0, low=104.0, close=105.0, volume=1.0),
            ]

    candles = module.get_candles_cached(Client(), "BTC-USDT-SWAP", "3m", 10, tmp_path)

    assert [candle.ts for candle in candles] == [0, 180_000, 360_000, 540_000, 720_000, 900_000]


def test_history_bars_for_years_counts_minute_bars():
    module = load_explore_module()

    assert module.history_bars_for_years("15m", 10.0) == 350_400
    assert module.history_bars_for_years("3m", 10.0) == 1_752_000


def test_build_strategy_timeframe_candidates_uses_fixed_strategy_set():
    module = load_explore_module()

    candidates = module.build_strategy_timeframe_candidates()

    assert [candidate.name for candidate in candidates] == [
        "bbmr-default",
        "bbsb-default",
        "donchian-e12-x6-s0.008",
        "rsi2-t50-l3.0-s97.0",
        "ema-pullback-f13-s34-b0.006",
        "range-momo-l10-tp0.006-sl0.004",
        "vwap-revert-w72-z2.0-sl0.006",
    ]