"""
Demonstrates how to use the sik-llm-eval framework to evaluate a RAG agent using Nvidia's ChatRAG
benchmark and datasets.

The datasets can be found here: https://huggingface.co/datasets/nvidia/ChatRAG-Bench/

The `source.py` script has functions to download the datasets, process them, and build evals for
sik-llm-eval. The `source.py` script also contains the `SimpleRAGAgent` class which is a simple
implementation of a RAG agent used demonstrate how an agent is evaluated using sik-llm-eval.

The F1 score, precision, and recall have been implemented in the sik-llm-eval framework. The primary
metrics that are used to evaluate the RAG agent are the F1 score and the retrieval accuracy. The
retrieval accuracy is the percentage of times the agent retrieves the correct document. The ground
truth document index is provided in the Nvidia datasets and is compared to the actual index that
the agent retrieves.

An exmample of how to analyze the results can be in
`rag_evals_via_nvidia_chatrag_bench_analysis.ipynb`.

This script can be run via `python -m examples.rag_evals_via_nvidia_chatrag_bench` from
the main directory of the repository.
"""
import time
import uuid
import itertools
from examples.nvidia_chatrag_bench.source import (
    SimpleRAGAgent,
    download_datasets,
    process_combine_datasets,
    build_evals,
)
from sik_llm_eval.candidates import Candidate, CandidateResponse
from sik_llm_eval.eval import CandidateRunResults, EvalHarness, Mode

from dotenv import load_dotenv
load_dotenv()


download_datasets()  # download nvidia datasets from huggingface
process_combine_datasets()  # combine datasets into single, consistant, dataset
build_evals(sample_size_per_dataset=20)  # transform dataset into "evals" for sik-llm-eval

params = {
    'model_name': ('gpt-4o-mini', 'gpt-3.5-turbo-0125'),
    'use_all_messages_in_retrieval': (True, False),
    'system_message': ('', 'Please respond with very concise and relevant information.'),
}
keys, values = zip(*params.items())
param_combinations = [dict(zip(keys, combination)) for combination in itertools.product(*values)]


class CustomRAGAgentCandidate(Candidate):
    """
    This is an example of how a candidate can be created in order to convert the inputs/outputs
    of the RAG agent to the format that the sik-llm-eval framework expects.

    The Eval passes the candidate a dictionary with the keys "documents" and "messages". The Agent
    is expected to store the documents in a vector database and use the messages to generate a
    response. The response is then returned to the Eval. The Eval needs to know the index of the
    document that was used to generate the response in order to check that the correct document
    was retrieved. The eval also needs the generated response so that it can calculate the F1
    score. The Eval expects the information in a specific format it it can consistently parse/score
    the results. Therefore, the candidate must return a dictionary with the keys
    "generated_response" and "relevant_document_id".
    """  # noqa: D404

    def __init__(self, model_name: str, use_all_messages_in_retrieval: bool, system_message: str):
        super().__init__()
        self.model_name = model_name
        self.use_all_messages_in_retrieval = use_all_messages_in_retrieval
        self.system_message = system_message
        self.metadata['uuid'] = str(uuid.uuid4())
        self.metadata['id'] = f"{model_name}|use-all-messages:{use_all_messages_in_retrieval}|has-system-message:{bool(system_message)}"  # noqa: E501
        self.metadata['candidate_model_name'] = self.model_name
        self.metadata['candidate_use_all_messages_in_retrieval'] = self.use_all_messages_in_retrieval  # noqa: E501
        self.metadata['candidate_system_message'] = self.system_message

    def __call__(self, input: dict) -> CandidateResponse:  # noqa: A002
        """
        The Evals expect the candidate to return a dictionary with the keys "generated_response"
        and "relevant_document_id". The "generated_response" is the response generated by the
        agent and will be used to calculate the F1 score. The "relevant_document_id" is the
        index of the document that was used to generate the response. This index is used in the
        MatchCheck for evals that have a ground-truth-document.

        This candidate also returns a "time_dict" in the metadata. This is used to track the time
        taken to generate the response.

        Note that candidate objects should be stateless, meaning it should not retain
        state/information across __call__ invocations because the same candidate object will be
        reused across all evals. Therefore, the code below creates a new SimpleRAGAgent object
        for each __call__ invocation rather than creating it in the __init__ method.
        """
        agent = SimpleRAGAgent(
            model_name=self.model_name,
            use_all_messages_in_retrieval=self.use_all_messages_in_retrieval,
            system_message=self.system_message,
        )
        agent.add_documents(input['documents'])
        relevant_document_id, generated_response, time_dict = agent(input['messages'])
        return CandidateResponse(
            response={
                "generated_response": generated_response,
                "relevant_document_id": str(int(relevant_document_id)) if relevant_document_id is not None else None,  # noqa: E501
            },
            metadata={
                "time_dict": time_dict,
            },
        )

candidates = [CustomRAGAgentCandidate(**params) for params in param_combinations]
harness = EvalHarness(response_mode=Mode.PARALLEL, eval_mode=Mode.PARALLEL)
harness.add_evals_from_files("examples/nvidia_chatrag_bench/evals/*.yaml")
harness.add_candidates(candidates)

print(f"Starting Evals: {len(harness.evals)} evals and {len(harness.candidates)} candidates (total {len(harness.evals) * len(harness.candidates)} eval-candidate pairs)")  # noqa: E501
print(f"Timestamp: {time.strftime('%Y-%m-%d %H:%M:%S')}")
results: list[CandidateRunResults] = harness()
print(f"Finished - Timestamp: {time.strftime('%Y-%m-%d %H:%M:%S')}")

for index, candidate_result in enumerate(results):
    print(f"Candidate {index} has errors: {candidate_result.num_errors > 0}")
    for eval_result, response_error, eval_error in candidate_result:
        if response_error:
            print(f"Response Error: {response_error}")
        if eval_error:
            print(f"Eval Error: {eval_error}")
        if eval_result:
            eval_file = f"examples/nvidia_chatrag_bench/results/{eval_result.eval.metadata['uuid']}-{eval_result.candidate.metadata['uuid']}.yaml"  # noqa: E501
            eval_result.to_yaml(eval_file)
        else:
            print(f"Skipping/error Eval=`{eval_result.eval.metadata['uuid']}` - Candidate=`{eval_result.candidate.metadata['uuid']}`")  # noqa: E501
