"""
Custom dataset template.

This module provides a template for creating custom benchmark datasets.
Copy this file and modify it to create your own dataset class.
"""

from loguru import logger
from typing import Any, Dict

from .base import BenchmarkDataset


class CustomDataset(BenchmarkDataset):
    """Template for a custom benchmark dataset.

    Copy this class and modify it to create your own dataset class.
    Replace 'Custom' with your dataset name and update the implementation.
    """

    @classmethod
    def get_dataset_info(cls) -> Dict[str, str]:
        """Get basic information about the dataset."""
        return {
            "id": "custom",  # Unique identifier for the dataset
            "name": "Custom Dataset",  # Human-readable name
            "description": "Template for a custom benchmark dataset",  # Description
            "url": cls.get_default_dataset_path(),  # Default URL or path
        }

    @classmethod
    def get_default_dataset_path(cls) -> str:
        """Get the default path or URL for the dataset."""
        return "path/to/your/dataset.csv"  # Replace with your dataset path

    def process_example(self, example: Dict[str, Any]) -> Dict[str, Any]:
        """Process a single example from the dataset.

        This is where you can transform, decrypt, or otherwise process
        the raw examples from your dataset.

        Args:
            example: Raw example from the dataset.

        Returns:
            Processed example ready for use.
        """
        # Make a copy to avoid modifying the original
        processed = dict(example)

        # TODO: Add your custom processing here
        # For example:
        # - Extract relevant fields
        # - Transform data formats
        # - Handle special cases
        # - Apply data cleaning

        # Ensure required fields are present
        if "problem" not in processed:
            logger.warning("Example missing 'problem' field")
            processed["problem"] = ""

        if "answer" not in processed:
            logger.warning("Example missing 'answer' field")
            processed["answer"] = ""

        # Add correct_answer field if not present
        if "correct_answer" not in processed:
            processed["correct_answer"] = processed["answer"]

        return processed

    def get_question(self, example: Dict[str, Any]) -> str:
        """Extract the question from an example.

        Override this method if your dataset stores the question in a
        different field than 'problem'.
        """
        # Example: return example.get("question", "")
        return example.get("problem", "")

    def get_answer(self, example: Dict[str, Any]) -> str:
        """Extract the answer from an example.

        Override this method if your dataset stores the answer in a
        different field than 'answer' or 'correct_answer'.
        """
        # Try correct_answer first, then fall back to answer
        return example.get("correct_answer", example.get("answer", ""))


# To register your dataset, add this at the bottom of your file:
# DatasetRegistry.register(CustomDataset)
#
# Then import your dataset in the __init__.py file:
# from .custom_dataset import CustomDataset
