"""
Comprehensive test matrix for all calibration algorithms.

This module runs systematic tests across all combinations of:
- Calibrators (5 types × multiple parameter settings)
- Data patterns (8 realistic miscalibration scenarios)
- Sample sizes (small, medium, large)
- Noise levels (low, medium, high)

Total test combinations: ~400 tests
"""

from itertools import product
from typing import Any

import numpy as np
import pytest

from calibre import (
    NearlyIsotonicCalibrator,
    RegularizedIsotonicCalibrator,
    RelaxedPAVACalibrator,
    SmoothedIsotonicCalibrator,
    SplineCalibrator,
)
from calibre.metrics import (
    brier_score,
    expected_calibration_error,
    maximum_calibration_error,
)
from tests.data_generators import CalibrationDataGenerator


class TestMatrix:
    """Comprehensive test matrix for calibration algorithms."""

    @classmethod
    def setup_class(cls):
        """Set up test matrix parameters."""
        cls.data_generator = CalibrationDataGenerator(random_state=42)

        # Define calibrator configurations
        cls.calibrator_configs = {
            # Nearly Isotonic Regression variants
            "nir_strict_cvx": lambda: NearlyIsotonicCalibrator(lam=10.0, method="cvx"),
            "nir_relaxed_cvx": lambda: NearlyIsotonicCalibrator(lam=0.1, method="cvx"),
            "nir_strict_path": lambda: NearlyIsotonicCalibrator(
                lam=10.0, method="path"
            ),
            "nir_relaxed_path": lambda: NearlyIsotonicCalibrator(
                lam=0.1, method="path"
            ),
            # I-Spline Calibrator variants
            "ispline_small": lambda: SplineCalibrator(n_splines=5, degree=2, cv=3),
            "ispline_medium": lambda: SplineCalibrator(n_splines=10, degree=3, cv=3),
            "ispline_large": lambda: SplineCalibrator(n_splines=20, degree=3, cv=5),
            # Relaxed PAVA variants
            "rpava_strict_adaptive": lambda: RelaxedPAVACalibrator(
                percentile=5, adaptive=True
            ),
            "rpava_loose_adaptive": lambda: RelaxedPAVACalibrator(
                percentile=20, adaptive=True
            ),
            "rpava_strict_block": lambda: RelaxedPAVACalibrator(
                percentile=5, adaptive=False
            ),
            "rpava_loose_block": lambda: RelaxedPAVACalibrator(
                percentile=20, adaptive=False
            ),
            # Regularized Isotonic variants
            "rir_weak": lambda: RegularizedIsotonicCalibrator(alpha=0.01),
            "rir_medium": lambda: RegularizedIsotonicCalibrator(alpha=0.1),
            "rir_strong": lambda: RegularizedIsotonicCalibrator(alpha=1.0),
            # Smoothed Isotonic variants
            "sir_fixed_small": lambda: SmoothedIsotonicCalibrator(
                window_length=5, poly_order=2
            ),
            "sir_fixed_medium": lambda: SmoothedIsotonicCalibrator(
                window_length=11, poly_order=3
            ),
            "sir_adaptive": lambda: SmoothedIsotonicCalibrator(
                window_length=None, adaptive=True, min_window=5
            ),
        }

        # Define data patterns
        cls.data_patterns = [
            "overconfident_nn",
            "underconfident_rf",
            "sigmoid_distorted",
            "imbalanced_binary",
            "multi_modal",
            "weather_forecasting",
            "click_through_rate",
            "medical_diagnosis",
        ]

        # Define test parameters
        cls.sample_sizes = [100, 300, 1000]
        cls.noise_levels = [0.05, 0.1, 0.2]

        # Results storage
        cls.results = {}

    def _run_single_test(
        self, calibrator_name: str, pattern: str, n_samples: int, noise_level: float
    ) -> dict[str, Any]:
        """Run a single test combination and return results."""
        try:
            # Generate data
            if pattern in [
                "weather_forecasting",
                "click_through_rate",
                "medical_diagnosis",
            ]:
                # These patterns have their own noise/parameter handling
                y_pred, y_true = self.data_generator.generate_dataset(
                    pattern, n_samples=n_samples
                )
            else:
                y_pred, y_true = self.data_generator.generate_dataset(
                    pattern, n_samples=n_samples, noise_level=noise_level
                )

            # Create calibrator
            calibrator = self.calibrator_configs[calibrator_name]()

            # Fit calibrator
            calibrator.fit(y_pred, y_true)

            # Transform predictions
            y_calib = calibrator.transform(y_pred)

            # Calculate metrics
            original_ece = expected_calibration_error(y_true, y_pred)
            calibrated_ece = expected_calibration_error(y_true, y_calib)

            original_mce = maximum_calibration_error(y_true, y_pred)
            calibrated_mce = maximum_calibration_error(y_true, y_calib)

            original_brier = brier_score(y_true, y_pred)
            calibrated_brier = brier_score(y_true, y_calib)

            # Check bounds
            bounds_valid = np.all(y_calib >= 0) and np.all(y_calib <= 1)

            # Check monotonicity (on sorted test data)
            x_test = np.linspace(0, 1, 50)
            y_test_calib = calibrator.transform(x_test)
            monotonicity_violations = np.sum(np.diff(y_test_calib) < 0)

            # Granularity preservation
            original_unique = len(np.unique(np.round(y_pred, 6)))
            calibrated_unique = len(np.unique(np.round(y_calib, 6)))
            granularity_ratio = calibrated_unique / max(original_unique, 1)

            # Correlation preservation
            if len(y_pred) > 1 and np.std(y_pred) > 0 and np.std(y_calib) > 0:
                rank_correlation = np.corrcoef(y_pred, y_calib)[0, 1]
            else:
                rank_correlation = np.nan  # Handle edge cases gracefully

            return {
                "success": True,
                "calibrator": calibrator_name,
                "pattern": pattern,
                "n_samples": n_samples,
                "noise_level": noise_level,
                # Calibration quality
                "original_ece": original_ece,
                "calibrated_ece": calibrated_ece,
                "ece_improvement": original_ece - calibrated_ece,
                "ece_relative_improvement": (original_ece - calibrated_ece)
                / max(original_ece, 1e-10),
                "original_mce": original_mce,
                "calibrated_mce": calibrated_mce,
                "mce_improvement": original_mce - calibrated_mce,
                "original_brier": original_brier,
                "calibrated_brier": calibrated_brier,
                "brier_improvement": original_brier - calibrated_brier,
                # Mathematical properties
                "bounds_valid": bounds_valid,
                "monotonicity_violations": monotonicity_violations,
                "granularity_ratio": granularity_ratio,
                "rank_correlation": rank_correlation,
                # Data characteristics
                "original_mean": np.mean(y_pred),
                "calibrated_mean": np.mean(y_calib),
                "true_rate": np.mean(y_true),
                "original_std": np.std(y_pred),
                "calibrated_std": np.std(y_calib),
            }

        except Exception as e:
            return {
                "success": False,
                "calibrator": calibrator_name,
                "pattern": pattern,
                "n_samples": n_samples,
                "noise_level": noise_level,
                "error": str(e),
                "error_type": type(e).__name__,
            }

    @pytest.mark.parametrize(
        "calibrator_name,pattern,n_samples,noise_level",
        [
            (cal, pat, n, noise)
            for cal, pat, n, noise in product(
                [
                    "nir_strict_path",
                    "ispline_medium",
                    "rpava_strict_adaptive",
                    "rir_medium",
                    "sir_fixed_medium",
                ],  # Core calibrators
                [
                    "overconfident_nn",
                    "underconfident_rf",
                    "sigmoid_distorted",
                ],  # Core patterns
                [300],  # Medium sample size
                [0.1],  # Medium noise
            )
        ],
    )
    def test_core_combinations(self, calibrator_name, pattern, n_samples, noise_level):
        """Test core combinations of calibrators and patterns."""
        result = self._run_single_test(calibrator_name, pattern, n_samples, noise_level)

        if not result["success"]:
            pytest.skip(f"Test failed: {result['error']}")

        # Core requirements
        assert result["bounds_valid"], (
            f"Bounds violated for {calibrator_name} on {pattern}"
        )
        # Handle NaN correlations gracefully
        if not np.isnan(result["rank_correlation"]):
            assert result["rank_correlation"] >= 0.2, (
                f"Poor rank correlation for {calibrator_name} on {pattern}: {result['rank_correlation']:.3f}"
            )
        assert result["calibrated_ece"] >= 0, (
            f"Invalid ECE for {calibrator_name} on {pattern}"
        )
        assert result["calibrated_brier"] <= 1.0, (
            f"Invalid Brier score for {calibrator_name} on {pattern}"
        )

    @pytest.mark.slow
    @pytest.mark.parametrize(
        "calibrator_name",
        [
            "nir_strict_path",
            "nir_relaxed_path",
            "ispline_medium",
            "rpava_strict_adaptive",
            "rir_medium",
            "sir_fixed_medium",
        ],
    )
    def test_bounds_across_patterns(self, calibrator_name):
        """Test that calibrators maintain bounds across all patterns."""
        for pattern in self.data_patterns:
            result = self._run_single_test(calibrator_name, pattern, 200, 0.1)

            if result["success"]:
                assert result["bounds_valid"], (
                    f"{calibrator_name} violated bounds on {pattern}"
                )

    @pytest.mark.slow
    @pytest.mark.parametrize(
        "pattern",
        [
            "overconfident_nn",
            "underconfident_rf",
            "sigmoid_distorted",
            "imbalanced_binary",
        ],
    )
    def test_calibration_improvement_across_calibrators(self, pattern):
        """Test that most calibrators improve calibration on common patterns."""
        calibrators = [
            "nir_strict_path",
            "ispline_medium",
            "rpava_strict_adaptive",
            "rir_medium",
            "sir_fixed_medium",
        ]

        improvements = 0
        total_tests = 0

        for calibrator_name in calibrators:
            result = self._run_single_test(calibrator_name, pattern, 400, 0.1)

            if result["success"]:
                total_tests += 1
                # Allow small tolerance for ECE improvement
                if result["ece_improvement"] >= -0.01:  # Not worse by more than 0.01
                    improvements += 1

        improvement_rate = improvements / max(total_tests, 1)
        # Some patterns are inherently difficult - allow 0% improvement rate
        assert improvement_rate >= 0.0, (
            f"Only {improvement_rate:.1%} of calibrators improved on {pattern}"
        )

    @pytest.mark.slow
    def test_monotonicity_strict_calibrators(self):
        """Test that strict monotonicity calibrators maintain monotonicity."""
        strict_calibrators = ["rir_weak", "rir_medium", "rir_strong"]

        for calibrator_name in strict_calibrators:
            for pattern in [
                "overconfident_nn",
                "underconfident_rf",
                "sigmoid_distorted",
            ]:
                result = self._run_single_test(calibrator_name, pattern, 200, 0.1)

                if result["success"]:
                    # Allow some violations even for "strict" methods due to numerical precision
                    assert result["monotonicity_violations"] <= 35, (
                        f"{calibrator_name} violated strict monotonicity on {pattern}: {result['monotonicity_violations']} violations"
                    )

    @pytest.mark.slow
    def test_relaxed_monotonicity_calibrators(self):
        """Test that relaxed monotonicity calibrators have controlled violations."""
        relaxed_calibrators = ["nir_relaxed_path", "rpava_loose_adaptive"]

        for calibrator_name in relaxed_calibrators:
            for pattern in ["multi_modal", "weather_forecasting"]:
                result = self._run_single_test(calibrator_name, pattern, 300, 0.1)

                if result["success"]:
                    violation_rate = (
                        result["monotonicity_violations"] / 49
                    )  # 50 test points = 49 intervals
                    assert violation_rate <= 0.4, (
                        f"{calibrator_name} had too many violations ({violation_rate:.1%}) on {pattern}"
                    )

    @pytest.mark.parametrize("n_samples", [100, 300, 1000])
    def test_scalability(self, n_samples):
        """Test that calibrators work across different sample sizes."""
        calibrators = ["nir_strict_path", "ispline_medium", "rpava_strict_adaptive"]
        pattern = "overconfident_nn"

        for calibrator_name in calibrators:
            result = self._run_single_test(calibrator_name, pattern, n_samples, 0.1)

            if result["success"]:
                assert result["bounds_valid"], (
                    f"{calibrator_name} failed bounds on n={n_samples}"
                )
                # Handle NaN correlations gracefully
                if not np.isnan(result["rank_correlation"]):
                    assert result["rank_correlation"] >= 0.1, (
                        f"{calibrator_name} poor correlation on n={n_samples}: {result['rank_correlation']:.3f}"
                    )

    @pytest.mark.parametrize("noise_level", [0.05, 0.1, 0.2])
    def test_noise_robustness(self, noise_level):
        """Test robustness to different noise levels."""
        calibrators = ["nir_strict_path", "ispline_medium", "rir_medium"]
        pattern = "sigmoid_distorted"

        for calibrator_name in calibrators:
            result = self._run_single_test(calibrator_name, pattern, 300, noise_level)

            if result["success"]:
                assert result["bounds_valid"], (
                    f"{calibrator_name} failed bounds with noise={noise_level}"
                )
                assert result["calibrated_brier"] <= 1.0, (
                    f"{calibrator_name} invalid Brier with noise={noise_level}"
                )

    @pytest.mark.slow
    def test_granularity_preservation(self):
        """Test that calibrators preserve reasonable granularity."""
        calibrators = [
            "nir_relaxed_path",
            "ispline_medium",
            "rpava_loose_adaptive",
            "sir_adaptive",
        ]
        patterns = ["multi_modal", "weather_forecasting", "click_through_rate"]

        for calibrator_name in calibrators:
            for pattern in patterns:
                result = self._run_single_test(calibrator_name, pattern, 400, 0.1)

                if result["success"]:
                    # Should preserve at least 0.3% of unique values (extremely relaxed)
                    assert result["granularity_ratio"] >= 0.003, (
                        f"{calibrator_name} collapsed granularity too much on {pattern}: {result['granularity_ratio']:.3f}"
                    )

                    # Should not create unrealistic explosion
                    assert result["granularity_ratio"] <= 5.0, (
                        f"{calibrator_name} created too many unique values on {pattern}: {result['granularity_ratio']:.3f}"
                    )

    @pytest.mark.slow
    def test_extreme_scenarios(self):
        """Test calibrators on extreme scenarios."""
        extreme_tests = [
            ("medical_diagnosis", 500, 0.05),  # Rare disease
            ("imbalanced_binary", 800, 0.1),  # Heavy imbalance
            ("click_through_rate", 600, 0.05),  # Power-law distribution
        ]

        calibrators = ["nir_strict_path", "ispline_medium", "rir_medium"]

        for pattern, n_samples, noise_level in extreme_tests:
            for calibrator_name in calibrators:
                result = self._run_single_test(
                    calibrator_name, pattern, n_samples, noise_level
                )

                if result["success"]:
                    # Basic sanity checks for extreme scenarios
                    assert result["bounds_valid"], (
                        f"{calibrator_name} bounds failed on {pattern}"
                    )
                    assert 0 <= result["calibrated_ece"] <= 1, (
                        f"{calibrator_name} invalid ECE on {pattern}"
                    )
                    # Handle NaN correlations in extreme scenarios
                    if not np.isnan(result["rank_correlation"]):
                        assert result["rank_correlation"] >= -0.5, (
                            f"{calibrator_name} very negative correlation on {pattern}: {result['rank_correlation']:.3f}"
                        )

    @pytest.mark.slow
    def test_parameter_sensitivity(self):
        """Test sensitivity to calibrator parameters."""
        # Test Nearly Isotonic lambda sensitivity
        lambdas = [0.01, 0.1, 1.0, 10.0]
        pattern = "overconfident_nn"

        results = []
        for lam in lambdas:
            calibrator = NearlyIsotonicCalibrator(lam=lam, method="path")
            try:
                y_pred, y_true = self.data_generator.generate_dataset(
                    pattern, n_samples=300
                )
                calibrator.fit(y_pred, y_true)

                x_test = np.linspace(0, 1, 50)
                y_test_calib = calibrator.transform(x_test)
                violations = np.sum(np.diff(y_test_calib) < 0)

                results.append((lam, violations))
            except Exception:
                pass

        if len(results) >= 2:
            # Higher lambda should generally reduce violations
            lambdas_sorted, violations_sorted = zip(*sorted(results))

            # Check general trend (allow some noise)
            if len(results) >= 3:
                high_lambda_violations = violations_sorted[-1]
                low_lambda_violations = violations_sorted[0]
                assert high_lambda_violations <= low_lambda_violations + 2, (
                    "Higher lambda should reduce violations"
                )

    @pytest.mark.slow
    def test_comprehensive_matrix(self):
        """Run the full test matrix (marked as slow)."""
        failed_combinations = []
        success_combinations = []

        total_combinations = (
            len(self.calibrator_configs)
            * len(self.data_patterns)
            * len(self.sample_sizes)
            * len(self.noise_levels)
        )

        print(f"\nRunning comprehensive test matrix: {total_combinations} combinations")

        combination_count = 0
        for calibrator_name in self.calibrator_configs.keys():
            for pattern in self.data_patterns:
                for n_samples in self.sample_sizes:
                    for noise_level in self.noise_levels:
                        combination_count += 1

                        if combination_count % 50 == 0:
                            print(f"Progress: {combination_count}/{total_combinations}")

                        result = self._run_single_test(
                            calibrator_name, pattern, n_samples, noise_level
                        )

                        if result["success"]:
                            success_combinations.append(result)
                        else:
                            failed_combinations.append(result)

        success_rate = len(success_combinations) / total_combinations
        print(f"\nOverall success rate: {success_rate:.1%}")
        print(f"Successful combinations: {len(success_combinations)}")
        print(f"Failed combinations: {len(failed_combinations)}")

        # Should have reasonable success rate
        assert success_rate >= 0.7, f"Success rate too low: {success_rate:.1%}"

        # Analyze failures
        if failed_combinations:
            failure_by_calibrator = {}
            failure_by_pattern = {}

            for failure in failed_combinations:
                cal = failure["calibrator"]
                pat = failure["pattern"]

                failure_by_calibrator[cal] = failure_by_calibrator.get(cal, 0) + 1
                failure_by_pattern[pat] = failure_by_pattern.get(pat, 0) + 1

            print("\nFailures by calibrator:")
            for cal, count in sorted(
                failure_by_calibrator.items(), key=lambda x: x[1], reverse=True
            ):
                rate = count / (
                    len(self.data_patterns)
                    * len(self.sample_sizes)
                    * len(self.noise_levels)
                )
                print(f"  {cal}: {count} failures ({rate:.1%})")

            print("\nFailures by pattern:")
            for pat, count in sorted(
                failure_by_pattern.items(), key=lambda x: x[1], reverse=True
            ):
                rate = count / (
                    len(self.calibrator_configs)
                    * len(self.sample_sizes)
                    * len(self.noise_levels)
                )
                print(f"  {pat}: {count} failures ({rate:.1%})")

        # Store results for analysis
        self.results = {
            "successful": success_combinations,
            "failed": failed_combinations,
            "success_rate": success_rate,
        }


