import hashlib
import typing as T
from abc import ABC, abstractmethod
from functools import cached_property
from pathlib import Path

import datasets
import pandas as pd
from datasets import Features, Sequence, Value
from llama_index.core import Document, SimpleDirectoryReader
from llama_index.core.evaluation.correctness import DEFAULT_USER_TEMPLATE
from overrides import overrides
from pydantic import BaseModel, field_validator
from ray._private.utils import get_ray_temp_dir

from syftr.configuration import cfg
from syftr.core import QAPair, QARTriplet
from syftr.utils.locks import distributed_lock


class PartitionMap(BaseModel):
    sample: str = "sample"
    train: str = "train"
    test: str = "test"
    holdout: str = "holdout"


class SyftrQADataset(BaseModel, ABC):
    """Container and utilities for dataset with remote storage.

    Instances of this class do _not_ store the dataset; only pointers to it.
    Therefore, this class can be safely passed to Ray tune.
    """

    # Canonical name for the dataset
    # Set this to a unique string Literal on your subclass
    # Must be present when loading a StudyConfig from a yaml
    xname: T.Literal["SyftrQADataset"] = "SyftrQADataset"

    # Partition names for this dataset as it is stored on disk
    storage_partitions: T.List[str] = ["sample", "train", "test", "holdout"]
    # How to map requested partition to the storage partitions
    # eg. MyDataset(partition_map={'test': 'sample'}) to run on the sample partition
    partition_map: PartitionMap = PartitionMap()

    # timeouts
    load_examples_timeout_s: int = 3600
    load_grounding_data_timeout_s: int = 3600

    # The evaluation user template specifies content and format of the information that is passed to the evaluator.
    # The template must include Python string templates for the following fields:
    #     'query', 'reference_answer', and 'generated_answer'.
    # The user may use this template to insert task-specific instructions to the evaluator to ensure the appropriate
    # criteria are being used.
    eval_user_template: str = DEFAULT_USER_TEMPLATE

    @property
    def name(self) -> str:
        """Subclasses may dynamically construct name."""
        return self.xname

    def _get_storage_partition(self, canonical_partition: str) -> str:
        return getattr(self.partition_map, canonical_partition)

    @abstractmethod
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        pass

    @abstractmethod
    def iter_grounding_data(self, partition="test") -> T.Iterator[Document]:
        pass


class InfiniteBenchHF(SyftrQADataset):
    xname: T.Literal["infinitebench_hf"] = "infinitebench_hf"  # type: ignore
    subset: T.Literal["longbook_qa_eng"] = "longbook_qa_eng"
    description: str = "The dataset contains a large number of books."

    def _load_raw_dataset(self) -> pd.DataFrame:
        ft = Features(
            {
                "id": Value("int64"),
                "context": Value("string"),
                "input": Value("string"),
                "answer": Sequence(Value("string")),
                "options": Sequence(Value("string")),
            }
        )

        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "xinrongzhang2022/InfiniteBench",
                features=ft,
                split=self.subset,
                cache_dir=cfg.paths.huggingface_cache,
            )
        df = dataset.to_pandas()
        return df

    def _add_partitions(self, df: pd.DataFrame) -> pd.DataFrame:
        df["book_id"] = df.context.factorize()[0]
        books = df.book_id.unique()
        book_partitions = {
            "sample": books[:1],
            "train": books[1:23],
            "test": books[23:46],
            "holdout": books[46:69],
        }

        def label_partition(book_id: int) -> str:
            for partition, book_range in book_partitions.items():
                if book_id in book_range:
                    return partition
            raise IndexError(f"Book id {book_id} is out of range of {book_partitions=}")

        df["partition"] = df.book_id.apply(label_partition)
        return df

    @property
    def _dataset(self) -> pd.DataFrame:
        df = self._load_raw_dataset()
        df = self._add_partitions(df)
        return df

    def _row_to_qapair(self, row: T.Dict[str, T.Any]) -> QAPair:
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["input"],
            answer=str(row["answer"]),
            id=str(row["id"]),
            context={"book_start": row["context"][:100]},
            supporting_facts=[],
            difficulty="",
            qtype="",
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        df = self._dataset
        partition = self._get_storage_partition(partition)
        df = df[df.partition == partition]
        for _, row in df.iterrows():
            yield self._row_to_qapair(row)

    @overrides
    def iter_grounding_data(self, partition="test") -> T.Iterator[Document]:
        df = self._dataset
        partition = self._get_storage_partition(partition)
        df = df[df.partition == partition]
        for book in df.context.unique():
            yield Document(
                text=book,
            )


