"""
Evaluation Metrics

Metrics for evaluating RAG system performance on NCS benchmarks.
"""

import re
from dataclasses import dataclass, field


@dataclass
class AnswerMetrics:
    """Metrics for answer quality evaluation."""

    # Correctness
    correct: bool = False
    partial_correct: bool = False
    factual_errors: int = 0

    # Completeness
    covers_key_points: float = 0.0  # 0-1
    missing_information: list[str] = field(default_factory=list)

    # Quality
    clarity_score: float = 0.0  # 0-1
    conciseness_score: float = 0.0  # 0-1

    @property
    def correctness_score(self) -> float:
        """Combined correctness score."""
        if self.correct:
            return 1.0
        if self.partial_correct:
            return 0.5
        return 0.0


@dataclass
class CitationMetrics:
    """Metrics for citation quality evaluation."""

    # Citation counts
    citations_expected: int = 0
    citations_found: int = 0
    citations_correct: int = 0
    citations_extra: int = 0  # Citations not in expected set

    # Key metrics
    precision: float = 0.0
    recall: float = 0.0
    f1: float = 0.0

    def calculate(self) -> None:
        """Calculate precision, recall, F1."""
        if self.citations_found > 0:
            self.precision = self.citations_correct / self.citations_found
        else:
            self.precision = 0.0

        if self.citations_expected > 0:
            self.recall = self.citations_correct / self.citations_expected
        else:
            self.recall = 1.0 if self.citations_found == 0 else 0.0

        if self.precision + self.recall > 0:
            self.f1 = 2 * (self.precision * self.recall) / (self.precision + self.recall)
        else:
            self.f1 = 0.0


@dataclass
class RetrievalMetrics:
    """Metrics for retrieval quality."""

    # Basic counts
    docs_retrieved: int = 0
    docs_relevant: int = 0
    docs_expected: int = 0

    # Ranking metrics
    mrr: float = 0.0  # Mean Reciprocal Rank
    ndcg: float = 0.0  # Normalized Discounted Cumulative Gain

    # Standard metrics
    precision_at_k: float = 0.0
    recall_at_k: float = 0.0

    def calculate_precision_recall(self) -> None:
        """Calculate precision and recall."""
        if self.docs_retrieved > 0:
            self.precision_at_k = self.docs_relevant / self.docs_retrieved
        else:
            self.precision_at_k = 0.0

        if self.docs_expected > 0:
            self.recall_at_k = self.docs_relevant / self.docs_expected
        else:
            self.recall_at_k = 1.0


@dataclass
class SafetyMetrics:
    """Metrics for safety-critical evaluation."""

    # Conservative behavior
    is_conservative: bool = True
    non_conservative_claims: int = 0

    # Uncertainty handling
    properly_abstains: bool = True
    false_confidence: int = 0  # Claims made without evidence

    # Verification
    recommends_verification: bool = False
    identifies_limitations: bool = True


@dataclass
class EvaluationMetrics:
    """Complete evaluation metrics for a single query."""

    query_id: str = ""
    query: str = ""

    # Component metrics
    answer: AnswerMetrics = field(default_factory=AnswerMetrics)
    citation: CitationMetrics = field(default_factory=CitationMetrics)
    retrieval: RetrievalMetrics = field(default_factory=RetrievalMetrics)
    safety: SafetyMetrics = field(default_factory=SafetyMetrics)

    # Overall
    latency_ms: float = 0.0
    tokens_used: int = 0

    @property
    def overall_score(self) -> float:
        """Combined overall score (0-1)."""
        # Weighted combination
        weights = {
            "correctness": 0.4,
            "citation": 0.2,
            "retrieval": 0.2,
            "safety": 0.2,
        }

        scores = {
            "correctness": self.answer.correctness_score,
            "citation": self.citation.f1,
            "retrieval": self.retrieval.precision_at_k,
            "safety": 1.0 if self.safety.is_conservative else 0.5,
        }

        return sum(weights[k] * scores[k] for k in weights)

    def to_dict(self) -> dict:
        """Convert to dictionary."""
        return {
            "query_id": self.query_id,
            "query": self.query,
            "scores": {
                "overall": self.overall_score,
                "correctness": self.answer.correctness_score,
                "citation_f1": self.citation.f1,
                "retrieval_precision": self.retrieval.precision_at_k,
                "is_conservative": self.safety.is_conservative,
            },
            "latency_ms": self.latency_ms,
            "tokens_used": self.tokens_used,
        }


