"""
Evaluation Framework - built-in benchmarking for HIEROMEM.

Provides:
1. memory.evaluate() API for easy benchmarking
2. Standard metrics (recall, NDCG, answer_hit)
3. Comparison methods (graph vs dense vs baseline)
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional
import numpy as np


@dataclass
class EvaluationExample:
    """A single evaluation example."""
    query: str
    relevant_ids: list[str]  # IDs of documents that should be retrieved
    answer: Optional[str] = None  # Ground truth answer (if applicable)
    query_vector: Optional[list[float]] = None  # Cached embedding
    metadata: dict = field(default_factory=dict)


@dataclass
class EvaluationResult:
    """Results from evaluating a retrieval method."""
    method: str
    recall_at_k: float
    precision_at_k: float
    ndcg_at_k: float
    answer_hit_rate: float
    mrr: float  # Mean Reciprocal Rank
    k: int
    num_examples: int
    per_example: list[dict] = field(default_factory=list)

    def __repr__(self) -> str:
        return (
            f"EvaluationResult({self.method}: "
            f"R@{self.k}={self.recall_at_k:.3f}, "
            f"P@{self.k}={self.precision_at_k:.3f}, "
            f"NDCG@{self.k}={self.ndcg_at_k:.3f}, "
            f"Answer={self.answer_hit_rate:.3f})"
        )


@dataclass
class ComparisonResult:
    """Results from comparing multiple retrieval methods."""
    results: dict[str, EvaluationResult]
    winner: str
    improvement: float  # Improvement of winner over baseline

    def summary(self) -> str:
        """Human-readable summary."""
        lines = ["=" * 60, "COMPARISON RESULTS", "=" * 60]
        for name, result in self.results.items():
            lines.append(f"{name}: R@{result.k}={result.recall_at_k:.3f}")
        lines.append("-" * 60)
        lines.append(f"Winner: {self.winner} (+{self.improvement:.1%})")
        return "\n".join(lines)


class Evaluator:
    """
    Evaluation engine for HIEROMEM retrieval.

    Usage:
        evaluator = Evaluator(memory)
        result = evaluator.evaluate(test_set, k=5)
        comparison = evaluator.compare(test_set, methods=['graph', 'dense'])
    """

    def __init__(self, memory):
        """
        Initialize evaluator with a Memory instance.

        Args:
            memory: HIEROMEM Memory instance
        """
        self.memory = memory

    def evaluate(
        self,
        examples: list[EvaluationExample],
        k: int = 10,
        use_graph: bool = True,
        alpha_dense: float = 1.0,
        beta_symbolic: float = 50.0,
    ) -> EvaluationResult:
        """
        Evaluate retrieval on a test set.

        Args:
            examples: List of evaluation examples
            k: Number of results to retrieve
            use_graph: Whether to use graph-enhanced retrieval
            alpha_dense: Dense weight for hybrid retrieval
            beta_symbolic: Symbolic weight for hybrid retrieval

        Returns:
            EvaluationResult with metrics
        """
        from hieromem_core.encoder import Encoder, EncoderConfig

        # Use memory's encoder if available, else create default
        encoder = getattr(self.memory, 'encoder', None) or Encoder(EncoderConfig())

        recalls = []
        precisions = []
        ndcgs = []
        answer_hits = []
        mrrs = []
        per_example = []

        for example in examples:
            query_vec = encoder.encode(example.query)
            # Ensure 1D vector (encoder may return 2D)
            if hasattr(query_vec, 'squeeze'):
                query_vec = query_vec.squeeze()

            # Build query graph if using graph enhancement
            query_graph = None
            if use_graph:
                query_graph = self.memory.build_query_graph(example.query)

            # Retrieve
            result = self.memory.dual_retrieve(
                query_vec,
                query_graph=query_graph,
                k=k,
                alpha_dense=alpha_dense,
                beta_symbolic=beta_symbolic,
            )

            # Get retrieved IDs
            retrieved_ids = [ep.id for ep in result.episodes]

            # Calculate metrics
            recall = self._recall_at_k(retrieved_ids, example.relevant_ids, k)
            precision = self._precision_at_k(retrieved_ids, example.relevant_ids, k)
            ndcg = self._ndcg_at_k(retrieved_ids, example.relevant_ids, k)
            mrr = self._mrr(retrieved_ids, example.relevant_ids)

            # Check answer hit
            answer_hit = 0.0
            if example.answer:
                for ep in result.episodes:
                    content = ep.meta.get('content_preview', '')
                    if example.answer.lower() in content.lower():
                        answer_hit = 1.0
                        break

            recalls.append(recall)
            precisions.append(precision)
            ndcgs.append(ndcg)
            answer_hits.append(answer_hit)
            mrrs.append(mrr)
            per_example.append({
                'query': example.query,
                'recall': recall,
                'precision': precision,
                'ndcg': ndcg,
                'answer_hit': answer_hit,
                'retrieved': retrieved_ids[:3],
            })

        method = "graph_enhanced" if use_graph else "dense_only"
        return EvaluationResult(
            method=method,
            recall_at_k=float(np.mean(recalls)),
            precision_at_k=float(np.mean(precisions)),
            ndcg_at_k=float(np.mean(ndcgs)),
            answer_hit_rate=float(np.mean(answer_hits)),
            mrr=float(np.mean(mrrs)),
            k=k,
            num_examples=len(examples),
            per_example=per_example,
        )

    def compare(
        self,
        examples: list[EvaluationExample],
        k: int = 10,
        methods: Optional[list[str]] = None,
    ) -> ComparisonResult:
        """
        Compare multiple retrieval methods.

        Args:
            examples: List of evaluation examples
            k: Number of results to retrieve
            methods: Methods to compare ('graph', 'dense')

        Returns:
            ComparisonResult with per-method results and winner
        """
        if methods is None:
            methods = ['graph', 'dense']

        results = {}

        if 'dense' in methods:
            results['dense'] = self.evaluate(
                examples, k=k, use_graph=False,
                alpha_dense=1.0, beta_symbolic=0.0
            )

        if 'graph' in methods:
            results['graph'] = self.evaluate(
                examples, k=k, use_graph=True,
                alpha_dense=1.0, beta_symbolic=50.0
            )

        # Find winner based on recall
        baseline = results.get('dense', list(results.values())[0])
        best_result = max(results.values(), key=lambda r: r.recall_at_k)
        winner = best_result.method
        improvement = (
            (best_result.recall_at_k - baseline.recall_at_k) /
            max(baseline.recall_at_k, 0.001)
        )

        return ComparisonResult(
            results=results,
            winner=winner,
            improvement=improvement,
        )

    # Metric implementations

    def _recall_at_k(
        self, retrieved: list[str], relevant: list[str], k: int
    ) -> float:
        """Recall@k: fraction of relevant docs retrieved in top-k."""
        if not relevant:
            return 0.0
        retrieved_set = set(retrieved[:k])
        relevant_set = set(relevant)
        return len(retrieved_set & relevant_set) / len(relevant_set)

    def _precision_at_k(
        self, retrieved: list[str], relevant: list[str], k: int
    ) -> float:
        """Precision@k: fraction of top-k that are relevant."""
        if k == 0:
            return 0.0
        retrieved_k = retrieved[:k]
        relevant_set = set(relevant)
        return sum(1 for r in retrieved_k if r in relevant_set) / k

    def _ndcg_at_k(
        self, retrieved: list[str], relevant: list[str], k: int
    ) -> float:
        """NDCG@k: Normalized Discounted Cumulative Gain."""
        if not relevant:
            return 0.0

        relevant_set = set(relevant)
        dcg = 0.0
        for i, doc_id in enumerate(retrieved[:k]):
            if doc_id in relevant_set:
                dcg += 1.0 / np.log2(i + 2)  # i+2 because position is 1-indexed

        # Ideal DCG (all relevant docs at top)
        idcg = sum(1.0 / np.log2(i + 2) for i in range(min(len(relevant), k)))

        return dcg / idcg if idcg > 0 else 0.0

    def _mrr(self, retrieved: list[str], relevant: list[str]) -> float:
        """Mean Reciprocal Rank: 1/position of first relevant result."""
        relevant_set = set(relevant)
        for i, doc_id in enumerate(retrieved):
            if doc_id in relevant_set:
                return 1.0 / (i + 1)
        return 0.0


def create_example(
    query: str,
    relevant_ids: list[str],
    answer: Optional[str] = None,
) -> EvaluationExample:
    """Helper to create an evaluation example."""
    return EvaluationExample(
        query=query,
        relevant_ids=relevant_ids,
        answer=answer,
    )
