from datasets import DatasetBuilder, DatasetInfo, SplitGenerators, DownloadManager
from datasets.features import Features, ClassLabel, Sequence, Value

class MyCustomDataset(DatasetBuilder):
    VERSION = datasets.Version("1.0.0")

    def _info(self):
        return DatasetInfo(
            description="My custom dataset for tracking objects.",
            features=Features({
                "prompting_type": Value("string"),
                "deception": Value("bool"),
                "story_length": Value("int32"),
                "question_order": Value("int32"),
                "sample_id": Value("int32"),
                "story": Value("string"),
                "question": Value("string"),
                "choices": Value("string"),
                "answer": Value("string"),
            }),
            supervised_keys=None,
            homepage="https://github.com/ying-hui-he/Hi-ToM_dataset",
            citation=CITATION,
        )

    def _split_generators(self, dl_manager: DownloadManager):
        downloaded_files = dl_manager.download_and_extract({
            "data_file": "https://github.com/ying-hui-he/Hi-ToM_dataset/blob/main/Hi-ToM_data.json"
        })

        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,
                gen_kwargs={
                    "filepath": downloaded_files["data_file"],
                },
            ),
        ]

    def _generate_examples(self, filepath):
        with open(filepath, encoding="utf-8") as f:
            data = json.load(f)
            for id, item in enumerate(data["data"]):
                yield id, {
                    "prompting_type": item["prompting_type"],
                    "deception": item["deception"],
                    "story_length": item["story_length"],
                    "question_order": item["question_order"],
                    "sample_id": item["sample_id"],
                    "story": item["story"],
                    "question": item["question"],
                    "choices": item["choices"],
                    "answer": item["answer"],
                }