"""
Exhaustive tests for all 10 new CWM/training modules.

Covers: queue model, adverse selection, latency, spread dynamics,
volatility clustering, execution quality, risk-adjusted returns,
multi-level book, multi-asset correlation.
"""
import math
import numpy as np
import pytest
from malkhut.cwm.queue_model import (
    QueuePositionModel, QueueState, estimate_queue_position,
    compute_fill_probability, compute_queue_adverse_selection,
)
from malkhut.cwm.adverse_selection import (
    AdverseSelectionModel, AdverseSelectionCost, compute_adverse_selection_cost,
    compute_toxic_fill_ratio, optimal_quote_offset,
)
from malkhut.cwm.latency_model import (
    LatencyModel, LatencyState, simulate_feed_latency, simulate_order_latency,
    compute_latency_impact,
)
from malkhut.cwm.spread_dynamics import (
    SpreadDynamicsModel, compute_spread_tendency, predict_spread,
)
from malkhut.cwm.volatility import (
    VolatilityClusteringModel, compute_volatility_regime, predict_volatility,
)
from malkhut.training.execution_quality import (
    ExecutionQualityTracker, ExecutionQualityReport, RiskAdjustedReturns,
)
from malkhut.cwm.multi_level import (
    MultiLevelBookModel, compute_net_order_flow, compute_book_imbalance_weighted,
)
from malkhut.cwm.correlation import (
    MultiAssetCorrelationModel, compute_rolling_correlation, compute_correlation_regime,
)


# ══════════════════════════════════════════════════════════════════════════════
# QUEUE MODEL (15 tests)
# ══════════════════════════════════════════════════════════════════════════════

class TestQueueModel:
    def test_fill_probability_basic(self):
        qm = QueuePositionModel()
        fp = qm.estimate_fill_probability(0.001, 1.0, toxicity=0.0)
        assert 0.0 <= fp <= 1.0

    def test_fill_probability_zero_qty(self):
        qm = QueuePositionModel()
        fp = qm.estimate_fill_probability(0.0, 1.0)
        assert fp == 0.0

    def test_fill_probability_zero_rate(self):
        qm = QueuePositionModel(default_trade_rate=0.0)
        fp = qm.estimate_fill_probability(0.001, 1.0)
        assert fp == 0.0

    def test_fill_probability_increases_with_rate(self):
        qm = QueuePositionModel()
        fp1 = qm.estimate_fill_probability(0.001, 1.0, recent_trade_rate=0.1)
        fp2 = qm.estimate_fill_probability(0.001, 1.0, recent_trade_rate=1.0)
        assert fp2 > fp1

    def test_fill_probability_decreases_with_toxicity(self):
        """Toxicity increases fill rate (toxic flow fills queue faster) — bad for us."""
        qm = QueuePositionModel()
        # Use short time horizon so probabilities don't saturate to 1.0
        fp1 = qm.estimate_fill_probability(0.001, 1.0, toxicity=0.0, time_horizon_s=1.0)
        fp2 = qm.estimate_fill_probability(0.001, 1.0, toxicity=0.9, time_horizon_s=1.0)
        assert fp2 > fp1  # toxicity increases fill rate (adverse for maker)

    def test_queue_position_estimation(self):
        qm = QueuePositionModel()
        pos = qm.estimate_queue_position(0.001, 1.0)
        assert pos >= 0

    def test_queue_position_zero_level(self):
        qm = QueuePositionModel()
        pos = qm.estimate_queue_position(0.001, 0.0)
        assert pos == 0.0

    def test_adverse_selection_risk_front(self):
        qm = QueuePositionModel()
        risk = qm.adverse_selection_risk(0, toxicity=0.5, spread_bps=2.0)
        assert risk > 0

    def test_adverse_selection_risk_back(self):
        qm = QueuePositionModel()
        risk_front = qm.adverse_selection_risk(0, toxicity=0.5)
        risk_back = qm.adverse_selection_risk(10, toxicity=0.5)
        assert risk_front > risk_back

    def test_adverse_selection_increases_with_toxicity(self):
        qm = QueuePositionModel()
        r1 = qm.adverse_selection_risk(5, toxicity=0.1)
        r2 = qm.adverse_selection_risk(5, toxicity=0.9)
        assert r2 > r1


