from aitrader.ai.ensemble import combine, single
from aitrader.ai.schema import TradeDecision


def _d(action, conf=0.8, size=0.5):
    return TradeDecision(action=action, confidence=conf, suggested_size_pct=size, reasoning="")


def test_consensus_buy_above_threshold():
    r = combine(_d("BUY", 0.8, 0.6), _d("BUY", 0.7, 0.4), min_conf=0.6)
    assert r.action == "BUY"
    assert r.suggested_size_pct == 0.4
    assert r.confidence == 0.7
    assert r.rationale == "consensus_buy"


def test_consensus_sell():
    r = combine(_d("SELL", 0.9, 0.3), _d("SELL", 0.65, 0.5), min_conf=0.6)
    assert r.action == "SELL"


def test_below_threshold_holds():
    r = combine(_d("BUY", 0.8, 0.5), _d("BUY", 0.4, 0.5), min_conf=0.6)
    assert r.action == "HOLD"
    assert r.rationale == "consensus_below_threshold"


def test_disagreement_holds():
    r = combine(_d("BUY", 0.9, 0.5), _d("SELL", 0.9, 0.5), min_conf=0.6)
    assert r.action == "HOLD"
    assert r.rationale == "disagreement"


def test_hold_consensus_holds():
    r = combine(_d("HOLD", 0.9, 0.0), _d("HOLD", 0.9, 0.0), min_conf=0.6)
    assert r.action == "HOLD"


def test_single_buy():
    r = single(_d("BUY", 0.8, 0.5), min_conf=0.6)
    assert r.action == "BUY"
    assert r.rationale == "single_buy"


def test_single_below_threshold():
    r = single(_d("BUY", 0.4, 0.5), min_conf=0.6)
    assert r.action == "HOLD"