class FinanceBenchHF(SyftrQADataset):
    xname: T.Literal["financebench_hf"] = "financebench_hf"  # type: ignore
    description: str = (
        "Financial dataset that contains everything about finance, including real-world financial documents, "
        "SEC filings, earning reports, call transcripts, and much more. "
        "It has all the financial live data, historical data, just about everything about finance, for instance, "
        "definitions and explanations of financial term, "
        "insights on company revenues, mergers, founders, or stock performance, "
        "details on financial laws, compliance, or government policies, "
        "information required to evaluated finance risk, and "
        "information about banking operations, credit systems, or loan structures."
    )

    def _load_grounding_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/financebench",
                "groundingdata",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/financebench",
                "qapairs",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    @overrides
    def iter_grounding_data(
        self, partition="test", **load_kwargs
    ) -> T.Iterator[Document]:
        assert partition in self.storage_partitions
        grounding_dataset = self._load_grounding_dataset()
        partition = self._get_storage_partition(partition)
        for row in grounding_dataset[partition]:
            yield Document(
                text=row["html"],
                metadata={"file_name": row["filename"]},
            )

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["question"],
            answer=row["answer"],
            id=str(row["financebench_id"]),
            context=row["evidence"],
            supporting_facts=[row["justification"]],
            difficulty="",
            qtype=row["question_type"],
            gold_evidence=row["html_evidence"],
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        partition = self._get_storage_partition(partition)
        qa_examples = self._load_qa_dataset()
        for row in qa_examples[partition]:
            yield self._row_to_qapair(row)


class SyntheticFinanceBenchHF(FinanceBenchHF):
    xname: T.Literal["synthetic_financebench_hf"] = "synthetic_financebench_hf"  # type: ignore

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/financebench",
                "qapairs_synthetic",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["query"],
            answer=row["reference_answer"],
            id=str(row["id"]),
            context={},
            supporting_facts=[],
            difficulty="default",
            qtype="default",
        )


class HotPotQAHF(SyftrQADataset):
    xname: T.Literal["hotpotqa_hf"] = "hotpotqa_hf"  # type: ignore
    subset: str = "dev"  # train_hard, dev
    description: str = (
        "This dataset is a vast collection of all kind of information that you can find on Wikipedia. "
        "It can be used, for instance, to retrieve straightforward facts from one or more documents, "
        "compare two entities based on shared attributes, "
        "identify relationships, roles, or attributes of entities, "
        "reason about dates, timelines, or chronological order, "
        "determine geographical relationships or locations, "
        "explain causes or sequences of events or processes, "
        "synthesize facts from multiple documents to infer answers, and "
        "validate or refute premises in the context of the question."
    )

    @property
    def name(self) -> str:
        return f"{self.xname}/{self.subset}"

    def _load_grounding_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/hotpotqa",
                f"groundingdata_{self.subset}",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/hotpotqa",
                f"qapairs_{self.subset}",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    @overrides
    def iter_grounding_data(
        self,
        partition="test",
    ) -> T.Iterator[Document]:
        assert partition in self.storage_partitions
        grounding_dataset = self._load_grounding_dataset()
        partition = self._get_storage_partition(partition)
        for row in grounding_dataset[partition]:
            yield Document(
                text=row["text"],
            )

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["question"],
            answer=row["answer"],
            id=str(row["id"]),
            context=[{title: sentence} for title, sentence in row["context"]],
            supporting_facts=row["supporting_facts"],
            difficulty=row["level"],
            qtype=row["type"],
            gold_evidence=[sentence for _, sentence in row["context"]],
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        partition = self._get_storage_partition(partition)
        qa_examples = self._load_qa_dataset()
        for row in qa_examples[partition]:
            yield self._row_to_qapair(row)


class SyntheticHotPotQAHF(HotPotQAHF):
    xname: T.Literal["synthetic_hotpotqa_hf"] = "synthetic_hotpotqa_hf"  # type: ignore

    @property
    def name(self) -> str:
        return f"{self.xname}/{self.subset}"

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/hotpotqa",
                f"qapairs_synthetic_{self.subset}",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["query"],
            answer=row["reference_answer"],
            id=str(row["id"]),
            context={},
            supporting_facts=[],
            difficulty="default",
            qtype="default",
        )


