"""
src/evaluation/metrics.py
--------------------------
Evaluation metrics for the Scientific RAG project.

Retrieval: Recall@K, Precision@K, MRR, Hit Rate
QA:        Exact Match, F1, Groundedness
Hallucination: hallucination rate, refusal accuracy
"""

import re
import math
from typing import Optional


# ===========================================================================
# Retrieval Metrics
# ===========================================================================

def recall_at_k(retrieved_ids: list, relevant_ids: list, k: int) -> float:
    """Recall@K: fraction of relevant docs found in top-k."""
    if not relevant_ids:
        return 0.0
    top_k = set(retrieved_ids[:k])
    return len(top_k & set(relevant_ids)) / len(relevant_ids)


def precision_at_k(retrieved_ids: list, relevant_ids: list, k: int) -> float:
    """Precision@K: fraction of top-k that are relevant."""
    if k == 0:
        return 0.0
    top_k = retrieved_ids[:k]
    relevant_set = set(relevant_ids)
    return sum(1 for d in top_k if d in relevant_set) / k


def mrr(retrieved_ids: list, relevant_ids: list) -> float:
    """Mean Reciprocal Rank: 1/rank of first relevant doc."""
    relevant_set = set(relevant_ids)
    for rank, doc_id in enumerate(retrieved_ids, start=1):
        if doc_id in relevant_set:
            return 1.0 / rank
    return 0.0


def hit_rate(retrieved_ids: list, relevant_ids: list, k: int) -> float:
    """Hit Rate@K: 1 if any relevant doc is in top-k, else 0."""
    relevant_set = set(relevant_ids)
    return float(any(d in relevant_set for d in retrieved_ids[:k]))


def compute_retrieval_metrics(
    retrieved_ids: list,
    relevant_ids: list,
    k_values: list[int] = (1, 3, 5, 10),
) -> dict:
    """Compute all retrieval metrics for one query."""
    results = {}
    for k in k_values:
        results[f"recall@{k}"] = round(recall_at_k(retrieved_ids, relevant_ids, k), 4)
        results[f"precision@{k}"] = round(precision_at_k(retrieved_ids, relevant_ids, k), 4)
        results[f"hit_rate@{k}"] = round(hit_rate(retrieved_ids, relevant_ids, k), 4)
    results["mrr"] = round(mrr(retrieved_ids, relevant_ids), 4)
    return results


def aggregate_retrieval_metrics(per_query_metrics: list[dict]) -> dict:
    """Average retrieval metrics across all queries."""
    if not per_query_metrics:
        return {}
    keys = per_query_metrics[0].keys()
    return {
        k: round(sum(m[k] for m in per_query_metrics) / len(per_query_metrics), 4)
        for k in keys
    }


# ===========================================================================
# QA Metrics (token-level, following SQuAD convention)
# ===========================================================================

def _normalize(text: str) -> str:
    """Lowercase, strip punctuation, collapse whitespace."""
    text = text.lower()
    text = re.sub(r"[^\w\s]", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def _tokenize(text: str) -> list[str]:
    return _normalize(text).split()


def exact_match(prediction: str, ground_truth: str) -> float:
    return float(_normalize(prediction) == _normalize(ground_truth))


def token_f1(prediction: str, ground_truth: str) -> float:
    pred_tokens = _tokenize(prediction)
    truth_tokens = _tokenize(ground_truth)
    if not pred_tokens or not truth_tokens:
        return float(pred_tokens == truth_tokens)
    common = set(pred_tokens) & set(truth_tokens)
    if not common:
        return 0.0
    precision = len(common) / len(pred_tokens)
    recall = len(common) / len(truth_tokens)
    return 2 * precision * recall / (precision + recall)


def groundedness_score(answer: str, context_chunks: list[str]) -> float:
    """
    Heuristic groundedness: fraction of answer bigrams that appear in context.
    A proper implementation would use an NLI model (e.g. cross-encoder/nli-deberta-v3).
    """
    def bigrams(text: str) -> set:
        tokens = _tokenize(text)
        return set(zip(tokens[:-1], tokens[1:]))

    answer_bigrams = bigrams(answer)
    if not answer_bigrams:
        return 0.0
    context = " ".join(context_chunks)
    context_bigrams = bigrams(context)
    overlap = answer_bigrams & context_bigrams
    return round(len(overlap) / len(answer_bigrams), 4)


def citation_accuracy(answer: str, source_titles: list[str]) -> float:
    """
    Fraction of cited paper titles (or substrings) that appear in the sources list.
    Looks for patterns like [Title, Year] or mentions of paper titles.
    """
    if not source_titles:
        return 0.0
    found = 0
    for title in source_titles:
        # Check for partial title match (first 5 words)
        key = " ".join(title.lower().split()[:5])
        if key in answer.lower():
            found += 1
    return round(found / len(source_titles), 4)


# ===========================================================================
# Hallucination Metrics
# ===========================================================================

REFUSAL_PHRASES = [
    "i cannot answer this question from the available scientific corpus",
    "je ne sais pas a partir des documents fournis",
    "i don't have enough information",
    "not found in the corpus",
    "the context does not contain",
]


def is_refusal(answer: str) -> bool:
    """Check if the model correctly refused an out-of-corpus question."""
    norm = _normalize(answer)
    return any(phrase in norm for phrase in REFUSAL_PHRASES)


def hallucination_rate(
    results: list[dict],
    context_chunks_key: str = "retrieved",
    answer_key: str = "answer",
    threshold: float = 0.2,
) -> float:
    """
    Estimate hallucination rate as fraction of answers with groundedness < threshold.
    """
    if not results:
        return 0.0
    hallucinated = sum(
        1 for r in results
        if groundedness_score(r[answer_key], r[context_chunks_key]) < threshold
    )
    return round(hallucinated / len(results), 4)


def refusal_accuracy(
    out_of_corpus_results: list[dict],
    answer_key: str = "answer",
) -> float:
    """
    Fraction of out-of-corpus questions where the model correctly refused.
    """
    if not out_of_corpus_results:
        return 0.0
    correct = sum(1 for r in out_of_corpus_results if is_refusal(r[answer_key]))
    return round(correct / len(out_of_corpus_results), 4)


# ===========================================================================
# Full Evaluation Runner
# ===========================================================================

def evaluate_qa_result(
    prediction: str,
    ground_truth: str,
    context_chunks: list[str],
    source_titles: list[str] = None,
) -> dict:
    """Compute all QA metrics for one (prediction, ground_truth) pair."""
    return {
        "exact_match": exact_match(prediction, ground_truth),
        "token_f1": round(token_f1(prediction, ground_truth), 4),
        "groundedness": groundedness_score(prediction, context_chunks),
        "citation_accuracy": citation_accuracy(prediction, source_titles or []),
        "is_refusal": is_refusal(prediction),
    }


def print_retrieval_report(agg: dict, label: str = ""):
    print(f"\n{'='*50}")
    print(f"Retrieval Metrics {label}")
    print(f"{'='*50}")
    for k, v in sorted(agg.items()):
        print(f"  {k:<20} {v:.4f}")


def print_qa_report(results: list[dict], label: str = ""):
    print(f"\n{'='*50}")
    print(f"QA Metrics {label}")
    print(f"{'='*50}")
    if not results:
        print("  No results.")
        return
    keys = ["exact_match", "token_f1", "groundedness", "citation_accuracy"]
    for k in keys:
        vals = [r.get(k, 0) for r in results]
        avg = sum(vals) / len(vals)
        print(f"  {k:<25} {avg:.4f}")