"""
Unit tests for the Kelly quant engine.
"""
import numpy as np
from builderbrain.quant_engine import KellyEngine, CorrelationMatrix, MarketEdge


def test_correlation_matrix_build():
    markets = [
        MarketEdge("trump_2024", "Will Trump win?", "politics", "YES", 0.08, 0.55, 0.63, 50000, "2024-11-05"),
        MarketEdge("musk_doge", "Will Musk lead DOGE?", "politics", "YES", 0.05, 0.60, 0.65, 30000, "2024-12-01"),
        MarketEdge("btc_100k", "BTC > $100k?", "crypto", "YES", 0.10, 0.40, 0.50, 100000, "2024-12-31"),
    ]
    corr = CorrelationMatrix()
    matrix = corr.build(markets)
    assert matrix.shape == (3, 3)
    assert np.allclose(np.diag(matrix), 1.0)
    # Politics-politics should be high
    assert matrix[0, 1] > 0.5
    # Politics-crypto should be low
    assert matrix[0, 2] < 0.2
    print("✅ test_correlation_matrix_build passed")


def test_kelly_engine_basic():
    engine = KellyEngine(bankroll_usd=10000)
    markets = [
        MarketEdge("m1", "Test 1", "sports", "YES", 0.10, 0.40, 0.50, 10000, "2024-12-01"),
        MarketEdge("m2", "Test 2", "sports", "NO", 0.05, 0.55, 0.50, 10000, "2024-12-01"),
    ]
    positions = engine.size_positions(markets)
    assert len(positions) > 0
    total = sum(p.fraction_of_bankroll for p in positions)
    assert total <= engine.max_leverage + 1e-6
    for p in positions:
        assert p.fraction_of_bankroll <= 0.25 + 1e-6
    print("✅ test_kelly_engine_basic passed")


def test_kelly_no_viable_edges():
    engine = KellyEngine(bankroll_usd=10000, min_edge=0.50)
    markets = [
        MarketEdge("m1", "Test", "sports", "YES", 0.01, 0.49, 0.50, 10000, "2024-12-01"),
    ]
    positions = engine.size_positions(markets)
    assert len(positions) == 0
    print("✅ test_kelly_no_viable_edges passed")


if __name__ == "__main__":
    test_correlation_matrix_build()
    test_kelly_engine_basic()
    test_kelly_no_viable_edges()
    print("\n🎉 All quant engine tests passed")