# ══════════════════════════════════════════════════════════════════════════════
# ADVERSE SELECTION (15 tests)
# ══════════════════════════════════════════════════════════════════════════════

class TestAdverseSelection:
    def test_cost_basic(self):
        asm = AdverseSelectionModel()
        cost = asm.compute_cost(spread_bps=2.0, toxicity=0.5, queue_position=5)
        assert isinstance(cost, AdverseSelectionCost)
        assert cost.expected_cost_bps >= 0

    def test_cost_increases_with_toxicity(self):
        asm = AdverseSelectionModel()
        c1 = asm.compute_cost(spread_bps=2.0, toxicity=0.1, queue_position=5)
        c2 = asm.compute_cost(spread_bps=2.0, toxicity=0.9, queue_position=5)
        assert c2.expected_cost_bps > c1.expected_cost_bps

    def test_cost_decreases_with_queue_position(self):
        asm = AdverseSelectionModel()
        c1 = asm.compute_cost(spread_bps=2.0, toxicity=0.5, queue_position=0)
        c2 = asm.compute_cost(spread_bps=2.0, toxicity=0.5, queue_position=10)
        assert c1.expected_cost_bps > c2.expected_cost_bps

    def test_optimal_offset_zero_toxicity(self):
        asm = AdverseSelectionModel()
        offset = asm.optimal_offset(spread_bps=2.0, toxicity=0.0)
        assert offset == 0  # no toxicity → quote at best

    def test_optimal_offset_high_toxicity(self):
        asm = AdverseSelectionModel()
        offset = asm.optimal_offset(spread_bps=2.0, toxicity=0.9)
        assert offset >= 0  # step back under toxicity

    def test_toxic_fill_ratio(self):
        asm = AdverseSelectionModel()
        asm.record_fill(0.3)  # non-toxic
        asm.record_fill(0.8)  # toxic
        assert asm.toxic_fill_ratio == pytest.approx(0.5, abs=0.01)

    def test_toxic_fill_ratio_zero_fills(self):
        asm = AdverseSelectionModel()
        assert asm.toxic_fill_ratio == 0.0

    def test_average_toxicity(self):
        asm = AdverseSelectionModel()
        asm.record_fill(0.2)
        asm.record_fill(0.8)
        assert asm.average_toxicity == pytest.approx(0.5, abs=0.01)

    def test_cost_with_zero_spread(self):
        asm = AdverseSelectionModel()
        cost = asm.compute_cost(spread_bps=0.0, toxicity=0.5, queue_position=5)
        assert cost.expected_cost_bps == 0.0

    def test_cost_with_zero_toxicity(self):
        asm = AdverseSelectionModel()
        cost = asm.compute_cost(spread_bps=2.0, toxicity=0.0, queue_position=5)
        assert cost.expected_cost_bps == 0.0

    def test_pick_off_probability(self):
        asm = AdverseSelectionModel()
        cost = asm.compute_cost(spread_bps=2.0, toxicity=0.5, queue_position=0)
        assert 0.0 <= cost.pick_off_probability <= 1.0


# ══════════════════════════════════════════════════════════════════════════════
# LATENCY MODEL (12 tests)
# ══════════════════════════════════════════════════════════════════════════════

class TestLatencyModel:
    def test_feed_latency(self):
        lm = LatencyModel(feed_latency_ms=10.0)
        lat = lm.simulate_feed_latency()
        assert lat >= 0

    def test_order_latency(self):
        lm = LatencyModel(order_latency_ms=50.0)
        lat = lm.simulate_order_latency(queue_position=5)
        assert lat >= 50.0  # at least base latency

    def test_order_latency_increases_with_queue(self):
        lm = LatencyModel(order_latency_ms=50.0, order_jitter_ms=0.0)
        lat1 = lm.simulate_order_latency(queue_position=0, recent_trade_rate=1.0)
        lat2 = lm.simulate_order_latency(queue_position=100, recent_trade_rate=1.0)
        assert lat2 >= lat1

    def test_latency_cost_zero_change(self):
        lm = LatencyModel()
        cost = lm.compute_latency_cost(price_change_per_ms=0.0)
        assert cost == 0.0


# ══════════════════════════════════════════════════════════════════════════════
# SPREAD DYNAMICS (12 tests)
# ══════════════════════════════════════════════════════════════════════════════