class CragTask3HF(SyftrQADataset):
    xname: T.Literal["crag_hf"] = "crag_hf"  # type: ignore
    subset: str = "sports"  # finance, movie, music, sports, open

    _descriptions = {
        "sports": (
            "This resource contains everything about sports, including sports news, sports events, "
            "sports statistics, sports schedules, and more. "
            "It contains historical data, live data, just about everything about sports, for instance, "
            "information about athlete achievements, teams, or career stats, "
            "information on match dates, scores, or tournaments, "
            "team performances or player statistics, "
            "key events or records in sports history, and "
            "information about rules or formats in specific sports."
        ),
        "finance": (
            "This resource contains everything about finance, including financial news, "
            "market data, financial reports, and more. "
            "It contains historical data, live data, just about everything about finance, for instance, "
            "stock prices, trends, or market indices, "
            "revenue, founders, or headquarters of companies, "
            "major financial events or policy changes, "
            "comparing returns or risks of different investments, and "
            "financial terms or concepts explained."
        ),
        "movie": (
            "This resource contains everything about movies, including movie news, reviews, "
            "box office data, and more. "
            "It contains historical data, live data, just about everything about movies, for instance, "
            "information about actors, directors, or crew roles, "
            "movies by a specific actor or director, "
            "key events or summaries from movie plots, "
            "recognition received by films or individuals, and "
            "dates or production details for movies."
        ),
        "music": (
            "This resource contains everything about music, including music news, charts, "
            "artist information, and more. "
            "It contains historical data, live data, just about everything about music, for instance, "
            "information on musicians, albums, or band members, "
            "albums or songs released by specific artists, "
            "information about song lyrics or their significance, "
            "Grammys or other recognitions for artists or albums, and "
            "details about music styles or trends over time."
        ),
        "open": (
            "This resource contains everything about open domain, including general knowledge, "
            "trivia, fun facts, and more. "
            "It contains historical data, live data, just about everything about open domain, for instance, "
            "general knowledge about diverse topics, "
            "facts that combine multiple subject areas, "
            "unique or interesting facts across domains (trivia), and "
            "further material useful for complex, unconstrained reasoning without a specific domain."
        ),
    }

    @property
    def name(self) -> str:
        return f"{self.xname}/{self.subset}"

    @property
    def description(self) -> str:
        return self._descriptions[self.subset]

    def _load_grounding_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/crag",
                f"groundingdata_{self.subset}",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/crag",
                f"qapairs_{self.subset}",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    @overrides
    def iter_grounding_data(self, partition="test") -> T.Iterator[Document]:
        assert partition in self.storage_partitions
        grounding_dataset = self._load_grounding_dataset()
        partition = self._get_storage_partition(partition)
        for row in grounding_dataset[partition]:
            yield Document(
                text=row["markdown"],
                metadata={"file_name": row["filename"]},
            )

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        context = {result["_hash"]: str(result) for result in row["search_results"]}
        return QAPair(
            question=row["query"],
            answer=row["answer"],
            id=str(row["interaction_id"]),
            context=context,
            supporting_facts=[],
            difficulty="default",
            qtype=row["question_type"],
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        partition = self._get_storage_partition(partition)
        qa_examples = self._load_qa_dataset()
        for row in qa_examples[partition]:
            yield self._row_to_qapair(row)


class SyntheticCragTask3HF(CragTask3HF):
    xname: T.Literal["synthetic_crag_hf"] = "synthetic_crag_hf"  # type: ignore

    @property
    def name(self) -> str:
        return f"{self.xname}/{self.subset}"

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/crag",
                f"qapairs_synthetic_{self.subset}",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["query"],
            answer=row["reference_answer"],
            id=str(row["id"]),
            context={},
            supporting_facts=[],
            difficulty="default",
            qtype="default",
        )


class DRDocsHF(SyftrQADataset):
    xname: T.Literal["drdocs_hf"] = "drdocs_hf"  # type: ignore
    description: str = (
        "The dataset contains comprehensive information about DataRobot, including its API, documentation, examples, "
        "key features, platform architecture, integrations, setup guides, data handling, feature engineering, EDA tools, "
        "automated machine learning, model management, deployment options, monitoring, REST API, batch predictions, "
        "real-time scoring, custom recipes, retraining, lifecycle management, bias detection, explainability, diagnostics, "
        "cross-validation, leaderboard insights, time series modeling, data governance, security, user roles, Python/R usage, "
        "custom blueprints, external model integration, Docker deployments, API reference, BI tool integration, workflow automation, "
        "multimodal modeling, NLP, image recognition, hyperparameter tuning, performance optimization, resource management, "
        "parallel processing, drift detection, retraining triggers, industry use cases, tutorials, case studies, common issues, "
        "debugging tips, FAQs, support access, community resources, and release notes."
    )

    def _load_grounding_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/drdocs",
                "groundingdata",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/drdocs",
                "qapairs",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset

    @overrides
    def iter_grounding_data(self, partition="notused") -> T.Iterator[Document]:
        # There is no partition. The grounding dataset is the same
        # across all partitions of the qa pairs.
        grounding_dataset = self._load_grounding_dataset()
        for row in grounding_dataset["train"]:
            yield Document(
                text=row["markdown"],
                metadata={"file_name": row["filename"]},
            )

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["question"],
            answer=row["answer"],
            id=str(row["id"]),
            context={},
            supporting_facts=[],
            difficulty="default",
            qtype="default",
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        assert partition in self.storage_partitions
        partition = self._get_storage_partition(partition)
        qa_examples = self._load_qa_dataset()
        for row in qa_examples[partition]:
            yield self._row_to_qapair(row)


