"""Tests for PredictionEvaluator.

This module tests the prediction accuracy evaluation framework.
"""

from __future__ import annotations

import json
import tempfile
from datetime import datetime
from pathlib import Path

import pytest

from sagellm_control.predictor.evaluator import (
    EvaluationMetrics,
    PredictionEvaluator,
    PredictionRecord,
)


class TestPredictionRecord:
    """Tests for PredictionRecord."""

    def test_initialization(self) -> None:
        """Test basic initialization."""
        record = PredictionRecord(
            request_id="req1",
            predicted_lifetime=10.0,
            actual_lifetime=12.0,
            predicted_at=datetime.now(),
            completed_at=datetime.now(),
        )
        assert record.request_id == "req1"
        assert record.predicted_lifetime == 10.0
        assert record.actual_lifetime == 12.0

    def test_absolute_error(self) -> None:
        """Test absolute error calculation."""
        record = PredictionRecord(
            request_id="req1",
            predicted_lifetime=10.0,
            actual_lifetime=12.0,
            predicted_at=datetime.now(),
            completed_at=datetime.now(),
        )
        assert abs(record.absolute_error - 2.0) < 0.001

    def test_relative_error(self) -> None:
        """Test relative error calculation."""
        record = PredictionRecord(
            request_id="req1",
            predicted_lifetime=10.0,
            actual_lifetime=12.0,
            predicted_at=datetime.now(),
            completed_at=datetime.now(),
        )
        # Error = |10-12|/12 * 100 = 16.67%
        assert abs(record.relative_error - 16.67) < 0.1

    def test_relative_error_zero_actual(self) -> None:
        """Test relative error when actual is zero."""
        record = PredictionRecord(
            request_id="req1",
            predicted_lifetime=10.0,
            actual_lifetime=0.0,
            predicted_at=datetime.now(),
            completed_at=datetime.now(),
        )
        assert record.relative_error == float("inf")

    def test_relative_error_both_zero(self) -> None:
        """Test relative error when both are zero."""
        record = PredictionRecord(
            request_id="req1",
            predicted_lifetime=0.0,
            actual_lifetime=0.0,
            predicted_at=datetime.now(),
            completed_at=datetime.now(),
        )
        assert record.relative_error == 0.0

    def test_is_accurate(self) -> None:
        """Test accuracy check."""
        # Accurate prediction (error < 20%)
        record1 = PredictionRecord(
            request_id="req1",
            predicted_lifetime=10.0,
            actual_lifetime=11.0,
            predicted_at=datetime.now(),
            completed_at=datetime.now(),
        )
        assert record1.relative_error < 20.0

        # Inaccurate prediction (error >= 20%)
        record2 = PredictionRecord(
            request_id="req2",
            predicted_lifetime=10.0,
            actual_lifetime=15.0,
            predicted_at=datetime.now(),
            completed_at=datetime.now(),
        )
        assert record2.relative_error >= 20.0


class TestEvaluationMetrics:
    """Tests for EvaluationMetrics."""

    def test_initialization(self) -> None:
        """Test default initialization."""
        metrics = EvaluationMetrics()
        assert metrics.mae == 0.0
        assert metrics.rmse == 0.0
        assert metrics.mape == 0.0
        assert metrics.num_predictions == 0

    def test_to_dict(self) -> None:
        """Test conversion to dictionary."""
        metrics = EvaluationMetrics(
            mae=1.5,
            rmse=2.0,
            mape=10.0,
            accuracy=75.0,
            num_predictions=100,
        )
        d = metrics.to_dict()
        assert d["mae"] == 1.5
        assert d["rmse"] == 2.0
        assert d["mape"] == 10.0
        assert d["accuracy"] == 75.0
        assert d["num_predictions"] == 100