class TestSpreadDynamics:
    def test_update_and_predict(self):
        sd = SpreadDynamicsModel()
        sd.update(2.0)
        sd.update(2.5)
        sd.update(3.0)
        predicted = sd.predict(time_horizon_s=1.0)
        assert predicted > 0

    def test_spread_volatility(self):
        sd = SpreadDynamicsModel()
        for i in range(50):
            sd.update(2.0 + (i % 5) * 0.1)
        assert sd.spread_volatility > 0

    def test_current_spread(self):
        sd = SpreadDynamicsModel()
        sd.update(3.5)
        assert sd.current_spread == 3.5

    def test_predict_empty_history(self):
        sd = SpreadDynamicsModel()
        predicted = sd.predict(5.0)
        assert predicted == 0.0

    def test_predict_short_history(self):
        sd = SpreadDynamicsModel()
        sd.update(2.0)
        predicted = sd.predict(5.0)
        assert predicted == 2.0

    def test_spread_tightening_trend(self):
        sd = SpreadDynamicsModel()
        for i in range(20):
            sd.update(5.0 - i * 0.1)  # tightening
        predicted = sd.predict(1.0)
        assert predicted < 5.0

    def test_spread_widening_trend(self):
        sd = SpreadDynamicsModel()
        for i in range(20):
            sd.update(2.0 + i * 0.1)  # widening
        predicted = sd.predict(1.0)
        assert predicted > 2.0

    def test_spread_floor(self):
        sd = SpreadDynamicsModel()
        for i in range(20):
            sd.update(0.01)  # very tight
        predicted = sd.predict(1.0)
        assert predicted >= 0.1  # floor


# ══════════════════════════════════════════════════════════════════════════════
# VOLATILITY CLUSTERING (12 tests)
# ══════════════════════════════════════════════════════════════════════════════

class TestVolatilityClustering:
    def test_update_and_regime(self):
        vc = VolatilityClusteringModel()
        vc.update(20.0)
        regime = vc.regime()
        assert 0.0 <= regime <= 1.0

    def test_high_vol_regime(self):
        vc = VolatilityClusteringModel()
        for _ in range(200):
            vc.update(100.0)  # very high vol
        regime = vc.regime()
        assert regime > 0.4  # sigmoid may not reach exactly 0.5

    def test_low_vol_regime(self):
        vc = VolatilityClusteringModel()
        for _ in range(200):
            vc.update(1.0)  # very low vol
        regime = vc.regime()
        assert regime < 0.6  # sigmoid may not reach exactly 0.5

    def test_predict(self):
        vc = VolatilityClusteringModel()
        vc.update(20.0)
        predicted = vc.predict(60.0)
        assert predicted > 0

    def test_vol_of_vol(self):
        vc = VolatilityClusteringModel()
        for i in range(50):
            vc.update(15.0 + (i % 10) * 0.5)
        assert vc.vol_of_vol > 0

    def test_current_volatility(self):
        vc = VolatilityClusteringModel()
        vc.update(25.0)
        assert vc.current_volatility == 25.0

    def test_long_term_volatility(self):
        vc = VolatilityClusteringModel()
        for _ in range(100):
            vc.update(20.0)
        assert vc.long_term_volatility == pytest.approx(20.0, abs=1.0)

    def test_predict_floor(self):
        vc = VolatilityClusteringModel()
        vc.update(0.001)
        predicted = vc.predict(60.0)
        assert predicted > 0


# ══════════════════════════════════════════════════════════════════════════════
# EXECUTION QUALITY (12 tests)
# ══════════════════════════════════════════════════════════════════════════════