class PhantomWikiv050(SyftrQADataset):
    xname: T.Literal["phantomwikiv050_hf"] = "phantomwikiv050_hf"  # type: ignore

    # subset choices are:
    # 'depth_20_size_25_seed_1',
    # 'depth_20_size_25_seed_2',
    # 'depth_20_size_25_seed_3',
    # 'depth_20_size_50_seed_1',
    # 'depth_20_size_50_seed_2',
    # 'depth_20_size_50_seed_3',
    # 'depth_20_size_100_seed_1',
    # 'depth_20_size_100_seed_2',
    # 'depth_20_size_100_seed_3',
    # 'depth_20_size_200_seed_1',
    # 'depth_20_size_200_seed_2',
    # 'depth_20_size_200_seed_3',
    # 'depth_20_size_300_seed_1',
    # 'depth_20_size_300_seed_2',
    # 'depth_20_size_300_seed_3',
    # 'depth_20_size_400_seed_1',
    # 'depth_20_size_400_seed_2',
    # 'depth_20_size_400_seed_3',
    # 'depth_20_size_500_seed_1',
    # 'depth_20_size_500_seed_2',
    # 'depth_20_size_500_seed_3',
    # 'depth_20_size_1000_seed_1',
    # 'depth_20_size_1000_seed_2',
    # 'depth_20_size_1000_seed_3',
    # 'depth_20_size_2500_seed_1',
    # 'depth_20_size_2500_seed_2',
    # 'depth_20_size_2500_seed_3',
    # 'depth_20_size_5000_seed_1',
    # 'depth_20_size_5000_seed_2',
    # 'depth_20_size_5000_seed_3',
    # 'depth_20_size_10000_seed_1',
    # 'depth_20_size_10000_seed_2',
    # 'depth_20_size_10000_seed_3'

    subset: str = "depth_20_size_10000_seed_3"
    description: str = (
        "This dataset contains data from PhantomWiki, which "
        "is a framework for generating unique, factually"
        "and consistent document corpora with diverse question-answer pairs."
        "Unlike prior work, PhantomWiki is neither a fixed dataset, nor is it"
        "based on any existing data. Instead, a new PhantomWiki instance is "
        "generated on demand for each evaluation. PhantomWiki generates a "
        "fictional universe of characters along with a set of facts. "
        "We reflect these facts in a large-scale corpus, mimicking the "
        "style of fan-wiki websites. Then we generate question-answer pairs "
        "with tunable difficulties, encapsulating the types of multi-hop "
        "questions commonly considered in the question-answering (QA) literature."
    )

    def _get_partition_range(self, partition: str):
        partition_ranges = {
            "sample": range(0, 10),
            "train": range(0, 100),
            "test": range(100, 400),
            "holdout": range(400, 500),
        }
        if partition in partition_ranges:
            return partition_ranges[partition]

    def _load_grounding_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "kilian-group/phantom-wiki-v050",
                "text-corpus",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset[self.subset]

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "kilian-group/phantom-wiki-v050",
                "question-answer",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset[self.subset]

    @overrides
    def iter_grounding_data(self, partition="notused") -> T.Iterator[Document]:
        # There is no partition. The grounding dataset is the same
        # across all partitions of the qa pairs.
        grounding_dataset = self._load_grounding_dataset()
        for row in grounding_dataset:
            yield Document(
                text=row["article"],
                metadata={"title": row["title"]},
            )

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["question"],
            answer=" ".join(row["answer"]),
            id=str(row["id"]),
            context={},
            supporting_facts=[],
            difficulty=str(row["difficulty"]),
            qtype=str(row["type"]),
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        assert partition in self.storage_partitions
        partition = self._get_storage_partition(partition)
        qa_examples = self._load_qa_dataset()
        partition_range = self._get_partition_range(partition)

        for i in partition_range:
            row = qa_examples[i]
            yield self._row_to_qapair(row)


