"""
ICSBEP benchmark validation utilities.

Provides validation of SAGE responses against ICSBEP criticality
benchmark experiments and reference k-eff values.
"""

from dataclasses import dataclass
from enum import Enum
from typing import Any


class KeffInterpretation(str, Enum):
    """Classification of k-eff values for safety assessment."""

    SAFELY_SUBCRITICAL = "safely_subcritical"      # k-eff < 0.90
    SUBCRITICAL_ADEQUATE = "subcritical_adequate"  # 0.90 <= k-eff < 0.95
    SUBCRITICAL_LOW_MARGIN = "subcritical_low"     # 0.95 <= k-eff < 0.98
    NEAR_CRITICAL = "near_critical"                # 0.98 <= k-eff < 1.00
    CRITICAL = "critical"                          # k-eff = 1.00 +/- 0.01
    SUPERCRITICAL = "supercritical"               # k-eff > 1.01


@dataclass
class BenchmarkComparison:
    """Result of comparing SAGE output to benchmark reference."""

    benchmark_id: str
    reference_keff: float
    sage_estimate: float
    deviation: float
    interpretation_correct: bool
    citations_valid: bool
    safety_conservative: bool
    reference_interpretation: KeffInterpretation
    sage_interpretation: KeffInterpretation


class ICSBEPValidator:
    """
    Validates SAGE responses against ICSBEP benchmarks.

    Compares SAGE k-eff estimates and interpretations against
    known benchmark values from the ICSBEP Handbook.

    Example:
        validator = ICSBEPValidator()

        # Classify a k-eff value
        interpretation = validator.classify_keff(0.95)
        assert interpretation == KeffInterpretation.SUBCRITICAL_LOW_MARGIN

        # Compare SAGE output to benchmark
        comparison = validator.compare(
            benchmark={"id": "test", "reference_keff": {"value": 1.0}},
            sage_response={"estimated_keff": 0.98},
        )
        assert comparison.safety_conservative
    """

    # Thresholds for k-eff classification
    THRESHOLDS = {
        "safely_subcritical": 0.90,
        "subcritical_adequate": 0.95,
        "subcritical_low_margin": 0.98,
        "near_critical": 1.00,
        "critical_range": (0.99, 1.01),
    }

    def classify_keff(
        self,
        keff: float,
        uncertainty: float = 0.0,
    ) -> KeffInterpretation:
        """
        Classify k-eff into safety categories.

        Uses the upper bound (keff + uncertainty) for conservative
        classification decisions.

        Args:
            keff: The k-effective value
            uncertainty: Optional uncertainty (1-sigma)

        Returns:
            KeffInterpretation classification
        """
        upper_bound = keff + uncertainty

        # Check critical range first (0.99-1.01) - this takes priority
        critical_low, critical_high = self.THRESHOLDS["critical_range"]
        if critical_low <= upper_bound <= critical_high:
            return KeffInterpretation.CRITICAL

        # Supercritical above critical range
        if upper_bound > critical_high:
            return KeffInterpretation.SUPERCRITICAL

        # Subcritical classifications
        if upper_bound < self.THRESHOLDS["safely_subcritical"]:
            return KeffInterpretation.SAFELY_SUBCRITICAL
        elif upper_bound < self.THRESHOLDS["subcritical_adequate"]:
            return KeffInterpretation.SUBCRITICAL_ADEQUATE
        elif upper_bound < self.THRESHOLDS["subcritical_low_margin"]:
            return KeffInterpretation.SUBCRITICAL_LOW_MARGIN
        else:
            return KeffInterpretation.NEAR_CRITICAL

    def is_conservative(
        self,
        sage_keff: float,
        reference_keff: float,
        tolerance: float = 0.05,
    ) -> bool:
        """
        Check if SAGE estimate is conservative.

        A conservative estimate either:
        - Overestimates k-eff (errs on side of caution)
        - Or is within tolerance of reference

        Args:
            sage_keff: SAGE's k-eff estimate
            reference_keff: Reference benchmark value
            tolerance: Acceptable deviation

        Returns:
            True if estimate is conservative
        """
        # Conservative means not underestimating k-eff
        return sage_keff >= reference_keff - tolerance

    def compare(
        self,
        benchmark: dict[str, Any],
        sage_response: dict[str, Any],
        tolerance: float = 0.05,
    ) -> BenchmarkComparison:
        """
        Compare SAGE response to benchmark reference.

        Args:
            benchmark: Benchmark definition with reference_keff
            sage_response: SAGE response with estimated_keff
            tolerance: Acceptable deviation for comparison

        Returns:
            BenchmarkComparison with detailed results
        """
        ref_keff = benchmark["reference_keff"]["value"]
        ref_uncertainty = benchmark["reference_keff"].get("uncertainty", 0.0)
        sage_keff = sage_response.get("estimated_keff", 0.0)

        deviation = abs(sage_keff - ref_keff)

        # Classify both values
        ref_interpretation = self.classify_keff(ref_keff, ref_uncertainty)
        sage_interpretation = self.classify_keff(sage_keff)

        # Check if interpretations match
        interpretation_correct = ref_interpretation == sage_interpretation

        # Check conservatism
        safety_conservative = self.is_conservative(sage_keff, ref_keff, tolerance)

        # Citation validation (can be extended)
        citations_valid = self._validate_citations(
            benchmark.get("expected_citations", []),
            sage_response.get("citations", []),
        )

        return BenchmarkComparison(
            benchmark_id=benchmark.get("id", "unknown"),
            reference_keff=ref_keff,
            sage_estimate=sage_keff,
            deviation=deviation,
            interpretation_correct=interpretation_correct,
            citations_valid=citations_valid,
            safety_conservative=safety_conservative,
            reference_interpretation=ref_interpretation,
            sage_interpretation=sage_interpretation,
        )

    def _validate_citations(
        self,
        expected: list[str],
        found: list[str],
    ) -> bool:
        """
        Validate that expected citations are present.

        Args:
            expected: List of expected citation strings
            found: List of citations found in response

        Returns:
            True if at least one expected citation is found
        """
        if not expected:
            return True

        found_lower = [c.lower() for c in found]
        for expected_cite in expected:
            if any(expected_cite.lower() in f for f in found_lower):
                return True

        return False

    def evaluate_batch(
        self,
        benchmarks: list[dict[str, Any]],
        sage_responses: list[dict[str, Any]],
        tolerance: float = 0.05,
    ) -> dict[str, Any]:
        """
        Evaluate a batch of benchmark comparisons.

        Args:
            benchmarks: List of benchmark definitions
            sage_responses: Corresponding SAGE responses
            tolerance: Acceptable deviation

        Returns:
            Summary statistics for the batch
        """
        if len(benchmarks) != len(sage_responses):
            raise ValueError("Benchmarks and responses must have same length")

        comparisons = []
        for bench, response in zip(benchmarks, sage_responses):
            comparisons.append(self.compare(bench, response, tolerance))

        # Calculate summary statistics
        total = len(comparisons)
        correct_interpretations = sum(1 for c in comparisons if c.interpretation_correct)
        conservative_count = sum(1 for c in comparisons if c.safety_conservative)
        valid_citations = sum(1 for c in comparisons if c.citations_valid)

        avg_deviation = sum(c.deviation for c in comparisons) / total if total > 0 else 0.0

        return {
            "total_benchmarks": total,
            "correct_interpretations": correct_interpretations,
            "interpretation_accuracy": correct_interpretations / total if total > 0 else 0.0,
            "conservative_count": conservative_count,
            "conservative_rate": conservative_count / total if total > 0 else 0.0,
            "valid_citations": valid_citations,
            "citation_rate": valid_citations / total if total > 0 else 0.0,
            "average_deviation": avg_deviation,
            "comparisons": comparisons,
        }