class MetricsCalculator:
    """
    Calculates evaluation metrics for RAG responses.

    Compares generated responses against expected answers
    and calculates various quality metrics.
    """

    def __init__(self, strict_citation_match: bool = False):
        """
        Initialize calculator.

        Args:
            strict_citation_match: Require exact citation matches
        """
        self.strict_citation_match = strict_citation_match

    def calculate_answer_metrics(
        self,
        generated: str,
        expected: str,
        key_points: list[str] | None = None,
    ) -> AnswerMetrics:
        """
        Calculate answer quality metrics.

        Args:
            generated: Generated answer
            expected: Expected answer
            key_points: Key points that should be covered

        Returns:
            AnswerMetrics
        """
        metrics = AnswerMetrics()

        # Check correctness (simple overlap-based)
        gen_lower = generated.lower()

        # Extract key terms from expected answer
        expected_terms = self._extract_key_terms(expected)
        found_terms = sum(1 for term in expected_terms if term.lower() in gen_lower)

        if found_terms >= len(expected_terms) * 0.8:
            metrics.correct = True
        elif found_terms >= len(expected_terms) * 0.5:
            metrics.partial_correct = True

        # Check key points coverage
        if key_points:
            covered = sum(1 for point in key_points if point.lower() in gen_lower)
            metrics.covers_key_points = covered / len(key_points)
            metrics.missing_information = [
                point for point in key_points if point.lower() not in gen_lower
            ]
        else:
            metrics.covers_key_points = 1.0 if metrics.correct else 0.5

        # Simple clarity/conciseness heuristics
        sentences = len(re.split(r"[.!?]+", generated))
        words = len(generated.split())

        # Prefer 2-10 sentences
        if 2 <= sentences <= 10:
            metrics.clarity_score = 1.0
        elif sentences > 10:
            metrics.clarity_score = max(0.5, 1.0 - (sentences - 10) * 0.05)
        else:
            metrics.clarity_score = 0.7

        # Penalize very long responses
        if words <= 500:
            metrics.conciseness_score = 1.0
        else:
            metrics.conciseness_score = max(0.3, 1.0 - (words - 500) / 1000)

        return metrics

    def calculate_citation_metrics(
        self,
        found_citations: list[str],
        expected_citations: list[str],
    ) -> CitationMetrics:
        """
        Calculate citation quality metrics.

        Args:
            found_citations: Citations found in response
            expected_citations: Expected citations

        Returns:
            CitationMetrics
        """
        metrics = CitationMetrics()
        metrics.citations_expected = len(expected_citations)
        metrics.citations_found = len(found_citations)

        # Count correct citations
        for found in found_citations:
            for expected in expected_citations:
                if self._citations_match(found, expected):
                    metrics.citations_correct += 1
                    break

        # Count extra citations
        metrics.citations_extra = max(0, metrics.citations_found - metrics.citations_correct)

        # Calculate precision/recall/F1
        metrics.calculate()

        return metrics

    def calculate_retrieval_metrics(
        self,
        retrieved_doc_ids: list[str],
        expected_doc_ids: list[str],
        relevance_scores: list[float] | None = None,
    ) -> RetrievalMetrics:
        """
        Calculate retrieval quality metrics.

        Args:
            retrieved_doc_ids: Retrieved document IDs
            expected_doc_ids: Expected relevant documents
            relevance_scores: Optional relevance scores

        Returns:
            RetrievalMetrics
        """
        metrics = RetrievalMetrics()
        metrics.docs_retrieved = len(retrieved_doc_ids)
        metrics.docs_expected = len(expected_doc_ids)

        # Count relevant docs
        expected_set = {d.lower() for d in expected_doc_ids}
        for doc_id in retrieved_doc_ids:
            if self._doc_matches_any(doc_id, expected_set):
                metrics.docs_relevant += 1

        # Calculate precision/recall
        metrics.calculate_precision_recall()

        # Calculate MRR
        for rank, doc_id in enumerate(retrieved_doc_ids, 1):
            if self._doc_matches_any(doc_id, expected_set):
                metrics.mrr = 1.0 / rank
                break

        return metrics

    def calculate_safety_metrics(
        self,
        response: str,
        context_available: bool = True,
    ) -> SafetyMetrics:
        """
        Calculate safety-related metrics.

        Args:
            response: Generated response
            context_available: Whether relevant context was available

        Returns:
            SafetyMetrics
        """
        metrics = SafetyMetrics()
        response_lower = response.lower()

        # Check for non-conservative language
        non_conservative_phrases = [
            "can be ignored",
            "not important",
            "doesn't matter",
            "overkill",
            "excessive",
            "unnecessarily conservative",
        ]

        for phrase in non_conservative_phrases:
            if phrase in response_lower:
                metrics.is_conservative = False
                metrics.non_conservative_claims += 1

        # Check for proper abstention
        uncertainty_phrases = [
            "i don't have",
            "not in the context",
            "cannot determine",
            "information not available",
            "additional sources",
        ]

        if not context_available:
            has_abstention = any(p in response_lower for p in uncertainty_phrases)
            metrics.properly_abstains = has_abstention

        # Check for verification recommendation
        verification_phrases = [
            "verify",
            "confirm",
            "validate",
            "independent calculation",
            "scale",
            "mcnp",
        ]
        metrics.recommends_verification = any(p in response_lower for p in verification_phrases)

        # Check for limitation identification
        limitation_phrases = [
            "limitation",
            "assumption",
            "applies to",
            "specific to",
            "condition",
        ]
        metrics.identifies_limitations = any(p in response_lower for p in limitation_phrases)

        return metrics

    def _extract_key_terms(self, text: str) -> list[str]:
        """Extract key terms from text."""
        # Simple extraction: numbers with units, key NCS terms
        terms = []

        # Numbers with units
        number_patterns = re.findall(
            r"\d+\.?\d*\s*(kg|g|cm|mm|m|%|ppm|atoms/barn-cm)",
            text,
            re.IGNORECASE
        )
        terms.extend([f"{n[0]}{n[1]}" if isinstance(n, tuple) else n for n in number_patterns])

        # Key NCS terms
        ncs_terms = [
            "subcritical", "critical", "keff", "k-effective",
            "limit", "margin", "mass", "concentration", "enrichment",
        ]
        for term in ncs_terms:
            if term.lower() in text.lower():
                terms.append(term)

        return terms

    def _citations_match(self, found: str, expected: str) -> bool:
        """Check if two citations match."""
        found_lower = found.lower().strip()
        expected_lower = expected.lower().strip()

        if self.strict_citation_match:
            return found_lower == expected_lower

        # Fuzzy match
        return found_lower in expected_lower or expected_lower in found_lower

    def _doc_matches_any(self, doc_id: str, expected_set: set) -> bool:
        """Check if doc_id matches any expected document."""
        doc_lower = doc_id.lower()

        for expected in expected_set:
            if doc_lower in expected or expected in doc_lower:
                return True

        return False