class MultiHopRAGHF(SyftrQADataset):
    xname: T.Literal["multihoprag_hf"] = "multihoprag_hf"  # type: ignore

    description: str = (
        "This resource contains a corpus of news dataset. "
        "The news data source comprises various English-language "
        "websites covering a range of news categories: "
        "entertainment, business, sports, technology, "
        "health and science. It contains news articles "
        "on these topics published between September to December 2023. "
        "Every news article is paired with metadata, "
        "including title, publish date, author, category, URL and news source."
    )

    def _get_partition_range(self, partition: str):
        partition_ranges = {
            "sample": range(0, 20),
            "train": range(0, 1000),
            "test": range(1000, 2000),
            "holdout": range(2000, 2556),
        }
        if partition in partition_ranges:
            return partition_ranges[partition]

    def _load_grounding_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "yixuantt/MultiHopRAG",
                "corpus",
                cache_dir=cfg.paths.huggingface_cache,
            )
        # all of the corpus is in a split named "train"
        return dataset["train"]

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "yixuantt/MultiHopRAG",
                "MultiHopRAG",
                cache_dir=cfg.paths.huggingface_cache,
            )

        # all qapairs are in a split named "train"
        return dataset["train"]

    @overrides
    def iter_grounding_data(self, partition="notused") -> T.Iterator[Document]:
        # There is no partition. The grounding dataset is the same
        # across all partitions of the qa pairs.
        grounding_dataset = self._load_grounding_dataset()
        for row in grounding_dataset:
            yield Document(
                text=row["body"],
                metadata={
                    "title": row["title"],
                    "author": row["author"],
                    "category": row["category"],
                    "published_at": row["published_at"],
                    "url": row["url"],
                    "source": row["source"],
                },
            )

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["query"],
            answer=row["answer"],
            id=hashlib.md5(row["query"].encode()).hexdigest(),
            context={},
            supporting_facts=row["evidence_list"],
            difficulty="",
            qtype=str(row["question_type"]),
            gold_evidence=[item["fact"] for item in row["evidence_list"]],
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        assert partition in self.storage_partitions
        partition = self._get_storage_partition(partition)
        qa_examples = self._load_qa_dataset()
        partition_range = self._get_partition_range(partition)

        for i in partition_range:
            row = qa_examples[i]
            yield self._row_to_qapair(row)


