"""
Tests for CAAE Cost Model.

Validates cost predictions against experimental results from Experiments 1-7.
"""

import pytest
from caae.cost_model import CostModel, CostEstimate


class TestCostModel:
    """Test suite for CostModel class."""

    @pytest.fixture
    def cost_model(self):
        """Create a cost model instance for testing."""
        return CostModel()

    # =========================================================================
    # Swap Cost Tests
    # =========================================================================

    def test_swap_cost_gen4_25k_tokens(self, cost_model):
        """
        Test: Swap cost for 25k tokens on Gen 4 hardware.

        From Experiment 5, Phase 4:
        - 25k tokens on Gen 4 (16 GB/s) should have swap cost ~10.1ms
        """
        swap_cost = cost_model.estimate_swap_cost(
            context_size_tokens=25000, bandwidth_gbps=16.0, queue_depth_ms=0.0
        )

        # Allow tolerance from experimental value
        assert 3.0 <= swap_cost <= 6.0, f"Expected ~4.1ms, got {swap_cost:.2f}ms"

    def test_swap_cost_gen5_100k_tokens(self, cost_model):
        """
        Test: Swap cost for 100k tokens on Gen 5 hardware.

        From Experiment 6:
        - 100k tokens on Gen 5 (128 GB/s) should have swap cost ~1.3ms
        """
        swap_cost = cost_model.estimate_swap_cost(
            context_size_tokens=100000, bandwidth_gbps=128.0, queue_depth_ms=0.0
        )

        # Allow tolerance
        assert 0.5 <= swap_cost <= 3.0, f"Expected ~1.3ms, got {swap_cost:.2f}ms"

    def test_swap_cost_includes_queue_depth(self, cost_model):
        """Test: Queue depth is added to swap cost."""
        base_cost = cost_model.estimate_swap_cost(
            context_size_tokens=10000, bandwidth_gbps=16.0, queue_depth_ms=0.0
        )

        with_queue = cost_model.estimate_swap_cost(
            context_size_tokens=10000, bandwidth_gbps=16.0, queue_depth_ms=5.0
        )

        assert with_queue - base_cost == pytest.approx(
            5.0, abs=0.1
        ), "Queue depth should add directly to swap cost"

    # =========================================================================
    # Recompute Cost Tests
    # =========================================================================

    def test_recompute_cost_25k_tokens(self, cost_model):
        """
        Test: Recompute cost for 25k tokens.

        From Experiment 5:
        - 25k tokens should have recompute cost ~122.7ms
        """
        recompute_cost = cost_model.estimate_recompute_cost(context_size_tokens=25000)

        # Allow tolerance for default layer/seq parameters
        assert (
            180.0 <= recompute_cost <= 280.0
        ), f"Expected ~230.1ms, got {recompute_cost:.2f}ms"

    def test_recompute_cost_early_layer_more_expensive(self, cost_model):
        """Test: Early layers (layer_depth=0) cost more to recompute."""
        early_layer = cost_model.estimate_recompute_cost(
            context_size_tokens=10000, layer_depth=0
        )

        late_layer = cost_model.estimate_recompute_cost(
            context_size_tokens=10000, layer_depth=79
        )

        assert (
            early_layer > late_layer
        ), "Early layers should have higher recompute cost"

    # =========================================================================
    # Full Estimate Tests
    # =========================================================================

    def test_estimate_returns_correct_structure(self, cost_model):
        """Test: Estimate returns CostEstimate with all fields."""
        estimate = cost_model.estimate(context_size_tokens=25000, bandwidth_gbps=16.0)

        assert isinstance(estimate, CostEstimate)
        assert estimate.swap_cost_ms > 0
        assert estimate.recompute_cost_ms > 0
        assert estimate.optimal_action in ("swap", "recompute")
        assert 0 <= estimate.confidence <= 1

    def test_estimate_gen4_25k_recommends_swap(self, cost_model):
        """
        Test: For 25k tokens on Gen 4, swap should be recommended.

        From Experiment 5: Swap (10.1ms) << Recompute (122.7ms)
        """
        estimate = cost_model.estimate(context_size_tokens=25000, bandwidth_gbps=16.0)

        assert (
            estimate.optimal_action == "swap"
        ), f"Expected swap, got {estimate.optimal_action}"
        assert (
            estimate.cost_ratio > 5.0
        ), f"Expected ratio > 5, got {estimate.cost_ratio}"

    def test_estimate_confidence_97_percent(self, cost_model):
        """
        Test: Model confidence should be ~97.2% (from Experiment 2).
        """
        estimate = cost_model.estimate(context_size_tokens=25000, bandwidth_gbps=16.0)

        assert (
            estimate.confidence >= 0.95
        ), f"Expected confidence >= 95%, got {estimate.confidence:.1%}"

    # =========================================================================
    # Crossover Threshold Tests
    # =========================================================================

    def test_crossover_threshold_gen4(self, cost_model):
        """
        Test: Gen 4 crossover threshold should be ~10k tokens.
        """
        threshold = cost_model.get_crossover_threshold(bandwidth_gbps=16.0)

        assert 5000 <= threshold <= 15000, f"Expected ~10k tokens, got {threshold}"

    def test_crossover_threshold_gen5(self, cost_model):
        """
        Test: Gen 5 crossover threshold should be ~20k tokens.
        """
        threshold = cost_model.get_crossover_threshold(bandwidth_gbps=128.0)

        assert 15000 <= threshold <= 25000, f"Expected ~20k tokens, got {threshold}"


class TestCostModelEdgeCases:
    """Edge case tests for CostModel."""

    @pytest.fixture
    def cost_model(self):
        return CostModel()

    def test_minimum_swap_cost(self, cost_model):
        """Test: Swap cost has minimum floor."""
        swap_cost = cost_model.estimate_swap_cost(
            context_size_tokens=1, bandwidth_gbps=128.0
        )

        assert swap_cost >= cost_model.MIN_SWAP_MS

    def test_minimum_recompute_cost(self, cost_model):
        """Test: Recompute cost has minimum floor."""
        recompute_cost = cost_model.estimate_recompute_cost(context_size_tokens=1)

        assert recompute_cost >= cost_model.MIN_RECOMPUTE_MS

    def test_large_context_handled(self, cost_model):
        """Test: Large context sizes (1M tokens) don't cause overflow."""
        estimate = cost_model.estimate(context_size_tokens=1000000, bandwidth_gbps=16.0)

        assert estimate.swap_cost_ms > 0
        assert estimate.recompute_cost_ms > 0
        assert not float("inf") == estimate.swap_cost_ms


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