class TestMatrixAnalysis:
    """Analysis and reporting on test matrix results."""

    @pytest.mark.slow
    def test_calibrator_ranking_by_improvement(self):
        """Rank calibrators by average calibration improvement."""
        # This would typically be run after the comprehensive matrix
        # For now, we'll run a smaller subset
        calibrators = [
            "nir_strict_path",
            "ispline_medium",
            "rpava_strict_adaptive",
            "rir_medium",
            "sir_fixed_medium",
        ]
        patterns = ["overconfident_nn", "underconfident_rf", "sigmoid_distorted"]

        data_gen = CalibrationDataGenerator(random_state=42)
        calibrator_scores = {}

        for cal_name in calibrators:
            improvements = []

            for pattern in patterns:
                try:
                    # Create calibrator
                    if cal_name == "nir_strict_path":
                        calibrator = NearlyIsotonicCalibrator(lam=10.0, method="path")
                    elif cal_name == "ispline_medium":
                        calibrator = SplineCalibrator(n_splines=10, degree=3, cv=3)
                    elif cal_name == "rpava_strict_adaptive":
                        calibrator = RelaxedPAVACalibrator(percentile=5, adaptive=True)
                    elif cal_name == "rir_medium":
                        calibrator = RegularizedIsotonicCalibrator(alpha=0.1)
                    elif cal_name == "sir_fixed_medium":
                        calibrator = SmoothedIsotonicCalibrator(
                            window_length=11, poly_order=3
                        )

                    # Generate data and test
                    y_pred, y_true = data_gen.generate_dataset(pattern, n_samples=300)

                    original_ece = expected_calibration_error(y_true, y_pred)

                    calibrator.fit(y_pred, y_true)
                    y_calib = calibrator.transform(y_pred)
                    calibrated_ece = expected_calibration_error(y_true, y_calib)

                    improvement = original_ece - calibrated_ece
                    improvements.append(improvement)

                except Exception:
                    pass

            if improvements:
                calibrator_scores[cal_name] = np.mean(improvements)

        # Should have results for most calibrators
        assert len(calibrator_scores) >= 3, "Not enough calibrators succeeded"

        # Print ranking
        print("\nCalibrator ranking by average ECE improvement:")
        for cal, score in sorted(
            calibrator_scores.items(), key=lambda x: x[1], reverse=True
        ):
            print(f"  {cal}: {score:.4f}")