class BrightHF(SyftrQADataset):
    xname: T.Literal["bright_hf"] = "bright_hf"  # type: ignore

    # NOTE: pony, leetcode, aops, theoremqa_theorems, theoremqa_questions
    # subsets do not have ground truth answers
    subset: T.Literal[
        "earth_science",
        "biology",
        "economics",
        "psychology",
        "robotics",
        "stackoverflow",
        "sustainable_living",
        "pony",
        "leetcode",
        "aops",
        "theoremqa_theorems",
        "theoremqa_questions",
    ] = "biology"

    _descriptions = {
        "earth_science": (
            "This resource contains everything about earth science, including geology, "
            "oceanography, meteorology, and more. "
            "It contains historical data, live data, just about everything about earth science, for instance, "
            "information about geological formations, ocean currents, or weather patterns, "
            "key events or phenomena in earth science history, "
            "information about climate change or environmental issues, and "
            "information about earth science research methods or techniques."
        ),
        "biology": (
            "This resource contains everything about biology, including genetics, "
            "ecology, evolution, and more. "
            "It contains historical data, live data, just about everything about biology, for instance, "
            "information about cellular structures, ecosystems, or evolutionary processes, "
            "key events or discoveries in biology history, "
            "information about genetic engineering or biotechnology, and "
            "information about biological research methods or techniques."
        ),
        "economics": (
            "This resource contains everything about economics, including microeconomics, "
            "macroeconomics, international trade, and more. "
            "It contains historical data, live data, just about everything about economics, for instance, "
            "information about supply and demand, market structures, or economic indicators, "
            "key events or theories in economics history, "
            "information about economic policies or regulations, and "
            "information about economic research methods or techniques."
        ),
        "psychology": (
            "This resource contains everything about psychology, including cognitive, "
            "developmental, social, and clinical psychology. "
            "It contains historical data, live data, just about everything about psychology, for instance, "
            "information about brain structures, psychological disorders, or therapeutic techniques, "
            "key events or theories in psychology history, "
            "information about psychological research methods or techniques, and "
            "information about psychological assessments or tests."
        ),
        "robotics": (
            "This resource contains everything about robotics, including robot design, "
            "control systems, and applications. "
            "It contains historical data, live data, just about everything about robotics, for instance, "
            "information about robot components, sensors, or actuators, "
            "key events or advancements in robotics history, "
            "information about robotic applications in various fields, and "
            "information about robotics research methods or techniques."
        ),
        "stackoverflow": (
            "This resource contains everything about StackOverflow, including programming, "
            "software development, and more. "
            "It contains historical data, live data, just about everything about StackOverflow, for instance, "
            "information about programming languages, software frameworks, or development methodologies, "
            "key events or trends in software development history, "
            "information about software development best practices or tools, and "
            "information about software development research methods or techniques."
        ),
        "sustainable_living": (
            "This resource contains everything about sustainable living, including "
            "environmental conservation, renewable energy, and more. "
            "It contains historical data, live data, just about everything about sustainable living, for instance, "
            "information about sustainable practices, technologies, or policies, "
            "key events or movements in sustainability history, "
            "information about environmental issues or challenges, and "
            "information about sustainability research methods or techniques."
        ),
        "pony": (
            "This resource contains everything about the Pony programming language, including its syntax, "
            "capabilities, compiler architecture, and more. "
            "It contains documentation, examples, just about everything about Pony, for instance, "
            "information about Pony's actor-based concurrency model, capabilities-secure type system, "
            "key concepts like reference capabilities and behavior guarantees, "
            "information about pattern matching, generics, and other language features, and "
            "information about Pony's memory management, garbage collection, and runtime performance."
        ),
        "leetcode": (
            "This resource contains everything about LeetCode, including its problems, "
            "solutions, and more. "
            "It contains documentation, examples, just about everything about LeetCode, for instance, "
            "information about data structures, algorithms, or coding challenges, "
            "key concepts like time complexity and space complexity, "
            "information about problem-solving techniques and strategies, and "
            "information about LeetCode's online judge system and coding environment."
        ),
        "aops": (
            "This resource contains everything about Art of Problem Solving (AoPS), including its curriculum, "
            "problems, and more. "
            "It contains documentation, examples, just about everything about AoPS, for instance, "
            "information about problem-solving techniques, mathematical concepts, or competition preparation, "
            "key events or competitions in mathematics history, "
            "information about AoPS's online community and resources, and "
            "information about AoPS's books and courses."
        ),
        "theoremqa_theorems": (
            "This resource contains everything about TheoremQA, including its theorems, "
            "proofs, and more. "
            "It contains documentation, examples, just about everything about TheoremQA, for instance, "
            "information about mathematical theorems, proofs, or concepts, "
            "key events or advancements in theorem proving history, "
            "information about theorem proving techniques and strategies, and "
            "information about TheoremQA's online community and resources."
        ),
        "theoremqa_questions": (
            "This resource contains everything about TheoremQA, including its questions, "
            "answers, and more. "
            "It contains documentation, examples, just about everything about TheoremQA, for instance, "
            "information about mathematical questions, answers, or concepts, "
            "key events or advancements in theorem proving history, "
            "information about theorem proving techniques and strategies, and "
            "information about TheoremQA's online community and resources."
        ),
    }

    @property
    def description(self) -> str:
        return self._descriptions[self.subset]

    def _get_partition_range(self, partition: str):
        subset_partition_ranges = {
            "earth_science": {
                "sample": range(0, 10),
                "train": range(0, 30),
                "test": range(30, 90),
                "holdout": range(90, 116),
            },
            "biology": {
                "sample": range(0, 10),
                "train": range(0, 30),
                "test": range(30, 90),
                "holdout": range(90, 103),
            },
            "economics": {
                "sample": range(0, 10),
                "train": range(0, 30),
                "test": range(30, 90),
                "holdout": range(90, 103),
            },
            "psychology": {
                "sample": range(0, 10),
                "train": range(0, 30),
                "test": range(30, 90),
                "holdout": range(90, 101),
            },
            "robotics": {
                "sample": range(0, 10),
                "train": range(0, 30),
                "test": range(30, 90),
                "holdout": range(90, 101),
            },
            "stackoverflow": {
                "sample": range(0, 10),
                "train": range(0, 30),
                "test": range(30, 90),
                "holdout": range(90, 117),
            },
            "sustainable_living": {
                "sample": range(0, 10),
                "train": range(0, 30),
                "test": range(30, 90),
                "holdout": range(90, 108),
            },
            "pony": {
                "sample": range(0, 10),
                "train": range(0, 30),
                "test": range(30, 90),
                "holdout": range(90, 112),
            },
            "leetcode": {
                "sample": range(0, 10),
                "train": range(0, 30),
                "test": range(30, 120),
                "holdout": range(120, 142),
            },
            "aops": {
                "sample": range(0, 10),
                "train": range(0, 30),
                "test": range(30, 90),
                "holdout": range(90, 111),
            },
            "theoremqa_theorems": {
                "sample": range(0, 10),
                "train": range(0, 20),
                "test": range(20, 50),
                "holdout": range(50, 76),
            },
            "theoremqa_questions": {
                "sample": range(0, 10),
                "train": range(0, 40),
                "test": range(40, 160),
                "holdout": range(160, 194),
            },
        }

        partition_ranges = subset_partition_ranges[self.subset]
        if partition in partition_ranges:
            return partition_ranges[partition]

    def _load_grounding_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "xlangai/BRIGHT",
                "documents",
                cache_dir=cfg.paths.huggingface_cache,
            )
        # all of the corpus is in a split named "train"
        return dataset[self.subset]

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "xlangai/BRIGHT",
                "examples",
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset[self.subset]

    @overrides
    def iter_grounding_data(self, partition="notused") -> T.Iterator[Document]:
        # There is no partition. The grounding dataset is the same
        # across all partitions of the qa pairs.
        grounding_dataset = self._load_grounding_dataset()
        for row in grounding_dataset:
            yield Document(
                text=row["content"],
                metadata={
                    "id": row["id"],
                },
            )

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["query"],
            answer=row["gold_answer"],
            id=row["id"],
            context={},
            supporting_facts=[],
            difficulty="",
            qtype="",
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        assert partition in self.storage_partitions
        partition = self._get_storage_partition(partition)
        qa_examples = self._load_qa_dataset()
        partition_range = self._get_partition_range(partition)

        for i in partition_range:
            row = qa_examples[i]
            yield self._row_to_qapair(row)