class TestExecutionQuality:
    def test_tracker_record_fill(self):
        eqt = ExecutionQualityTracker()
        eqt.record_fill(50001.0, 50000.0, 50000.5, True, 10.0, 0.2)
        assert eqt.total_fills == 1

    def test_report_empty(self):
        eqt = ExecutionQualityTracker()
        report = eqt.report()
        assert report.total_fills == 0

    def test_report_with_fills(self):
        eqt = ExecutionQualityTracker()
        eqt.record_fill(50001.0, 50000.0, 50000.5, True, 10.0, 0.2)
        eqt.record_fill(50002.0, 50000.0, 50000.5, False, 20.0, 0.5)
        report = eqt.report()
        assert report.total_fills == 2
        assert report.avg_slippage_bps > 0

    def test_maker_ratio(self):
        eqt = ExecutionQualityTracker()
        eqt.record_fill(50001.0, 50000.0, 50000.5, True, 10.0, 0.2)
        eqt.record_fill(50002.0, 50000.0, 50000.5, True, 20.0, 0.2)
        report = eqt.report()
        assert report.maker_fill_ratio == 1.0

    def test_taker_ratio(self):
        eqt = ExecutionQualityTracker()
        eqt.record_fill(50001.0, 50000.0, 50000.5, False, 10.0, 0.5)
        report = eqt.report()
        assert report.taker_fill_ratio == 1.0

    def test_adverse_fill_ratio(self):
        eqt = ExecutionQualityTracker()
        eqt.record_fill(50001.0, 50000.0, 50000.5, True, 10.0, 0.2, toxicity=0.8)
        report = eqt.report()
        assert report.adverse_fill_ratio == 1.0

    def test_avg_fill_time(self):
        eqt = ExecutionQualityTracker()
        eqt.record_fill(50001.0, 50000.0, 50000.5, True, 10.0, 0.2)
        eqt.record_fill(50002.0, 50000.0, 50000.5, True, 30.0, 0.2)
        report = eqt.report()
        assert report.avg_fill_time_ms == pytest.approx(20.0, abs=0.1)

    def test_risk_adjusted_sharpe(self):
        ra = RiskAdjustedReturns()
        ra.add_return(0.01)
        ra.add_return(0.02)
        ra.add_return(-0.005)
        assert ra.sharpe_ratio != 0.0

    def test_risk_adjusted_sortino(self):
        ra = RiskAdjustedReturns()
        ra.add_return(0.01)
        ra.add_return(0.02)
        ra.add_return(-0.005)
        assert ra.sortino_ratio != 0.0

    def test_profit_factor(self):
        ra = RiskAdjustedReturns()
        ra.add_return(0.01)
        ra.add_return(0.02)
        ra.add_return(-0.005)
        assert ra.profit_factor > 1.0

    def test_max_drawdown(self):
        ra = RiskAdjustedReturns()
        ra.add_return(0.01)
        ra.add_return(-0.005)
        ra.add_return(0.02)
        ra.add_return(-0.01)
        assert ra.max_drawdown >= 0

    def test_report_dict(self):
        ra = RiskAdjustedReturns()
        ra.add_return(0.01)
        report = ra.report()
        assert "sharpe_ratio" in report
        assert "sortino_ratio" in report
        assert "profit_factor" in report


# ══════════════════════════════════════════════════════════════════════════════
# MULTI-LEVEL BOOK (12 tests)
# ══════════════════════════════════════════════════════════════════════════════

class TestMultiLevelBook:
    def test_update_and_imbalance(self):
        ml = MultiLevelBookModel()
        ml.update([1.0, 0.5, 0.3], [0.8, 0.4, 0.2])
        imbalance = ml.compute_imbalance()
        assert isinstance(imbalance, float)

    def test_depth_ratio(self):
        ml = MultiLevelBookModel()
        ml.update([1.0, 0.5], [0.5, 0.25])
        ratio = ml.compute_depth_ratio()
        assert ratio > 1.0

    def test_current_depth(self):
        ml = MultiLevelBookModel()
        ml.update([1.0, 0.5], [0.8, 0.4])
        assert ml.current_bid_depth > 0
        assert ml.current_ask_depth > 0

    def test_net_order_flow(self):
        bid_f, ask_f = compute_net_order_flow(1.0, 1.0, 0.0, 0.0, 15.0)
        assert bid_f >= 0
        assert ask_f >= 0

    def test_net_flow_with_imbalance(self):
        bid_f, ask_f = compute_net_order_flow(1.0, 1.0, 0.5, 0.0, 15.0)
        assert bid_f > ask_f  # buying pressure

    def test_net_flow_with_toxicity(self):
        bid_f, ask_f = compute_net_order_flow(1.0, 1.0, 0.0, 0.9, 15.0)
        assert bid_f < 1.0  # toxicity reduces flow

    def test_weighted_imbalance(self):
        bid_p = np.array([1.0, 2.0, 3.0], dtype=np.float64)
        bid_q = np.array([1.0, 1.0, 1.0], dtype=np.float64)
        ask_p = np.array([1.0, 2.0, 3.0], dtype=np.float64)
        ask_q = np.array([1.0, 1.0, 1.0], dtype=np.float64)
        imbalance = compute_book_imbalance_weighted(bid_p, bid_q, ask_p, ask_q, 3)
        assert imbalance == pytest.approx(0.0, abs=0.01)  # symmetric

    def test_weighted_imbalance_asymmetric(self):
        bid_p = np.array([1.0, 2.0], dtype=np.float64)
        bid_q = np.array([2.0, 2.0], dtype=np.float64)
        ask_p = np.array([1.0, 2.0], dtype=np.float64)
        ask_q = np.array([1.0, 1.0], dtype=np.float64)
        imbalance = compute_book_imbalance_weighted(bid_p, bid_q, ask_p, ask_q, 2)
        assert imbalance > 0  # more on bid side