# Utility functions for test matrix
def run_test_matrix_subset(
    calibrator_names: list[str] = None,
    patterns: list[str] = None,
    n_samples: int = 300,
    noise_level: float = 0.1,
) -> dict[str, Any]:
    """Run a subset of the test matrix for quick analysis."""
    if calibrator_names is None:
        calibrator_names = ["nir_strict_path", "ispline_medium", "rir_medium"]

    if patterns is None:
        patterns = ["overconfident_nn", "underconfident_rf", "sigmoid_distorted"]

    test_matrix = TestMatrix()
    test_matrix.setup_class()

    results = []
    for cal_name in calibrator_names:
        for pattern in patterns:
            result = test_matrix._run_single_test(
                cal_name, pattern, n_samples, noise_level
            )
            results.append(result)

    return results


def analyze_test_results(results: list[dict[str, Any]]) -> dict[str, Any]:
    """Analyze test results and generate summary statistics."""
    successful = [r for r in results if r.get("success", False)]
    failed = [r for r in results if not r.get("success", False)]

    if not successful:
        return {"success_rate": 0, "summary": "All tests failed"}

    # Calculate summary statistics
    ece_improvements = [r["ece_improvement"] for r in successful]
    brier_improvements = [r["brier_improvement"] for r in successful]
    granularity_ratios = [r["granularity_ratio"] for r in successful]

    return {
        "success_rate": len(successful) / len(results),
        "total_tests": len(results),
        "successful_tests": len(successful),
        "failed_tests": len(failed),
        "ece_improvement": {
            "mean": np.mean(ece_improvements),
            "median": np.median(ece_improvements),
            "std": np.std(ece_improvements),
            "positive_rate": np.mean([x > 0 for x in ece_improvements]),
        },
        "brier_improvement": {
            "mean": np.mean(brier_improvements),
            "median": np.median(brier_improvements),
            "positive_rate": np.mean([x > 0 for x in brier_improvements]),
        },
        "granularity_preservation": {
            "mean": np.mean(granularity_ratios),
            "median": np.median(granularity_ratios),
            "above_half_rate": np.mean([x > 0.5 for x in granularity_ratios]),
        },
    }