class PhantomWikiV001HF(SyftrQADataset):
    xname: T.Literal["phantomwikiv001_hf"] = "phantomwikiv001_hf"  # type: ignore

    # subset choices are:
    # depth_20_size_100_seed_595    depth_20_size_2000_seed_595   depth_20_size_400_seed_595
    # depth_20_size_1000_seed_595   depth_20_size_250_seed_595    depth_20_size_4000_seed_595
    # depth_20_size_10000_seed_595  depth_20_size_2500_seed_595   depth_20_size_450_seed_595
    # depth_20_size_100000_seed_595 depth_20_size_300_seed_595    depth_20_size_4500_seed_595
    # depth_20_size_150_seed_595    depth_20_size_3000_seed_595   depth_20_size_50_seed_595
    # depth_20_size_1500_seed_595   depth_20_size_350_seed_595    depth_20_size_500_seed_595
    # depth_20_size_200_seed_595    depth_20_size_3500_seed_595   depth_20_size_5000_seed_595

    subset: str = "depth_20_size_1000_seed_595"

    description: str = (
        "This dataset contains data from PhantomWiki, which "
        "is a framework for generating unique, factually"
        "and consistent document corpora with diverse question-answer pairs."
        "Unlike prior work, PhantomWiki is neither a fixed dataset, nor is it"
        "based on any existing data. Instead, a new PhantomWiki instance is "
        "generated on demand for each evaluation. PhantomWiki generates a "
        "fictional universe of characters along with a set of facts. "
        "We reflect these facts in a large-scale corpus, mimicking the "
        "style of fan-wiki websites. Then we generate question-answer pairs "
        "with tunable difficulties, encapsulating the types of multi-hop "
        "questions commonly considered in the question-answering (QA) literature."
    )

    def _get_partition_range(self, partition: str):
        partition_ranges = {
            "sample": range(0, 10),
            "train": range(0, 100),
            "test": range(100, 400),
            "holdout": range(400, 500),
        }
        if partition in partition_ranges:
            return partition_ranges[partition]

    def _load_grounding_dataset(self) -> datasets.Dataset:
        subset_name = self.subset + "_groundingdata"
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/phantomwiki-001",
                subset_name,
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset["train"]

    def _load_qa_dataset(self) -> datasets.Dataset:
        subset_name = self.subset + "_qapairs"
        # Load the dataset with the specified subset name
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/phantomwiki-001",
                subset_name,
                cache_dir=cfg.paths.huggingface_cache,
            )

        return dataset["train"]

    @overrides
    def iter_grounding_data(self, partition="notused") -> T.Iterator[Document]:
        # There is no partition. The grounding dataset is the same
        # across all partitions of the qa pairs.
        grounding_dataset = self._load_grounding_dataset()
        for row in grounding_dataset:
            yield Document(
                text=row["article"],
                metadata={"title": row["title"]},
            )

    def _row_to_qapair(self, row):
        """Dataset-specific conversion of row to QAPair struct.

        Invoked by iter_examples.

        Default implementation assumes row is already in QAPair format.
        """
        return QAPair(
            question=row["question"],
            answer=" ".join(row["answer"]),
            id=str(row["id"]),
            context={},
            supporting_facts=[],
            difficulty=str(row["difficulty"]),
            qtype=str(row["type"]),
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        assert partition in self.storage_partitions
        partition = self._get_storage_partition(partition)
        qa_examples = self._load_qa_dataset()
        partition_range = self._get_partition_range(partition)

        for i in partition_range:
            row = qa_examples[i]
            yield self._row_to_qapair(row)


class JudgeEvalHF(SyftrQADataset):
    xname: T.Literal["judge_eval"] = "judge_eval"  # type: ignore

    description: str = (
        "Contains human labeled Question-Answer-Response triplets "
        "for evaluation of LLM-a-judge flows."
    )

    def _get_partition_range(self, partition: str):
        partition_ranges = {
            "sample": range(0, 20),
            "train": range(0, 300),
            "test": range(300, 700),
            "holdout": range(700, 807),
        }
        if partition in partition_ranges:
            return partition_ranges[partition]

    def _load_grounding_dataset(self) -> datasets.Dataset:
        raise NotImplementedError("judge_eval does not have a grounding dataset")

    def _load_qa_dataset(self) -> datasets.Dataset:
        with distributed_lock(
            self.name, timeout_s=self.load_examples_timeout_s, host_only=True
        ):
            dataset = datasets.load_dataset(
                "DataRobot-Research/judge-eval",
                cache_dir=cfg.paths.huggingface_cache,
            )

        # all qapairs are in a split named "train"
        return dataset["train"]

    @overrides
    def iter_grounding_data(self, partition="notused") -> T.Iterator[Document]:
        raise NotImplementedError("judge_eval does not have a grounding dataset")

    def _row_to_qartriplet(self, row):
        """Dataset-specific conversion of row to QARTriplet."""
        return QARTriplet(
            question=row["question"],
            answer=row["answer"],
            id=str(row["index"]),
            context={},
            supporting_facts=[],
            difficulty="",
            qtype="",
            gold_evidence=[],
            response=row["response"],
            label=bool(int(float(row["Human Passing"]))),
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        assert partition in self.storage_partitions
        partition = self._get_storage_partition(partition)
        qa_examples = self._load_qa_dataset()
        partition_range = self._get_partition_range(partition)

        for i in partition_range:
            row = qa_examples[i]
            yield self._row_to_qartriplet(row)


class CustomDataset(SyftrQADataset):
    """
    Custom dataset class that loads data from local files.

    This class allows you to evaluate a RAG pipeline using local documents and QA pairs
    stored on your machine.

    The qa_csv and grounding_data_path fields must be absolute paths so that they can be
    resolved by both Ray workers and local scripts.

    The CustomDataset class is not intended for use in distributed Ray environments. To
    use this dataset in Ray, you are responsible for ensuring the dataset exists at the
    specified path on all worker nodes.

    In your study YAML file, specify this dataset as follows:

        dataset:
          xname: custom_dataset                             # Must use this exact string
          qa_csv_path: "/path/to/dataset.csv"               # Must be an absolute path
          grounding_data_dir: "/path/to/grounding/data/"    # Must be an absolute path
    """

    xname: T.Literal["custom_dataset"] = "custom_dataset"  # type: ignore

    qa_csv_path: Path = Path("/tmp/custom_qa_data.csv")
    grounding_data_dir: Path = Path("/tmp/grounding_docs")

    description: str = (
        "Custom local dataset for testing purposes."
        "This dataset loads QA pairs from a local CSV file and grounding data from a local directory."
    )

    sample_size: int = 10
    training_size: int = 100
    testing_size: int = 100
    holdout_size: int = 0

    def _get_partition_range(self, partition: str):
        partition_ranges = {
            "sample": range(0, self.sample_size),
            "train": range(0, self.training_size),
            "test": range(self.training_size, self.training_size + self.testing_size),
            "holdout": range(
                self.training_size + self.testing_size,
                self.training_size + self.testing_size + self.holdout_size,
            ),
        }
        if partition in partition_ranges:
            return partition_ranges[partition]
        raise ValueError(f"Unknown partition: {partition}")

    @field_validator("qa_csv_path", "grounding_data_dir", mode="after")
    @classmethod
    def check_paths_are_absolute(cls, v: Path) -> Path:
        if not v.is_absolute():
            raise ValueError(
                f"The provided dataset path `{v}` must be an absolute path"
            )
        return v

    @cached_property
    def data_root(self) -> Path:
        return Path(get_ray_temp_dir()) / "data" / "local_data"

    def _get_qa_abspath(self) -> Path:
        """
        Ray worker-visible QA CSV path.
        """
        return self.qa_csv_path.absolute()

    def _get_grounding_dir_abspath(self) -> Path:
        """
        Ray worker-visible grounding directory path.
        """
        return self.grounding_data_dir.absolute()

    def _load_qa_dataset(self, partition: str = "test") -> datasets.Dataset:
        """
        This method reads a CSV file of QA pairs and returns data of the requested partition.
        """

        df: pd.DataFrame = pd.read_csv(self._get_qa_abspath())
        partition_range = self._get_partition_range(partition)
        df_partition = df.iloc[list(partition_range)]
        return datasets.Dataset.from_pandas(df_partition)

    def _load_grounding_dataset(self) -> datasets.DatasetDict:
        reader = SimpleDirectoryReader(input_dir=self._get_grounding_dir_abspath())
        docs = reader.load_data()
        dataset = datasets.Dataset.from_pandas(
            pd.DataFrame(
                [
                    {
                        "markdown": doc.text,
                        "filename": doc.metadata.get("file_name", "unknown"),
                    }
                    for doc in docs
                ]
            )
        )
        return datasets.DatasetDict({"test": dataset})

    def _row_to_qapair(self, row) -> QAPair:
        """Dataset-specific conversion of row to QAPair struct."""
        return QAPair(
            question=row["question"],
            answer=row["answer"],
            id=str(row["id"]),
            context={},
            supporting_facts=[],
            difficulty="default",
            qtype="default",
        )

    @overrides
    def iter_examples(self, partition="test") -> T.Iterator[QAPair]:
        assert partition in self.storage_partitions
        partition = self._get_storage_partition(partition)
        qa_dataset = self._load_qa_dataset(partition)
        for row in qa_dataset:
            yield self._row_to_qapair(row)

    @overrides
    def iter_grounding_data(self, partition="test") -> T.Iterator[Document]:
        grounding_dataset = self._load_grounding_dataset()
        for row in grounding_dataset["test"]:
            yield Document(
                text=row["markdown"],
                metadata={"file_name": row["filename"]},
            )
