TradeBot/tests/test_calibrate_thresholds.py

from __future__ import annotations

from types import SimpleNamespace

from tools.calibrate_torch_thresholds import (
    CalibrationResult,
    ForecastRecord,
    _average_selected_predictions,
    _apply_platt_calibration,
    _choose_recommendation,
    _full_backtest,
    _fit_platt_calibration,
    _entry_validation_skill,
)
from tools.train_torch_recurrent_forecaster import _ensemble_candidate


def _result(*, trades: int, average: float, total: float, profit_factor: float) -> CalibrationResult:
    return CalibrationResult(
        edge=0.05,
        probability=0.52,
        confidence=0.4,
        trades=trades,
        wins=max(0, trades // 2),
        win_rate=0.5,
        total_net_percent=total,
        average_net_percent=average,
        max_drawdown_percent=1.0,
        profit_factor=profit_factor,
        score=1.0,
    )


def _record(index: int, probability: float, future: float) -> ForecastRecord:
    return ForecastRecord(
        symbol="BTCUSDT",
        index=index,
        timestamp=index,
        close=100.0,
        high=101.0,
        low=99.0,
        next_open=100.0,
        next_timestamp=index + 1,
        atr=1.0,
        expected_percent=0.1,
        probability_up=probability,
        confidence=0.5,
        skill=0.1,
        q50_percent=0.1,
        block_entry=False,
        future_net_percent=future,
        benchmark_entry=False,
        benchmark_exit=False,
    )


def test_calibration_does_not_fallback_to_too_few_trades() -> None:
    selected = _choose_recommendation(
        [_result(trades=1, average=2.0, total=2.0, profit_factor=999.0)],
        min_trades=30,
    )

    assert selected is None


def test_calibration_selects_only_viable_result() -> None:
    viable = _result(trades=30, average=0.2, total=6.0, profit_factor=1.4)

    assert _choose_recommendation([viable], min_trades=30) is viable


def test_platt_calibration_learns_probability_direction_from_train_records() -> None:
    records = [
        _record(index, 0.8 if index % 2 else 0.2, -1.0 if index % 2 else 1.0)
        for index in range(100)
    ]

    calibration = _fit_platt_calibration(records)
    calibrated = _apply_platt_calibration(
        [_record(101, 0.8, -1.0), _record(102, 0.2, 1.0)],
        calibration,
    )

    assert calibration["slope"] < 0
    assert calibrated[0].probability_up < calibrated[1].probability_up


def test_entry_quality_never_falls_back_to_holdout_skill() -> None:
    entry = {"validation_skill": 0.12, "skill": 0.99, "holdout_skill": 0.99}

    assert _entry_validation_skill(entry) == 0.12
    assert _entry_validation_skill({"skill": 0.99, "holdout_skill": 0.99}) == 0.0


def test_batched_ensemble_averages_decoded_predictions() -> None:
    averaged = _average_selected_predictions(
        [
            {"expected_return": 0.01, "q50": 0.02, "probability_up": 0.6},
            {"expected_return": 0.03, "q50": 0.04, "probability_up": 0.8},
        ]
    )

    assert averaged == {
        "expected_return": 0.02,
        "q50": 0.03,
        "probability_up": 0.7,
    }


def test_multi_seed_export_does_not_duplicate_first_member_weights() -> None:
    members = [
        {
            "validation_mae": 0.1,
            "state_dict": {"weight": [seed]},
            "head_weight": [[seed]],
            "head_bias": [seed],
        }
        for seed in (7, 19)
    ]

    exported = _ensemble_candidate(members, [7, 19])

    assert exported["ensemble_size"] == 2
    assert exported["ensemble_seeds"] == [7, 19]
    assert len(exported["ensemble_members"]) == 2
    assert "state_dict" not in exported
    assert "head_weight" not in exported


def test_single_seed_export_keeps_only_top_level_weights() -> None:
    exported = _ensemble_candidate(
        [
            {
                "validation_mae": 0.1,
                "state_dict": {"weight": [7]},
                "head_weight": [[7]],
                "head_bias": [7],
            }
        ],
        [7],
    )

    assert exported["ensemble_size"] == 1
    assert exported["state_dict"] == {"weight": [7]}
    assert "ensemble_members" not in exported


def test_full_backtest_never_uses_global_threshold_for_ineligible_symbol() -> None:
    btc = [_record(index, 0.8, 1.0) for index in range(3)]
    eth = [_record(index, 0.8, 1.0) for index in range(3)]
    for record in eth:
        record.symbol = "ETHUSDT"
    thresholds = _result(trades=3, average=1.0, total=3.0, profit_factor=999.0)

    replay = _full_backtest(
        btc + eth,
        thresholds,
        horizon=3,
        round_trip_cost=0.0,
        settings=SimpleNamespace(
            stop_loss_percent=0.04,
            take_profit_percent=0.035,
            stop_loss_exit_enabled=True,
            atr_trailing_multiplier=2.2,
        ),
        symbol_thresholds={"BTCUSDT": thresholds},
        require_symbol_thresholds=True,
    )

    assert {row["symbol"] for row in replay["symbol_breakdown"]} == {"BTCUSDT"}