class TestPredictionEvaluator:
    """Tests for PredictionEvaluator."""

    @pytest.fixture
    def evaluator(self) -> PredictionEvaluator:
        """Create a basic evaluator."""
        return PredictionEvaluator(report_interval=0)

    def test_initialization_default(self) -> None:
        """Test default initialization."""
        evaluator = PredictionEvaluator()
        assert evaluator.accuracy_threshold == 20.0
        assert evaluator.report_interval == 100
        assert evaluator.enable_logging is True
        assert len(evaluator) == 0

    def test_initialization_custom(self) -> None:
        """Test custom initialization."""
        evaluator = PredictionEvaluator(
            accuracy_threshold=15.0,
            report_interval=50,
            enable_logging=False,
        )
        assert evaluator.accuracy_threshold == 15.0
        assert evaluator.report_interval == 50
        assert evaluator.enable_logging is False

    def test_invalid_accuracy_threshold(self) -> None:
        """Test invalid accuracy threshold."""
        with pytest.raises(ValueError, match="accuracy_threshold must be"):
            PredictionEvaluator(accuracy_threshold=0.0)

        with pytest.raises(ValueError, match="accuracy_threshold must be"):
            PredictionEvaluator(accuracy_threshold=150.0)

    def test_invalid_report_interval(self) -> None:
        """Test invalid report interval."""
        with pytest.raises(ValueError, match="report_interval must be >= 0"):
            PredictionEvaluator(report_interval=-1)

    def test_record_prediction_basic(self, evaluator: PredictionEvaluator) -> None:
        """Test basic prediction recording."""
        evaluator.record_prediction(
            request_id="req1",
            predicted_lifetime=10.0,
            actual_lifetime=12.0,
        )
        assert len(evaluator) == 1

        records = evaluator.get_records()
        assert len(records) == 1
        assert records[0].request_id == "req1"
        assert records[0].predicted_lifetime == 10.0
        assert records[0].actual_lifetime == 12.0

    def test_record_prediction_with_metadata(self, evaluator: PredictionEvaluator) -> None:
        """Test prediction recording with metadata."""
        evaluator.record_prediction(
            request_id="req1",
            predicted_lifetime=10.0,
            actual_lifetime=12.0,
            request_type="llm_chat",
            prompt_length=100,
            output_length=50,
            model_name="Qwen2-7B",
            metadata={"priority": 1},
        )
        assert len(evaluator) == 1

        records = evaluator.get_records()
        assert records[0].request_type == "llm_chat"
        assert records[0].prompt_length == 100
        assert records[0].output_length == 50
        assert records[0].model_name == "Qwen2-7B"
        assert records[0].metadata["priority"] == 1

    def test_record_prediction_invalid_lifetime(self, evaluator: PredictionEvaluator) -> None:
        """Test recording with invalid lifetimes."""
        with pytest.raises(ValueError, match="predicted_lifetime must be >= 0"):
            evaluator.record_prediction(
                request_id="req1",
                predicted_lifetime=-1.0,
                actual_lifetime=10.0,
            )

        with pytest.raises(ValueError, match="actual_lifetime must be >= 0"):
            evaluator.record_prediction(
                request_id="req1",
                predicted_lifetime=10.0,
                actual_lifetime=-1.0,
            )

    def test_compute_metrics_empty(self, evaluator: PredictionEvaluator) -> None:
        """Test computing metrics with no data."""
        metrics = evaluator.compute_metrics()
        assert metrics.num_predictions == 0
        assert metrics.mae == 0.0
        assert metrics.rmse == 0.0

    def test_compute_metrics_single_prediction(self, evaluator: PredictionEvaluator) -> None:
        """Test computing metrics with single prediction."""
        evaluator.record_prediction(
            request_id="req1",
            predicted_lifetime=10.0,
            actual_lifetime=12.0,
        )

        metrics = evaluator.compute_metrics()
        assert metrics.num_predictions == 1
        assert abs(metrics.mae - 2.0) < 0.001
        assert abs(metrics.rmse - 2.0) < 0.001
        assert abs(metrics.mape - 16.67) < 0.1
        assert abs(metrics.median_error - 2.0) < 0.001

    def test_compute_metrics_multiple_predictions(self, evaluator: PredictionEvaluator) -> None:
        """Test computing metrics with multiple predictions."""
        # Perfect prediction
        evaluator.record_prediction("req1", 10.0, 10.0)
        # 10% error
        evaluator.record_prediction("req2", 10.0, 11.0)
        # 20% error
        evaluator.record_prediction("req3", 10.0, 12.0)
        # 30% error
        evaluator.record_prediction("req4", 10.0, 13.0)

        metrics = evaluator.compute_metrics()
        assert metrics.num_predictions == 4

        # MAE = (0 + 1 + 2 + 3) / 4 = 1.5
        assert abs(metrics.mae - 1.5) < 0.001

        # RMSE = sqrt((0^2 + 1^2 + 2^2 + 3^2) / 4) = sqrt(14/4) = 1.87
        assert abs(metrics.rmse - 1.871) < 0.01

        # MAPE = (0 + 9.09 + 16.67 + 23.08) / 4 ≈ 12.21%
        assert 10.0 < metrics.mape < 15.0

        # Accuracy: 3 out of 4 have error < 20% = 75%
        assert abs(metrics.accuracy - 75.0) < 0.1

    def test_compute_metrics_accuracy_threshold(self) -> None:
        """Test accuracy calculation with different thresholds."""
        evaluator1 = PredictionEvaluator(accuracy_threshold=10.0, report_interval=0)
        evaluator2 = PredictionEvaluator(accuracy_threshold=30.0, report_interval=0)

        # Add same predictions to both
        for evaluator in [evaluator1, evaluator2]:
            evaluator.record_prediction("req1", 10.0, 10.0)  # 0% error
            evaluator.record_prediction("req2", 10.0, 11.0)  # 9% error
            evaluator.record_prediction("req3", 10.0, 12.0)  # 17% error
            evaluator.record_prediction("req4", 10.0, 13.0)  # 23% error

        # With 10% threshold: only 2 accurate (0%, 9%)
        metrics1 = evaluator1.compute_metrics()
        assert abs(metrics1.accuracy - 50.0) < 0.1

        # With 30% threshold: all 4 accurate
        metrics2 = evaluator2.compute_metrics()
        assert abs(metrics2.accuracy - 100.0) < 0.1

    def test_compute_grouped_metrics_by_type(self, evaluator: PredictionEvaluator) -> None:
        """Test grouped metrics by request type."""
        evaluator.record_prediction("req1", 10.0, 12.0, request_type="llm_chat")
        evaluator.record_prediction("req2", 10.0, 11.0, request_type="llm_chat")
        evaluator.record_prediction("req3", 10.0, 15.0, request_type="llm_generate")
        evaluator.record_prediction("req4", 10.0, 14.0, request_type="llm_generate")

        grouped = evaluator.compute_grouped_metrics("request_type")

        assert "llm_chat" in grouped
        assert "llm_generate" in grouped

        chat_metrics = grouped["llm_chat"]
        assert chat_metrics.num_predictions == 2
        # Chat: errors are 2.0 and 1.0, MAE = 1.5
        assert abs(chat_metrics.mae - 1.5) < 0.001

        gen_metrics = grouped["llm_generate"]
        assert gen_metrics.num_predictions == 2
        # Generate: errors are 5.0 and 4.0, MAE = 4.5
        assert abs(gen_metrics.mae - 4.5) < 0.001

    def test_compute_grouped_metrics_by_model(self, evaluator: PredictionEvaluator) -> None:
        """Test grouped metrics by model."""
        evaluator.record_prediction("req1", 10.0, 12.0, model_name="Qwen2-7B")
        evaluator.record_prediction("req2", 10.0, 11.0, model_name="Qwen2-7B")
        evaluator.record_prediction("req3", 10.0, 15.0, model_name="Llama-3-8B")

        grouped = evaluator.compute_grouped_metrics("model")

        assert "Qwen2-7B" in grouped
        assert "Llama-3-8B" in grouped

        qwen_metrics = grouped["Qwen2-7B"]
        assert qwen_metrics.num_predictions == 2

        llama_metrics = grouped["Llama-3-8B"]
        assert llama_metrics.num_predictions == 1

    def test_compute_grouped_metrics_by_length_bucket(self, evaluator: PredictionEvaluator) -> None:
        """Test grouped metrics by length bucket."""
        evaluator.record_prediction("req1", 10.0, 12.0, prompt_length=50, output_length=30)  # 80
        evaluator.record_prediction("req2", 10.0, 11.0, prompt_length=200, output_length=150)  # 350
        evaluator.record_prediction(
            "req3", 10.0, 15.0, prompt_length=600, output_length=500
        )  # 1100

        grouped = evaluator.compute_grouped_metrics("length_bucket")

        assert "0-100" in grouped
        assert "100-500" in grouped
        assert "1000-5000" in grouped

    def test_compute_metrics_with_grouping(self, evaluator: PredictionEvaluator) -> None:
        """Test computing metrics with specific grouping."""
        evaluator.record_prediction("req1", 10.0, 12.0, request_type="llm_chat")
        evaluator.record_prediction("req2", 10.0, 11.0, request_type="llm_chat")
        evaluator.record_prediction("req3", 10.0, 15.0, request_type="llm_generate")

        # Get metrics for only llm_chat
        chat_metrics = evaluator.compute_metrics(group_by="request_type", group_value="llm_chat")
        assert chat_metrics.num_predictions == 2

        # Get metrics for only llm_generate
        gen_metrics = evaluator.compute_metrics(group_by="request_type", group_value="llm_generate")
        assert gen_metrics.num_predictions == 1

    def test_compute_metrics_invalid_grouping(self, evaluator: PredictionEvaluator) -> None:
        """Test invalid grouping parameters."""
        evaluator.record_prediction("req1", 10.0, 12.0)

        # Invalid group_by
        with pytest.raises(ValueError, match="Invalid group_by"):
            evaluator.compute_metrics(group_by="invalid_dimension")

        # Missing group_value
        with pytest.raises(ValueError, match="group_value required"):
            evaluator.compute_metrics(group_by="request_type")

    def test_export_metrics(self, evaluator: PredictionEvaluator) -> None:
        """Test exporting metrics to JSON."""
        evaluator.record_prediction("req1", 10.0, 12.0, request_type="llm_chat")
        evaluator.record_prediction("req2", 10.0, 11.0, request_type="llm_generate")

        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "metrics.json"
            evaluator.export_metrics(output_path)

            assert output_path.exists()

            with open(output_path) as f:
                data = json.load(f)

            assert "timestamp" in data
            assert "config" in data
            assert "overall" in data
            assert "by_request_type" in data
            assert "recent_predictions" in data

            assert data["overall"]["num_predictions"] == 2
            assert "llm_chat" in data["by_request_type"]
            assert "llm_generate" in data["by_request_type"]

    def test_get_records_all(self, evaluator: PredictionEvaluator) -> None:
        """Test getting all records."""
        evaluator.record_prediction("req1", 10.0, 12.0)
        evaluator.record_prediction("req2", 10.0, 11.0)
        evaluator.record_prediction("req3", 10.0, 13.0)

        records = evaluator.get_records()
        assert len(records) == 3

    def test_get_records_limited(self, evaluator: PredictionEvaluator) -> None:
        """Test getting limited records."""
        evaluator.record_prediction("req1", 10.0, 12.0)
        evaluator.record_prediction("req2", 10.0, 11.0)
        evaluator.record_prediction("req3", 10.0, 13.0)
        evaluator.record_prediction("req4", 10.0, 14.0)

        # Get last 2
        records = evaluator.get_records(limit=2)
        assert len(records) == 2
        assert records[0].request_id == "req3"
        assert records[1].request_id == "req4"

    def test_clear(self, evaluator: PredictionEvaluator) -> None:
        """Test clearing all records."""
        evaluator.record_prediction("req1", 10.0, 12.0)
        evaluator.record_prediction("req2", 10.0, 11.0)

        assert len(evaluator) == 2

        evaluator.clear()

        assert len(evaluator) == 0
        metrics = evaluator.compute_metrics()
        assert metrics.num_predictions == 0

    def test_len(self, evaluator: PredictionEvaluator) -> None:
        """Test __len__ method."""
        assert len(evaluator) == 0

        evaluator.record_prediction("req1", 10.0, 12.0)
        assert len(evaluator) == 1

        evaluator.record_prediction("req2", 10.0, 11.0)
        assert len(evaluator) == 2

    def test_repr(self, evaluator: PredictionEvaluator) -> None:
        """Test __repr__ method."""
        repr_str = repr(evaluator)
        assert "PredictionEvaluator" in repr_str
        assert "num_predictions" in repr_str
        assert "accuracy_threshold" in repr_str

    def test_print_report_empty(self, evaluator: PredictionEvaluator, capsys) -> None:
        """Test printing report with no data."""
        evaluator.print_report()
        captured = capsys.readouterr()
        assert "No predictions" in captured.out

    def test_print_report_basic(self, evaluator: PredictionEvaluator, capsys) -> None:
        """Test printing basic report."""
        evaluator.record_prediction("req1", 10.0, 12.0)
        evaluator.record_prediction("req2", 10.0, 11.0)

        evaluator.print_report()
        captured = capsys.readouterr()

        assert "PREDICTION EVALUATOR REPORT" in captured.out
        assert "Overall Metrics" in captured.out
        assert "MAE:" in captured.out
        assert "RMSE:" in captured.out
        assert "MAPE:" in captured.out
        assert "Accuracy:" in captured.out

    def test_print_report_detailed(self, evaluator: PredictionEvaluator, capsys) -> None:
        """Test printing detailed report."""
        evaluator.record_prediction("req1", 10.0, 12.0, request_type="llm_chat")
        evaluator.record_prediction("req2", 10.0, 11.0, request_type="llm_generate")

        evaluator.print_report(detailed=True)
        captured = capsys.readouterr()

        assert "Metrics by Request Type:" in captured.out
        assert "llm_chat" in captured.out
        assert "llm_generate" in captured.out

    def test_automatic_reporting(self, capsys) -> None:
        """Test automatic report generation based on interval."""
        evaluator = PredictionEvaluator(report_interval=2)

        # First prediction - no report
        evaluator.record_prediction("req1", 10.0, 12.0)
        captured = capsys.readouterr()
        assert "PREDICTION EVALUATOR REPORT" not in captured.out

        # Second prediction - should trigger report
        evaluator.record_prediction("req2", 10.0, 11.0)
        captured = capsys.readouterr()
        assert "PREDICTION EVALUATOR REPORT" in captured.out

        # Third prediction - no report
        evaluator.record_prediction("req3", 10.0, 13.0)
        captured = capsys.readouterr()
        assert "PREDICTION EVALUATOR REPORT" not in captured.out

        # Fourth prediction - should trigger report again
        evaluator.record_prediction("req4", 10.0, 14.0)
        captured = capsys.readouterr()
        assert "PREDICTION EVALUATOR REPORT" in captured.out

    def test_integration_with_predictor_workflow(self, evaluator: PredictionEvaluator) -> None:
        """Test integration with typical predictor workflow."""
        # Simulate predictor workflow
        requests = [
            ("req1", 100, 50, 10.0, 12.0),  # predicted, actual
            ("req2", 200, 100, 20.0, 22.0),
            ("req3", 150, 75, 15.0, 14.0),
            ("req4", 300, 150, 30.0, 35.0),
            ("req5", 250, 125, 25.0, 24.0),
        ]

        for req_id, prompt_len, output_len, predicted, actual in requests:
            evaluator.record_prediction(
                request_id=req_id,
                predicted_lifetime=predicted,
                actual_lifetime=actual,
                prompt_length=prompt_len,
                output_length=output_len,
                model_name="Qwen2-7B",
            )

        # Check overall metrics
        metrics = evaluator.compute_metrics()
        assert metrics.num_predictions == 5
        assert metrics.mae > 0
        assert metrics.rmse >= metrics.mae
        assert 0 <= metrics.accuracy <= 100

        # Check grouping by length
        length_metrics = evaluator.compute_grouped_metrics("length_bucket")
        assert len(length_metrics) > 0


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