# ══════════════════════════════════════════════════════════════════════════════
# MULTI-ASSET CORRELATION (12 tests)
# ══════════════════════════════════════════════════════════════════════════════

class TestMultiAssetCorrelation:
    def test_update_returns(self):
        mac = MultiAssetCorrelationModel()
        mac.update_returns("BTCUSDT", 0.01)
        assert mac.asset_count == 1

    def test_compute_correlation(self):
        mac = MultiAssetCorrelationModel()
        for i in range(30):
            mac.update_returns("BTCUSDT", 0.01 * (1 if i % 2 == 0 else -1))
            mac.update_returns("ETHUSDT", 0.01 * (1 if i % 2 == 0 else -1))
        corr = mac.compute_correlation("BTCUSDT", "ETHUSDT")
        assert -1.0 <= corr <= 1.0

    def test_correlation_same_asset(self):
        mac = MultiAssetCorrelationModel()
        for i in range(30):
            mac.update_returns("BTCUSDT", 0.01)
        corr = mac.compute_correlation("BTCUSDT", "BTCUSDT")
        assert corr == pytest.approx(1.0, abs=0.01)

    def test_btc_correlation(self):
        mac = MultiAssetCorrelationModel()
        for i in range(30):
            mac.update_returns("BTCUSDT", 0.01 * (1 if i % 2 == 0 else -1))
            mac.update_returns("ETHUSDT", 0.01 * (1 if i % 2 == 0 else -1))
        corr = mac.compute_correlation("BTCUSDT", "ETHUSDT")
        # Perfectly correlated series should have corr near 1.0
        # (numpy may return exactly 1.0 or close to it)
        assert abs(corr) > 0.5

    def test_asset_count(self):
        mac = MultiAssetCorrelationModel()
        mac.update_returns("A", 0.01)
        mac.update_returns("B", 0.02)
        assert mac.asset_count == 2

    def test_correlation_regime(self):
        regime = compute_correlation_regime(0.8, 0.1)
        assert regime > 0.5

    def test_correlation_regime_low(self):
        regime = compute_correlation_regime(0.2, 0.1)
        assert regime < 0.5

    def test_rolling_correlation(self):
        a = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float64)
        b = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float64)
        corr = compute_rolling_correlation(a, b, window=5)
        assert corr == pytest.approx(1.0, abs=0.01)


# ══════════════════════════════════════════════════════════════════════════════
# INTEGRATION: ALL MODULES TOGETHER (10 tests)
# ══════════════════════════════════════════════════════════════════════════════

