"""
An Experiment is a collection of queries, and manages their retrieved documents, agent answers and their respective
evaluations created by Evaluators.
For most use cases, this class should be the main entry point to interact with RAGElo.
"""

from __future__ import annotations

import csv
import json
import logging
import os
import warnings
from collections import defaultdict
from pathlib import Path
from typing import Any, Literal, overload

import rich
from pydantic import BaseModel

from ragelo.logger import CLILogHandler, logger
from ragelo.types.evaluables import AgentAnswer, Document
from ragelo.types.query import Query
from ragelo.types.results import AnswerEvaluatorResult, EloTournamentResult, EvaluatorResult, RetrievalEvaluatorResult


class Experiment:
    """
    A class to manage and evaluate an experiment. An experiment is composed of multiple queries and their evaluations.
    This is the main class for interacting with RAGElo, and manages all the queries to which agents retrieve documents
      and answer questions.
    Attributes:
        experiment_name (str): An unique name for your experiment.
        queries: (dict[str, Query]): A dictionary with the queries for this experiment.
            Can also be accessed directly with __getitem__.
        elo_tournaments: (list[EloTournamentResult]): A list with the results for Elo Tournaments
            played over this set of queries.
    Methods:
        add_query(query, query_id, doc_id, metadata, force): Adds a query
            (either a Query object or a string) to the current experiment. Returns the id of the added query.
        add_retrieved_doc(doc, query_id, doc_id, score, agent, force): Adds a retrieved document to a query
        add_agent_answer(answer, agent, query_id, force): Adds an agent's answer to a query.
        add_evaluation(evaluation, should_save, force): Adds an evaluation to a retrieved document,
            an answer or a pairwise game.
        evaluate_retrieval(metrics, relevance_threshold): Evaluates the retrieval performance of the agents
            using the specified metrics and the evaluations generated by a `RetrievalEvaluator`.
            This requires the ir_measures package to be installed.
        get_qrels(relevance_key, relevance_threshold, output_path, output_format): Gets the qrels
            (dictionary of relevance) for the queries. Documents with a score above relevance_threshold are considered
            relevant. Optionally saves the result to output_path with output_format.
        get_runs(agents, output_path, output_format): Gets the runs (retrieval results) for the provided agents for
            all the queries. Optionally saves the result to output_path with output_format.
        save(): Saves the current state of the queries object to disk as a JSON object to `cache_path`.
        save_results(result): Saves an evaluation result to the `evaluations_cache_path` file.
        add_retrieved_docs_from_runfile(run_file_path, corpus, top_k): Adds retrieved documents from a trec-formatted
            run file to the queries. Will create the Document objects as needed.

        __len__(): Returns the number of queries.
        __iter__(): Returns an iterator over the queries.
        __getitem__(key: str) -> Query: Gets a query by its key.
        keys(): Returns the keys (query_ids) of all queries.
    """

    def __init__(
        self,
        experiment_name: str,
        save_path: str | None = None,
        evaluations_cache_path: str | None = None,
        cache_evaluations: bool = True,
        save_on_disk: bool = True,
        verbose: bool = False,
        rich_print: bool = True,
        clear_evaluations: bool = False,
        queries_csv_path: str | None = None,
        csv_query_id_col: str | None = None,
        csv_query_text_col: str = "query",
        documents_csv_path: str | None = None,
        csv_document_id_col: str | None = None,
        csv_document_text_col: str = "document",
        answers_csv_path: str | None = None,
        csv_agent_col: str = "agent",
        csv_answer_text_col: str = "answer",
    ):
        """
        Initialize the experiment object with optional data and set up paths and caches.

        Args:
            experiment_name (str): The name of your experiment
            save_path (Optional[str]): A JSON file path to persist the experiment on disk.
                If not set, results will not be saved locally. If it already exists,
                we will try to load the state from that file.
            evaluations_cache_path (Optional[str]): A JSON Lines file path to persist the evaluators result on disk,
                to avoid re-computing evaluations. If not set, will create one based on experiment_name.
            cache_evaluations (bool, defaults to True): Whether or not to cache the evaluations on evaluations_cache_path,
            save_on_disk (bool, defaults to True): Whether to save the experiment to disk using the experiment_name.
            verbose (bool, defaults to False): Whether to print information about the experiment.
            rich_print (bool, defaults to True): Whether to print information about the experiment using the rich library.
            clear_evaluations (bool, defaults to False): If set to True, will clear all existing evaluations and
                re-compute as needed.
            queries_csv_path (Optional[str]): An optional path to a CSV file with the queries to be used.
            csv_query_id_col (Optional[str]): The column name for query ID in the CSV files.
            csv_query_text_col (str): The column name for query text in the queries CSV file.
            documents_csv_path (Optional[str]): An optional path to a CSV file with the retrieved documents for
                the queries.
            csv_document_id_col (Optional[str]): The column name for document ID in the documents CSV file.
            csv_document_text_col (str): The column name for document text in the documents CSV file.
            answers_csv_path (Optional[str]): An optional path to a CSV file with the answers for the queries.
            csv_agent_col (str, defaults to "agent"): The column name for agent name in the documents and answers
                CSV files.
            csv_answer_text_col (str, defaults to "answer"): The column name for answer text in the answers CSV file.



        The initialization process involves setting up paths and loading queries from different sources:
        1. If `csv_path` is provided, queries are loaded from the specified CSV file.
        2. If `csv_path` is not provided but `save_path` exists and points to a valid file, queries are loaded from
            the cache file.
        3. If both `csv_path` and `save_path` are provided, the method ensures that the queries loaded from the CSV
            file match those in the cache file.
        The method prioritizes loading queries from the CSV file over the cache file. If neither source is available,
            it relies on the `queries` attribute if it has been set.
        """
        self.cache_evaluations = cache_evaluations
        self.evaluations_cache_path = evaluations_cache_path
        self.experiment_name = experiment_name
        self.queries: dict[str, Query] = {}
        self.elo_tournaments: list[EloTournamentResult] = []
        self.rich_print = rich_print
        self.save_on_disk = save_on_disk
        self.save_path = save_path
        if verbose:
            logger.setLevel(logging.INFO)
            if len(logger.handlers) == 0:
                logger.addHandler(CLILogHandler(use_rich=rich_print))
        elif len(logger.handlers) > 0:
            logger.setLevel(logging.WARNING)
            logger.handlers = []
            logger.addHandler(CLILogHandler(use_rich=rich_print))
        if self.save_on_disk and not self.save_path:
            self.save_path = f"ragelo_cache/{self.experiment_name}.json"
            if not os.path.exists("ragelo_cache"):
                logger.info(f"Creating a cache file for the experiment at ragelo_cache/{self.experiment_name}.json")
                Path("ragelo_cache").mkdir(parents=True, exist_ok=True)

        elif self.save_path is None and not self.save_on_disk:
            logger.info("No save path provided. Will not save the experiment to disk.")

        if self.cache_evaluations and not self.evaluations_cache_path:
            if self.save_path:
                self.evaluations_cache_path = self.save_path.replace(".json", "_results.jsonl")
            else:
                self.evaluations_cache_path = f"ragelo_cache/{self.experiment_name}_results.jsonl"
                if not os.path.exists("ragelo_cache"):
                    os.makedirs("ragelo_cache", exist_ok=True)
            if not os.path.isfile(self.evaluations_cache_path):
                logger.info(f"Creating a cache file for the experiment's evaluations at{self.evaluations_cache_path}")
                Path(self.evaluations_cache_path).touch()

        if self.save_path and os.path.isfile(self.save_path):
            self.queries = self._load_from_cache(self.save_path)
        else:
            self.queries = {}
        if queries_csv_path:
            self.add_queries_from_csv(queries_csv_path, csv_query_id_col, csv_query_text_col, exist_ok=True)
        if documents_csv_path:
            self.add_documents_from_csv(
                documents_csv_path,
                csv_document_id_col,
                csv_document_text_col,
                csv_query_id_col,
                csv_agent_col,
                exist_ok=True,
            )
        if answers_csv_path:
            self.add_agent_answers_from_csv(
                answers_csv_path, csv_agent_col, csv_answer_text_col, csv_query_id_col, exist_ok=True
            )

        if self.evaluations_cache_path and os.path.isfile(self.evaluations_cache_path):
            self._load_results_from_cache(self.evaluations_cache_path)
        if clear_evaluations:
            self.__clear_all_evaluations()
        self.save()

    def add_query(
        self,
        query: Query | str,
        query_id: str | None = None,
        metadata: dict | None = None,
        force: bool = False,
    ) -> str:
        """
        Adds a query to the collection of queries.
        Args:
            query (Query | str): The query to be added. It can be an instance of the Query class or a string.
            query_id (str | None, optional): The identifier for the query.
                If not provided, a default ID will be generated.
            metadata (dict | None, optional): Additional metadata for the query. Defaults to None.
            force (bool, optional): Whether to overwrite the query if it already exists. Defaults to False.
        """

        if isinstance(query, Query):
            if query_id is not None and query.qid != query_id:
                warnings.warn(f"Query ID mismatch. Using ID from query object: {query.qid}")
            query_id = query.qid
            query_obj = query
        elif query_id is None:
            query_id = f"query_{len(self.queries) + 1}"
            warnings.warn(f"Query ID not provided. Using default query ID: {query_id}.")
            query_obj = Query(qid=query_id, query=query, metadata=metadata)
        else:
            query_obj = Query(qid=query_id, query=query, metadata=metadata)
            query_id = query_obj.qid
        if query_id in self.queries and not force:
            logger.info(f'Query with ID "{query_id}" already exists. Use force=True to overwrite')
            return query_id
        if query_id in self.queries and force:
            logger.info(f'Query with ID "{query_id}" already exists, but force was set to True. Overwriting.')
        self.queries[query_id] = query_obj
        return query_id

    def add_retrieved_doc(
        self,
        doc: Document | str,
        query_id: str | None = None,
        doc_id: str | None = None,
        score: float | None = None,
        agent: str | None = None,
        force: bool = False,
        exist_ok: bool = False,
    ):
        """
        Adds a retrieved document to a query.
        Args:
            doc (Document | str): The document to be added. It can be an instance of the Document class or a string.
            query_id (Optional[str]): The query used to retrieve the document.
            doc_id (Optional[str]): The identifier for the document. If not provided, a default ID will be generated.
            score (Optional[float]): The score of the document. Defaults to None.
            agent (Optional[str]): The agent that retrieved the document. Defaults to None.
            force (bool): Whether to overwrite the document if it already exists. Defaults to False.
            exist_ok (bool): If True, will not raise an error if the document already exists.
                Defaults to False.
        """
        if isinstance(doc, Document):
            if doc.qid != query_id and query_id is not None:
                logger.warning(f"Query ID mismatch. Using ID from document object: {doc.qid}")
                query_id = doc.qid
            if doc.did != doc_id and doc_id is not None:
                logger.warning(f"Document ID mismatch. Using ID from document object: {doc.did}")
            doc_id = doc.did
            query_id = doc.qid
        elif doc_id is None:
            raise ValueError("Document was provided as string, but Document ID was not provided")
        if query_id is None:
            raise ValueError(f"ID of the query to which document {doc_id} was retrieved was not provided.")
        if isinstance(doc, str):
            doc = Document(qid=query_id, did=doc_id, text=doc)
            query_id = doc.qid
        if query_id not in self.queries:
            raise ValueError(f"Query {query_id} not found in queries")
        self.queries[query_id].add_retrieved_doc(doc, score, agent, force, exist_ok)

    def add_agent_answer(
        self,
        answer: AgentAnswer | str,
        agent: str | None = None,
        query_id: str | None = None,
        force: bool = False,
        exist_ok: bool = False,
    ):
        if isinstance(answer, AgentAnswer):
            if answer.qid != query_id and query_id is not None:
                logger.warning(f"Query ID mismatch. Using ID from answer object: {answer.qid}")
            query_id = answer.qid
            if answer.agent != agent and agent is not None:
                logger.warning(f"Agent mismatch. Using agent from answer object: {answer.agent}")
        elif query_id is None:
            raise ValueError("Query ID not provided")
        else:
            if agent is None:
                agent = f"agent_{len(self.queries[query_id].answers) + 1}"
                logger.warning(f"Agent not provided. Using default agent name: {agent}.")
            if query_id is None:
                raise ValueError("Query ID not provided")
            answer = AgentAnswer(qid=query_id, agent=agent, text=answer)
        query_id = answer.qid
        if query_id not in self.queries:
            raise ValueError(f"Query {query_id} not found in queries")
        self.queries[query_id].add_agent_answer(answer, force=force, exist_ok=exist_ok)

    @overload
    def add_evaluation(
        self,
        evaluation: RetrievalEvaluatorResult,
        should_save: bool = True,
        should_print: bool = True,
        force: bool = False,
        exist_ok: bool = False,
    ): ...

    @overload
    def add_evaluation(
        self,
        evaluation: AnswerEvaluatorResult,
        should_save: bool = True,
        should_print: bool = True,
        force: bool = False,
        exist_ok: bool = False,
    ): ...

    @overload
    def add_evaluation(
        self,
        evaluation: EloTournamentResult,
        should_save: bool = True,
        should_print: bool = True,
        force: bool = False,
        exist_ok: bool = False,
    ): ...

    def add_evaluation(
        self,
        evaluation: RetrievalEvaluatorResult | AnswerEvaluatorResult | EloTournamentResult,
        should_save: bool = True,
        should_print: bool = False,
        force: bool = False,
        exist_ok: bool = False,
    ):
        """
        Add an evaluation to the queries and optionally save the result.
        Args:
            evaluation (RetrievalEvaluatorResult | AnswerEvaluatorResult | EloTournamentResult): The evaluation result to be added.
            should_save (bool): Whether to save the result to disk. Defaults to True.
            should_print (bool): Whether to print the result. Defaults to True.
            force (bool): Whether to overwrite an existing evaluation. Defaults to False.
            exist_ok (bool): Whether to warn if an evaluation already exists. Defaults to False.
        """
        if isinstance(evaluation, EloTournamentResult):
            self.elo_tournaments.append(evaluation)
            added = True
        else:
            qid = evaluation.qid
            if qid not in self.queries:
                raise ValueError(f"Query {qid} not found in queries")
            added = self.queries[qid].add_evaluation(evaluation, force=force, exist_ok=exist_ok)
        if not added:
            return
        if should_save:
            self.save_results(evaluation)
        if not should_print:
            return
        if isinstance(evaluation, EloTournamentResult):
            self._print_elo_tournament_result(evaluation)
        else:
            self._print_response(evaluation)

    def _print_elo_tournament_result(self, result: EloTournamentResult):
        scores = sorted(result.scores.items(), key=lambda x: x[1], reverse=True)
        if self.rich_print:
            rich.print("-------[bold white] Agents Elo Ratings [/bold white]-------")
        else:
            print("------- Agents Elo Ratings -------")
        for agent, rating in scores:
            std_dev = result.std_dev.get(agent, 0)
            if self.rich_print:
                rich.print(f"[bold white]{agent:<15}[/bold white]: {rating:.1f}(±{std_dev:.1f})")
            else:
                print(f"{agent:<15}: {rating:.1f} (±{std_dev:.1f})")

    def _print_response(
        self,
        response: EvaluatorResult,
    ):
        answer: float | str | dict | int
        if isinstance(response.answer, dict):
            # Print the answer in a more readable format
            answer = json.dumps(response.answer, indent=4, ensure_ascii=False)
        elif isinstance(response.answer, BaseModel):
            answer = response.answer.model_dump_json(indent=4)
        else:
            answer = str(response.answer)
        response_dict = response.model_dump()
        agent_a = response_dict.get("agent_a")
        agent_b = response_dict.get("agent_b")
        agent = response_dict.get("agent")
        qid = response_dict.get("qid")
        did = response_dict.get("did")
        raw_answer = response_dict.get("raw_answer")

        if self.rich_print:
            rich.print(f"[bold blue]🔎 Query ID[/bold blue]: {qid}")
            did = response_dict.get("did")
            if did:
                rich.print(f"[bold blue]📜 Document ID[/bold blue]: {did}")
            if agent_a and agent_b:
                rich.print(f"[bold bright_cyan] {agent_a:<18} [/bold bright_cyan] 🆚  [bold red] {agent_b}[/bold red]")
            elif agent:
                rich.print(f"[bold bright_cyan]🕵️ Agent[/bold bright_cyan]: {agent}")
            rich.print(f"[bold blue]Parsed Answer[/bold blue]: {answer}")
            rich.print("")
            return
        print(f"Query ID: {qid}")
        if did:
            print(f"Document ID: {did}")
        if agent_a and agent_b:
            print(f"{agent_a} vs {agent_b}")
        elif agent:
            print(f"Agent: {agent}")
        print(f"Raw Answer: {raw_answer}")
        print(f"Parsed Answer: {answer}")
        print("")

    def evaluate_retrieval(
        self,
        relevance_key: str = "relevance",
        metrics: list[str] = ["Precision@10", "nDCG@10", "Judged@10"],
        relevance_threshold: int = 0,
    ) -> dict[str, dict[str, float]]:
        """
        Evaluate the retrieval performance of agents using specified metrics.
        Args:
            relevance_key (str): The key in the answer object that contains an integer with the relevance of the
                document.
            metrics (list[str]): A list of metrics to use. defaults to ["Precision@10", "nDCG@10", "Judged@10"].
            relevance_threshold (int): The threshold above which a document is considered relevant (default is 0).
        Returns:
            dict[str, dict[str, float]]: A dictionary where keys are agent names and values are dictionaries of scores.
        Raises:
            ImportError: If the `ir_measures` package is not installed.
            ValueError: If a specified metric is not found in the `ir_measures` registry.
        Notes:
        - The function uses the `ir_measures` package to calculate the metrics.
        - If the `rich` package is installed, the results are printed in a formatted table.
        - If the `rich` package is not installed, the results are printed using plain `print`.

        Example:
        >>> results = evaluate_retrieval(relevance_key="relevance", metrics=["Precision@10", "nDCG@10"], relevance_threshold=1)
        >>> print(results)
        {
            "agent1": {
                "Precision@10": 0.5,
                "nDCG@10": 0.6
            },
            "agent2": {
                "Precision@10": 0.4,
                "nDCG@10": 0.5
            }
        }
        """

        try:
            import ir_measures
            from ir_measures import parse_measure
        except ImportError:
            raise ImportError("ir_measures is not installed. Please install it with `pip install ir-measures`")
        qrels = self.get_qrels(
            relevance_key=relevance_key,
            relevance_threshold=relevance_threshold,
        )
        runs = self.get_runs()
        measures = []
        for metric in metrics:
            try:
                measure = parse_measure(metric)
            except NameError:
                valid_metrics = list(ir_measures.measures.registry.keys())
                raise ValueError(f"Metric {metric} not found. Valid metrics are: {valid_metrics}")
            measures.append(measure)
        metrics = [str(m) for m in measures]
        results = {}
        for agent, run in runs.items():
            # Transform the keys of the results back to strings
            results[agent] = {
                str(k): v
                for k, v in ir_measures.calc_aggregate(measures, qrels, run).items()  # type: ignore
            }

        key_metric = metrics[0]
        max_agent_len = max([len(agent) for agent in results.keys()]) + 3
        max_metric_len = max([len(metric) for metric in metrics])
        sorted_agents = sorted(results.items(), key=lambda x: x[1][key_metric], reverse=True)
        if self.rich_print:
            rich.print("---[bold cyan] Retrieval Scores [/bold cyan] ---")
            if relevance_threshold > 0:
                rich.print(f"[bold yellow]Relevance threshold: {relevance_threshold}[/bold yellow]")
            header = f"[bold magenta]{'Agent Name':<{max_agent_len}}"
            header += "\t".join([f"{m:<{max_metric_len}}" for m in metrics])
            header += "[/bold magenta]"
            rich.print(f"[bold cyan]{header}[/bold cyan]")
            for agent, scores in sorted_agents:
                row = f"[bold white]{agent:<{max_agent_len}}[/bold white]"
                row += "\t".join([f"{scores[metric]:<{max_metric_len},.4f}" for metric in metrics])
                rich.print(row)

        else:
            print(results)
        return results

    def get_qrels(
        self,
        relevance_key: str | None = "relevance",
        relevance_threshold: int = 0,
        output_path: str | None = None,
        output_format: Literal["trec", "json"] = "trec",
    ) -> dict[str, dict[str, int]]:
        """
        Retrieve the qrels (query relevance judgments) for the queries.
        Args:
            relevance_key (str | None): The key to use for determining relevance. Defaults to "relevance".
            relevance_threshold (int): The threshold value for relevance. Defaults to 0.
        Returns:
            dict[str, dict[str, int]]: A dictionary where each key is a query ID and the value is another dictionary
                                       mapping document IDs to their relevance scores.
        """

        qrels = {}
        for qid, query in self.queries.items():
            qrels[qid] = query.get_qrels(relevance_key, relevance_threshold)
        if output_path:
            with open(output_path, "w") as f:
                if output_format.lower() == "trec":
                    for qid, qrel in qrels.items():
                        for did, rel in qrel.items():
                            f.write(f"{qid} Q0 {did} {rel}\n")
                elif output_format.lower() == "json":
                    json.dump(qrels, f)
                else:
                    raise ValueError(
                        f"Invalid output format for QRELS: {output_format}Valid options are 'trec' and 'json'"
                    )
        return qrels

    def get_runs(
        self,
        agents: list[str] | None = None,
        output_path: str | None = None,
        output_format: str = "trec",
    ) -> dict[str, dict[str, dict[str, float]]]:
        """
        Retrieve runs for specified agents.
        This method aggregates runs from all queries and organizes them by agent.
        Args:
            agents (list[str] | None): A list of agent names to filter the runs.
                If None, runs for all agents are retrieved.
            output_path (str | None): The path to save the runs to disk. If None, the runs are not saved.
            output_format (str): The format to save the runs. Defaults to "trec". If TREC, output_path should be a
                directory, and each agent will have a separate file.
        Returns:
            dict[str, dict[str, dict[str, float]]]: A nested dictionary where the first level keys are agent names,
            the second level keys are run identifiers, and the third level keys are metric names with their
                corresponding float values.
        """

        runs_by_agent: dict[str, dict[str, dict[str, float]]] = defaultdict(dict)
        for query in self.queries.values():
            runs = query.get_runs(agents)
            for agent, run in runs.items():
                runs_by_agent[agent].update(run)
        if output_path:
            if output_format.lower() == "trec":
                if not os.path.isdir(output_path):
                    os.makedirs(output_path, exist_ok=True)
                for agent, run in runs_by_agent.items():
                    with open(f"{output_path}/{agent}.run", "w") as f:
                        for qid, docs in run.items():
                            sorted_scores = sorted(docs.items(), key=lambda x: x[1], reverse=True)
                            for idx, (did, score) in enumerate(sorted_scores):
                                f.write(f"{qid} Q0 {did} {idx + 1} {score} {agent}\n")
            elif output_format.lower() == "json":
                with open(output_path, "w") as f:
                    json.dump(runs_by_agent, f)

        return runs_by_agent

    def save(self, output_path: str | None = None):
        """
        Persist the current state of the model to disk if caching is enabled.
        This method checks if caching is enabled by evaluating the `save_cache` attribute.
        If caching is enabled, it attempts to write the model's current state to the file
        specified by `cache_path`. If `cache_path` is not set, a ValueError is raised.
        Raises:
            ValueError: If `cache_path` is None and caching is enabled.
        """
        if not self.save_on_disk:
            return
        if output_path is None:
            if self.save_path is None:
                raise ValueError("Cannot save experiment without a save path")
            output_path = self.save_path
        output_dict: dict[str, Any] = {}

        output_dict["queries"] = {qid: query.model_dump() for qid, query in self.queries.items()}
        output_dict["experiment_name"] = self.experiment_name
        output_dict["elo_tournaments"] = [tournament.model_dump() for tournament in self.elo_tournaments]

        with open(output_path, "w") as f:
            json.dump(output_dict, f, indent=4, ensure_ascii=False)

    def save_results(
        self,
        result: EvaluatorResult | EloTournamentResult,
    ):
        """
        Persist the evaluation result to a cache file.
        This method writes the given evaluation result to a cache file specified by
        `self.evaluations_cache_path`. The result is serialized to JSON format and appended
        to the file. The type of the result (either "answer" or "retrieval") is included
        in the serialized data.
        Args:
            result (EvaluatorResult | EloTournamentResult): The evaluation result
                to be persisted to the cache.
        Raises:
            ValueError: If `self.evaluations_cache_path` is not set or if the type of `result`
                is not recognized.
        """

        if not self.cache_evaluations:
            return
        if self.evaluations_cache_path is None:
            raise ValueError("Results cache path not set. Cannot dump result")
        with open(self.evaluations_cache_path, "a+") as f:
            if isinstance(result, AnswerEvaluatorResult):
                result_type = "answer"
            elif isinstance(result, RetrievalEvaluatorResult):
                result_type = "retrieval"
            elif isinstance(result, EloTournamentResult):
                result_type = "elo_tournament"
            else:
                raise ValueError(f"Cannot save evaluation of type {type(result)} to cache")
            f.write(json.dumps({result_type: result.model_dump()}, ensure_ascii=False) + "\n")

    def add_queries_from_csv(
        self,
        file_path: str,
        query_id_column: str | None = None,
        query_text_column: str = "query",
        exist_ok: bool = False,
    ):
        if not os.path.isfile(file_path):
            raise FileNotFoundError(f"CSV file with queries {file_path} not found")
        read_queries = set()
        try:
            query_id_column = self.__infer_query_id_column(file_path)
        except ValueError:
            query_id_column = None

        for idx, row in enumerate(csv.DictReader(open(file_path))):
            if query_id_column is None:
                query_id_column = f"query_{idx}"
            else:
                qid = row.get(query_id_column, f"query_{idx}")
            if qid in read_queries or qid in self.queries:
                if not exist_ok:
                    logger.warning(f"Query with ID {qid} already read. Skipping")
                continue
            query_text = row[query_text_column].strip()
            metadata = {k: v for k, v in row.items() if k not in [query_id_column, query_text_column]}
            self.queries[qid] = Query(qid=qid, query=query_text, metadata=metadata)
            read_queries.add(qid)
        if len(read_queries) > 0:
            logger.info(f"Loaded {len(read_queries)} queries from {file_path}")

    def add_documents_from_csv(
        self,
        file_path: str,
        document_id_col: str | None = "did",
        document_text_col: str = "document",
        query_id_col: str | None = "qid",
        agent_col: str | None = None,
        exist_ok: bool = False,
    ):
        """
        Loads a list of retrieved documents from a csv file. The csv file should have the following columns:
        query_id, document_id, document_text, [agent], [metadata columns]
        Args:
            file_path (str): Path to the CSV file with the retrieved documents.
            document_id_col (str): Name of the column with the document id. Defaults to 'did'.
            query_id_col (str): Name of the column with the query id to which the document was retrieved.
                Defaults to 'qid'.
            document_text_col (str): Name of the column with the document text. Defaults to 'document_text'.
            agent_col (str): Name of the column with the agent name. If None, will not add the agent name.
                Defaults to None.
            exist_ok (bool): If True, will not raise an error if the document already exists.
                Defaults to False.
        """
        documents_read = 0
        if not os.path.isfile(file_path):
            raise FileNotFoundError(f"CSV file {file_path} not found")
        query_id_col = self.__infer_query_id_column(file_path, query_id_col)
        document_id_col = self.__infer_document_id_column(file_path, document_id_col)
        warned_queries = set()

        for line in csv.DictReader(open(file_path)):
            qid = line[query_id_col].strip()
            did = line[document_id_col].strip()
            text = line[document_text_col].strip()
            agent = line.get(agent_col)
            metadata = {
                k: v for k, v in line.items() if k not in [query_id_col, document_id_col, document_text_col, agent_col]
            }
            if qid not in self.queries:
                if qid not in warned_queries:
                    warned_queries.add(qid)
                    logger.warning(f"Query {qid} found in {file_path} but not found in queries. Skipping")
                continue
            if (doc_obj := self.queries[qid].retrieved_docs.get(did)) is None:
                doc_obj = Document(qid=qid, did=did, text=text)
                doc_obj.add_metadata(metadata)
                documents_read += 1
            if agent is not None:
                doc_obj.add_retrieved_by(agent, exist_ok=exist_ok)
            self.add_retrieved_doc(doc_obj, exist_ok=exist_ok)
        if documents_read > 0:
            logger.info(f"Loaded {documents_read} new documents from {file_path}")

    def add_retrieved_docs_from_runfile(
        self,
        run_file_path: str,
        corpus: dict[str, Document] | dict[str, str],
        top_k: int = 1000,
    ):
        """
        Adds the retrieved documents from a run file to the queries.

        A run file has, traditionally, the following format:
            qid, Q0, docid, rank, score, run_name
        The run name will be considered as the agent that retrieved the documents.
            The second column is usually ignored.
        Args:
            run_file_path (str): Path to the run file.
            corpus (dict[str, Document] | dict[str, str]): A dictionary with the documents.
                If a dictionary with just docid and text is provided, we will use it as a corpus.
                Otherwise, we will assume it's a dictionary with docid as key and Document object as value.
            top_k (int): The maximum number of documents to add for each query.
                Assumes that the run file is ordered by rank.
        """
        documents_read = set()
        missing_docs = set()
        docs_per_query = {}
        warned_queries = set()

        if isinstance(list(corpus.values())[0], str):
            corpus = {k: str(v) for k, v in corpus.items()}
            # corpus = {k: Document(did=k, text=str(v)) for k, v in corpus.items()}

        if not os.path.isfile(run_file_path):
            raise FileNotFoundError(f"Run file {run_file_path} not found")
        for line in open(run_file_path):
            qid, _, did, _, score, agent = line.strip().split()
            if did not in corpus:
                missing_docs.add(did)
                continue
            doc = corpus[did]
            if isinstance(doc, str):
                doc = Document(qid=qid, did=did, text=doc)
            if qid not in self.queries:
                if qid not in warned_queries:
                    warned_queries.add(qid)
                    logger.warning(f"Query {qid} not found in queries. Skipping")
                continue
            if qid not in docs_per_query:
                docs_per_query[qid] = 0
            if docs_per_query[qid] < top_k:
                self.add_retrieved_doc(doc, qid, score=float(score), agent=agent)
                docs_per_query[qid] += 1
                documents_read.add(did)
        if len(missing_docs) > 0:
            logger.warning(f"Loaded {len(documents_read)} documents. {len(missing_docs)} missing docs")

    def add_agent_answers_from_csv(
        self,
        file_path: str,
        agent_col: str = "agent",
        answer_col: str = "answer",
        query_id_col: str | None = None,
        exist_ok: bool = False,
    ):
        """
        Adds agent answers from a CSV file to the queries.
        The CSV file should have the following columns:
        query_id, agent, answer, [metadata columns]
        Args:
            csv_path (str): Path to the CSV file with the agent answers.
            agent_col (str): Name of the column with the agent name. Defaults to 'agent'.
            answer_col (str): Name of the column with the answer text. Defaults to 'answer'.
            query_id_col (str): Name of the column with the query id.
                If None, will try to infer the query id column. Defaults to None.
            exist_ok (bool): If True, will not raise an error if the answer already exists.
                Defaults to False.
        """
        if not os.path.isfile(file_path):
            raise FileNotFoundError(f"Answers file {file_path} not found")
        answers_read = 0
        warned_queries = set()
        query_id_col = self.__infer_query_id_column(file_path, query_id_col)
        for line in csv.DictReader(open(file_path)):
            qid = line[query_id_col].strip()
            agent = line[agent_col].strip()
            answer = line[answer_col].strip()
            metadata = {k: v for k, v in line.items() if k not in [query_id_col, agent_col, answer_col]}
            if qid not in self.queries:
                if qid not in warned_queries:
                    warned_queries.add(qid)
                    logger.warning(f"Query {qid} found in {file_path} but not found in queries. Skipping")
                continue
            if agent not in self.queries[qid].answers:
                answer_obj = AgentAnswer(qid=qid, agent=agent, text=answer)
                answer_obj.add_metadata(metadata)
                self.add_agent_answer(answer_obj, exist_ok=exist_ok)
                answers_read += 1
            elif not exist_ok:
                logger.info(f"Answer from agent {agent} already exists in query {qid}. Skipping")
        if answers_read > 0:
            logger.info(f"Loaded {answers_read} answers from {file_path}")

    def add_answers_from_multiple_csv(
        self,
        answers_files: list[str],
        query_id_col: str | None = None,
        query_text_col: str = "query",
        answer_text_col: str = "answer",
    ):
        """Loads queries and answers from a list of CSVs with queries and answers generated by agents.
        We assume that each csv was generated by a single agent, and name each agent according to the file name.
        Args:
            answers_files (list[str]): List of paths to the CSV files with the answers.
            query_id_col (str): The column name with the query ID.
                If None, will try to infer the query ID column. Defaults to None.
            query_text_col (str): The column name with the query text. Defaults to 'query'.
            answer_text_col (str): The column name with the answer text. Defaults to 'answer'.
        """
        for f in answers_files:
            if query_id_col is None:
                query_id_col = self.__infer_query_id_column(f)
            if query_id_col is None:
                query_id_col = query_text_col
            p = Path(f)
            agent_name = p.stem

            for line in csv.DictReader(open(p)):
                qid = line[query_id_col].strip()
                query = line[query_text_col].strip()
                answer = line[answer_text_col].strip()
                if qid not in self.queries:
                    self.add_query(query, qid)
                self.add_agent_answer(answer, agent_name, qid)

    def __clear_all_evaluations(self, should_save: bool = False):
        logger.warning(f"Clearing all evaluations for {len(self)} queries")
        doc_eval_count = 0
        game_eval_count = 0
        answer_eval_count = 0
        for query in self.queries.values():
            for doc in query.retrieved_docs.values():
                if doc.evaluation is not None:
                    doc_eval_count += 1
                doc.evaluation = None
            for game in query.pairwise_games:
                if game.evaluation is not None:
                    game_eval_count += 1
                game.evaluation = None
            for answer in query.answers.values():
                if answer.evaluation is not None:
                    answer_eval_count += 1
                answer.evaluation = None
        self.elo_tournaments = []
        if doc_eval_count > 0 or game_eval_count > 0 or answer_eval_count > 0:
            logger.info(
                f"Cleared {doc_eval_count} document evaluations, {game_eval_count} game evaluations, "
                f"and {answer_eval_count} answer evaluations"
            )

        if should_save:
            if self.evaluations_cache_path and os.path.isfile(self.evaluations_cache_path):
                with open(self.evaluations_cache_path, "w"):
                    pass

            self.save()

    def _load_from_cache(self, cache_path: str) -> dict[str, Query]:
        if not os.path.isfile(cache_path):
            raise FileNotFoundError(f"Cache file {cache_path} not found")
        logger.info(f"Loading existing experiment from {cache_path}")
        with open(cache_path) as f:
            try:
                data = json.load(f)
            except json.JSONDecodeError:
                raise ValueError(f"Cache file {cache_path} is not a valid JSON file")
        read_experiment_name = data.get("experiment_name")
        if read_experiment_name and read_experiment_name != self.experiment_name:
            logger.warning(f"Experiment name mismatch. Expected {self.experiment_name}. Found {read_experiment_name}")
        queries = {k: Query(**v) for k, v in data.get("queries").items()}
        for q in queries.values():
            for doc in q.retrieved_docs.values():
                if doc.evaluation is not None and not isinstance(doc.evaluation, RetrievalEvaluatorResult):
                    doc_eval_dict = doc.evaluation.model_dump()
                    doc_eval_dict["did"] = doc.did
                    doc.evaluation = RetrievalEvaluatorResult(**doc_eval_dict)
            for answer in q.answers.values():
                if answer.evaluation is not None and not isinstance(answer.evaluation, AnswerEvaluatorResult):
                    answer_eval_dict = answer.evaluation.model_dump()
                    answer.evaluation = AnswerEvaluatorResult(**answer_eval_dict)
        if "elo_tournaments" in data:
            self.elo_tournaments = [EloTournamentResult(**t) for t in data["elo_tournaments"]]
        return queries

    def _load_results_from_cache(self, cache_path: str):
        assert os.path.isfile(cache_path)
        for line in open(cache_path):
            result = json.loads(line)
            result_type = list(result.keys())[0]
            if result_type == "answer":
                result = AnswerEvaluatorResult(**result["answer"])
                if result.qid not in self.queries:
                    continue
            elif result_type == "retrieval":
                result = RetrievalEvaluatorResult(**result["retrieval"])
                if result.qid not in self.queries:
                    continue
            elif result_type == "elo_tournament":
                result = EloTournamentResult(**result["elo_tournament"])
            self.add_evaluation(result, should_save=False, should_print=False, exist_ok=True)

    def __len__(self):
        return len(self.queries)

    def __iter__(self):
        return iter(self.queries.values())

    def __getitem__(self, key: str) -> Query:
        return self.queries[key]

    def keys(self):
        return self.queries.keys()

    @staticmethod
    def __infer_query_id_column(file_path: str, query_id_col: str | None = None) -> str:
        """Infer the column name with the query id from a CSV file."""
        if query_id_col is not None:
            return query_id_col
        possible_qid_columns = [
            "query_id",
            "qid",
            "question_id",
            "q_id",
            "queryid",
            "questionid",
        ]
        with open(file_path) as f:
            reader = csv.DictReader(f)
            columns = reader.fieldnames
            if columns is not None:
                for col in possible_qid_columns:
                    if col in columns:
                        return col
        raise ValueError(f"Could not identify Query ID column for CSV file {file_path}")

    @staticmethod
    def __infer_document_id_column(file_path: str, doc_id_col: str | None = None) -> str:
        """Infer the column name with the doc id from a CSV file."""
        if doc_id_col is not None:
            return doc_id_col
        possible_did_columns = [
            "document_id",
            "did",
            "d_id",
            "documentid",
            "passage_id",
            "pid",
            "p_id",
            "passageid",
        ]
        with open(file_path) as f:
            reader = csv.DictReader(f)
            columns = reader.fieldnames
            if columns is not None:
                for col in possible_did_columns:
                    if col in columns:
                        return col
        raise ValueError(f"Could not identify Document ID column for CSV file {file_path}")

    @classmethod
    def load(cls, experiment_name: str, path: str):
        return cls(experiment_name=experiment_name, save_path=path)
