import os
import pickle
import numpy as np
import faiss
import torch
from datasets import load_dataset
import evaluate

# Import RAG setup and retrieval logic from app.py
from app import setup_rag, retrieve


def retrieval_recall(dataset, passages, embedder, index, k=20, rerank_k=None, num_samples=100):
    """
    Compute raw Retrieval Recall@k on the first num_samples examples.
    If rerank_k is set, also apply cross-encoder reranking.
    """
    hits = 0
    for ex in dataset.select(range(num_samples)):
        question = ex["question"]
        gold_answers = ex["answers"]["text"]
        # get top-k retrieved contexts
        if rerank_k:
            ctxs, _ = retrieve(question, passages, embedder, index, k=k, rerank_k=rerank_k)
        else:
            # skip reranking: use top-k directly
            q_emb = embedder.encode([question], convert_to_numpy=True)
            distances, idxs = index.search(q_emb, k)
            ctxs = [passages[i] for i in idxs[0]]
        # check if any gold span appears
        if any(any(ans in ctx for ctx in ctxs) for ans in gold_answers):
            hits += 1
    recall = hits / num_samples
    print(f"Retrieval Recall@{k} (rerank_k={rerank_k}): {recall:.3f} ({hits}/{num_samples})")
    return recall


def retrieval_recall_answerable(dataset, passages, embedder, index, k=20, rerank_k=None, num_samples=100):
    """
    Retrieval Recall@k evaluated only on answerable questions.
    """
    hits, total = 0, 0
    for ex in dataset.select(range(num_samples)):
        if not ex["answers"]["text"]:
            continue
        total += 1
        question = ex["question"]
        if rerank_k:
            ctxs, _ = retrieve(question, passages, embedder, index, k=k, rerank_k=rerank_k)
        else:
            q_emb = embedder.encode([question], convert_to_numpy=True)
            distances, idxs = index.search(q_emb, k)
            ctxs = [passages[i] for i in idxs[0]]
        if any(any(ans in ctx for ctx in ctxs) for ans in ex["answers"]["text"]):
            hits += 1
    recall = hits / total if total > 0 else 0.0
    print(f"Retrieval Recall@{k} on answerable only (rerank_k={rerank_k}): {recall:.3f} ({hits}/{total})")
    return recall


def qa_eval_answerable(dataset, passages, embedder, reranker, index, qa_pipe, k=20, num_samples=100):
    """
    End-to-end QA EM/F1 on answerable subset using the retrieve_and_answer logic.
    """
    squad_metric = evaluate.load("squad")
    preds, refs = [], []
    for ex in dataset.select(range(num_samples)):
        if not ex["answers"]["text"]:
            continue
        qid = ex["id"]
        # retrieve and generate
        answer, _ = retrieve_and_answer(ex["question"], passages, embedder, reranker, index, qa_pipe)
        preds.append({"id": qid, "prediction_text": answer})
        refs.append({"id": qid, "answers": ex["answers"]})
    results = squad_metric.compute(predictions=preds, references=refs)
    print(f"Answerable-only QA EM: {results['exact_match']:.2f}, F1: {results['f1']:.2f}")
    return results


def main():
    # Setup RAG components
    passages, embedder, reranker, index, qa_pipe = setup_rag()
    # Load SQuAD v2 validation set
    squad = load_dataset("rajpurkar/squad_v2", split="validation")

    # Run evaluations
    retrieval_recall(squad, passages, embedder, index, k=20, rerank_k=5, num_samples=100)
    retrieval_recall_answerable(squad, passages, embedder, index, k=20, rerank_k=5, num_samples=100)
    qa_eval_answerable(squad, passages, embedder, reranker, index, qa_pipe, k=20, num_samples=100)


if __name__ == "__main__":
    main()