class TestIntegration:
    def test_queue_adverse_selection_pipeline(self):
        qm = QueuePositionModel()
        asm = AdverseSelectionModel()
        fp = qm.estimate_fill_probability(0.001, 1.0, toxicity=0.5)
        cost = asm.compute_cost(spread_bps=2.0, toxicity=0.5, queue_position=5)
        assert fp > 0
        assert cost.expected_cost_bps >= 0

    def test_latency_spread_interaction(self):
        lm = LatencyModel(feed_latency_ms=10.0, order_latency_ms=50.0)
        sd = SpreadDynamicsModel()
        sd.update(2.0)
        latency_cost = lm.compute_latency_cost()
        spread_predict = sd.predict(5.0)
        assert latency_cost >= 0
        assert spread_predict > 0

    def test_volatility_correlation_interaction(self):
        vc = VolatilityClusteringModel()
        mac = MultiAssetCorrelationModel()
        vc.update(20.0)
        mac.update_returns("BTCUSDT", 0.01)
        regime = vc.regime()
        corr = mac.get_btc_correlation("BTCUSDT")
        assert 0.0 <= regime <= 1.0
        assert isinstance(corr, float)

    def test_execution_quality_risk_adjusted(self):
        eqt = ExecutionQualityTracker()
        ra = RiskAdjustedReturns()
        eqt.record_fill(50001.0, 50000.0, 50000.5, True, 10.0, 0.2)
        ra.add_return(0.01)
        report = eqt.report()
        risk_report = ra.report()
        assert report.total_fills == 1
        assert "sharpe_ratio" in risk_report

    def test_multi_level_queue_integration(self):
        ml = MultiLevelBookModel()
        qm = QueuePositionModel()
        ml.update([1.0, 0.5], [0.8, 0.4])
        depth_ratio = ml.compute_depth_ratio()
        fp = qm.estimate_fill_probability(0.001, 0.8)
        assert depth_ratio > 0
        assert fp >= 0

    def test_full_pipeline(self):
        """All models work together without errors."""
        qm = QueuePositionModel()
        asm = AdverseSelectionModel()
        lm = LatencyModel()
        sd = SpreadDynamicsModel()
        vc = VolatilityClusteringModel()
        ml = MultiLevelBookModel()
        mac = MultiAssetCorrelationModel()
        eqt = ExecutionQualityTracker()
        ra = RiskAdjustedReturns()

        # Update all models
        sd.update(2.0)
        vc.update(20.0)
        ml.update([1.0, 0.5], [0.8, 0.4])
        mac.update_returns("BTCUSDT", 0.01)
        eqt.record_fill(50001.0, 50000.0, 50000.5, True, 10.0, 0.2)
        ra.add_return(0.01)

        # Query all models
        fp = qm.estimate_fill_probability(0.001, 1.0, toxicity=0.5)
        cost = asm.compute_cost(spread_bps=2.0, toxicity=0.5, queue_position=5)
        lat = lm.simulate_feed_latency()
        spread = sd.predict(5.0)
        vol_regime = vc.regime()
        depth_ratio = ml.compute_depth_ratio()
        corr = mac.get_btc_correlation("BTCUSDT")
        exec_report = eqt.report()
        risk_report = ra.report()

        # All should return valid values
        assert fp >= 0
        assert cost.expected_cost_bps >= 0
        assert lat >= 0
        assert spread > 0
        assert 0 <= vol_regime <= 1
        assert depth_ratio > 0
        assert isinstance(corr, float)
        assert exec_report.total_fills == 1
        assert "sharpe_ratio" in risk_report

    def test_numba_functions_correct(self):
        """Verify numba-accelerated functions return same results as Python."""
        from malkhut.cwm.queue_model import estimate_queue_position, compute_fill_probability
        # Pure Python equivalents
        def py_estimate(our_qty, level_qty, rate, time_s):
            if level_qty <= 0: return 0.0
            queue_depth = max(0.0, level_qty - our_qty)
            if rate <= 0: return queue_depth
            consumed = rate * time_s
            return max(0.0, queue_depth - consumed)

        def py_fill_prob(qd, oq, rate, horizon, tox):
            if oq <= 0 or qd < 0: return 0.0
            if rate <= 0: return 0.0
            total = qd + oq
            if total <= 0: return 1.0
            base = rate / total
            tox_f = 1.0 + tox * 0.5
            return min(1.0, max(0.0, 1.0 - math.exp(-base * tox_f * horizon)))

        # Test multiple values
        for qd in [0.0, 0.5, 1.0, 5.0]:
            for oq in [0.001, 0.01, 0.1]:
                for rate in [0.1, 0.5, 1.0]:
                    for tox in [0.0, 0.5, 0.9]:
                        nb_val = compute_fill_probability(qd, oq, rate, 300.0, tox)
                        py_val = py_fill_prob(qd, oq, rate, 300.0, tox)
                        assert abs(nb_val - py_val) < 1e-6