#!/usr/bin/env python3
"""
Tests for error_case_compare.py module

Run with:
    pytest chengyue_script/dvlm/evaluation/tests/test_error_case_compare.py -v
"""

import sys
import json
import tempfile
from pathlib import Path

import pytest
import pandas as pd

# Add parent directory to path
test_dir = Path(__file__).parent
eval_dir = test_dir.parent
if str(eval_dir) not in sys.path:
    sys.path.insert(0, str(eval_dir))

from error_case_compare import (
    find_experiment_paths,
    get_checkpoint_dir,
    select_checkpoint,
    find_results_xlsx,
    analyze_benchmark_errors,
    compare_experiments,
    generate_markdown_report,
    generate_json_report,
    ExperimentInfo,
    BenchmarkResult,
    ComparisonResult,
)


class TestExperimentDiscovery:
    """Test experiment path discovery."""
    
    def test_experiment_info_creation(self):
        """Test ExperimentInfo dataclass."""
        info = ExperimentInfo(name="test_exp")
        assert info.name == "test_exp"
        assert info.checkpoints == []
        assert info.vlmeval_dir is None
    
    def test_find_nonexistent_experiment(self, tmp_path, monkeypatch):
        """Test finding an experiment that doesn't exist."""
        # Monkeypatch PROJECT_ROOT to use tmp_path
        import error_case_compare
        monkeypatch.setattr(error_case_compare, "PROJECT_ROOT", tmp_path)
        
        info = find_experiment_paths("nonexistent_experiment")
        assert info.name == "nonexistent_experiment"
        assert info.vlmeval_dir is None
        assert info.checkpoints == []


class TestCheckpointSelection:
    """Test checkpoint selection logic."""
    
    def test_select_latest(self):
        """Test selecting latest checkpoint."""
        info = ExperimentInfo(
            name="test",
            checkpoints=["checkpoint-100", "checkpoint-200", "checkpoint-300"]
        )
        selected = select_checkpoint(info, "latest")
        assert selected == "checkpoint-300"
    
    def test_select_specific_step(self):
        """Test selecting specific step."""
        info = ExperimentInfo(
            name="test",
            checkpoints=["checkpoint-100", "checkpoint-200", "checkpoint-300"]
        )
        selected = select_checkpoint(info, "step:200")
        assert selected == "checkpoint-200"
        
        # Non-existent step
        selected = select_checkpoint(info, "step:999")
        assert selected is None
    
    def test_select_empty_checkpoints(self):
        """Test selecting from empty checkpoints."""
        info = ExperimentInfo(name="test", checkpoints=[])
        selected = select_checkpoint(info, "latest")
        assert selected is None


class TestBenchmarkAnalysis:
    """Test benchmark analysis with sample data."""
    
    def test_analyze_benchmark_errors(self, sample_results_df):
        """Test error analysis on sample data."""
        result = analyze_benchmark_errors(
            df=sample_results_df,
            benchmark="TextVQA",
            source_file="/path/to/results.xlsx",
            checkpoint="checkpoint-1000",
            max_details=100,
        )
        
        assert isinstance(result, BenchmarkResult)
        assert result.benchmark == "TextVQA"
        assert result.checkpoint == "checkpoint-1000"
        assert result.total_samples == 10
        assert result.correct_count == 0  # All have eval_score=0
        assert result.accuracy == 0.0
        assert len(result.pattern_counts) > 0
        assert len(result.pattern_details) > 0
        
        # Check that specific patterns were detected
        patterns = result.pattern_counts
        assert "empty_or_null" in patterns or "yes_no_flip" in patterns
    
    def test_analyze_with_correct_samples(self):
        """Test analysis with some correct samples."""
        df = pd.DataFrame({
            "index": [1, 2, 3],
            "question": ["Q1", "Q2", "Q3"],
            "answer": ["A1", "A2", "A3"],
            "prediction": ["A1", "wrong", "A3"],
            "eval_score": [1.0, 0.0, 1.0],
        })
        
        result = analyze_benchmark_errors(
            df=df,
            benchmark="Test",
            source_file="test.xlsx",
            checkpoint="checkpoint-100",
        )
        
        assert result.total_samples == 3
        assert result.correct_count == 2
        assert result.accuracy == pytest.approx(66.67, rel=0.1)
        # Only 1 error case
        assert len(result.pattern_details) == 1


class TestReportGeneration:
    """Test report generation."""
    
    def test_markdown_report(self):
        """Test markdown report generation."""
        result = ComparisonResult(
            exp_a="exp1",
            exp_b="exp2",
            checkpoint_a="checkpoint-100",
            checkpoint_b="checkpoint-200",
            benchmarks={
                "TextVQA": {
                    "exp_a": {
                        "benchmark": "TextVQA",
                        "checkpoint": "checkpoint-100",
                        "total_samples": 100,
                        "correct_count": 70,
                        "accuracy": 70.0,
                        "pattern_counts": {"empty_or_null": 5, "yes_no_flip": 10},
                        "pattern_details": [],
                        "source_file": "/path/a.xlsx",
                    },
                    "exp_b": {
                        "benchmark": "TextVQA",
                        "checkpoint": "checkpoint-200",
                        "total_samples": 100,
                        "correct_count": 60,
                        "accuracy": 60.0,
                        "pattern_counts": {"empty_or_null": 10, "yes_no_flip": 15},
                        "pattern_details": [],
                        "source_file": "/path/b.xlsx",
                    },
                    "delta": {
                        "accuracy_diff": 10.0,
                        "correct_diff": 10,
                    }
                }
            },
            sources=["/path/a.xlsx", "/path/b.xlsx"],
            generated_at="2025-01-21T12:00:00",
        )
        
        md = generate_markdown_report(result)
        
        # Check structure
        assert "# Error Case Analysis" in md
        assert "exp1" in md
        assert "exp2" in md
        assert "Scoreboard" in md
        assert "TextVQA" in md
        assert "Sources" in md
        assert "Pattern Legend" in md
    
    def test_json_report(self):
        """Test JSON report generation."""
        result = ComparisonResult(
            exp_a="exp1",
            exp_b="exp2",
            checkpoint_a="checkpoint-100",
            checkpoint_b="checkpoint-200",
            benchmarks={},
            sources=[],
            generated_at="2025-01-21T12:00:00",
        )
        
        report = generate_json_report(result)
        
        assert report["exp_a"] == "exp1"
        assert report["exp_b"] == "exp2"
        assert "pattern_descriptions" in report
        
        # Should be JSON serializable
        json_str = json.dumps(report)
        assert len(json_str) > 0


class TestXlsxDiscovery:
    """Test XLSX file discovery."""
    
    def test_find_results_xlsx(self, temp_xlsx_file, temp_experiment_dir):
        """Test finding XLSX in benchmark directory."""
        benchmark_dir = temp_experiment_dir["benchmark_dir"]
        found = find_results_xlsx(benchmark_dir)
        assert found is not None
        assert found.exists()
        assert found.suffix == ".xlsx"
    
    def test_find_results_xlsx_not_found(self, tmp_path):
        """Test when no XLSX exists."""
        empty_dir = tmp_path / "empty"
        empty_dir.mkdir()
        found = find_results_xlsx(empty_dir)
        assert found is None


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

