import os
import tempfile
from types import SimpleNamespace
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pandas as pd
import pytest
from fsspec.implementations.memory import MemoryFileSystem

from datasets import load_from_disk
from datasets.arrow_dataset import Dataset
from datasets.dataset_dict import DatasetDict, IterableDatasetDict
from datasets.features import ClassLabel, Features, List, Value
from datasets.iterable_dataset import IterableDataset
from datasets.splits import NamedSplit, SplitInfo

from .utils import (
    assert_arrow_memory_doesnt_increase,
    assert_arrow_memory_increases,
    require_numpy1_on_windows,
    require_polars,
    require_tf,
    require_torch,
)


class DatasetDictTest(TestCase):
    def _create_dummy_dataset(self, multiple_columns=False, int_to_float=False):
        if multiple_columns:
            data = {"col_1": [3, 2, 1, 0], "col_2": ["a", "b", "c", "d"]}
            dset = Dataset.from_dict(data)
        elif int_to_float:
            data = {
                "text": ["text1", "text2", "text3", "text4"],
                "labels": [[1, 1, 1, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 1, 1], [0, 0, 0, 1, 0]],
            }
            dset = Dataset.from_dict(data)
        else:
            dset = Dataset.from_dict(
                {"filename": ["my_name-train" + "_" + f"{x:03d}" for x in np.arange(30).tolist()]}
            )
        return dset

    def _create_dummy_dataset_dict(self, multiple_columns=False, int_to_float=False) -> DatasetDict:
        return DatasetDict(
            {
                "train": self._create_dummy_dataset(multiple_columns=multiple_columns, int_to_float=int_to_float),
                "test": self._create_dummy_dataset(multiple_columns=multiple_columns, int_to_float=int_to_float),
            }
        )

    def _create_dummy_iterable_dataset(self, multiple_columns=False) -> IterableDataset:
        def gen():
            if multiple_columns:
                data = {"col_1": [3, 2, 1, 0], "col_2": ["a", "b", "c", "d"]}
                for v1, v2 in zip(data["col_1"], data["col_2"]):
                    yield {"col_1": v1, "col_2": v2}
            else:
                for x in range(30):
                    yield {"filename": "my_name-train" + "_" + f"{x:03d}"}

        return IterableDataset.from_generator(gen)

    def _create_dummy_iterable_dataset_dict(self, multiple_columns=False) -> IterableDatasetDict:
        return IterableDatasetDict(
            {
                "train": self._create_dummy_iterable_dataset(multiple_columns=multiple_columns),
                "test": self._create_dummy_iterable_dataset(multiple_columns=multiple_columns),
            }
        )

    def test_flatten(self):
        dset_split = Dataset.from_dict(
            {"a": [{"b": {"c": ["text"]}}] * 10, "foo": [1] * 10},
            features=Features({"a": {"b": {"c": List(Value("string"))}}, "foo": Value("int64")}),
        )
        dset = DatasetDict({"train": dset_split, "test": dset_split})
        dset = dset.flatten()
        self.assertDictEqual(dset.column_names, {"train": ["a.b.c", "foo"], "test": ["a.b.c", "foo"]})
        self.assertListEqual(sorted(dset["train"].features.keys()), ["a.b.c", "foo"])
        self.assertDictEqual(dset["train"].features, Features({"a.b.c": List(Value("string")), "foo": Value("int64")}))
        del dset

    def test_set_format_numpy(self):
        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset.set_format(type="numpy", columns=["col_1"])
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0]), 1)
            self.assertIsInstance(dset_split[0]["col_1"], np.int64)
            self.assertEqual(dset_split[0]["col_1"].item(), 3)

        dset.reset_format()
        with dset.formatted_as(type="numpy", columns=["col_1"]):
            for dset_split in dset.values():
                self.assertEqual(len(dset_split[0]), 1)
                self.assertIsInstance(dset_split[0]["col_1"], np.int64)
                self.assertEqual(dset_split[0]["col_1"].item(), 3)

        for dset_split in dset.values():
            self.assertEqual(dset_split.format["type"], None)
            self.assertEqual(dset_split.format["format_kwargs"], {})
            self.assertEqual(dset_split.format["columns"], dset_split.column_names)
            self.assertEqual(dset_split.format["output_all_columns"], False)

        dset.set_format(type="numpy", columns=["col_1"], output_all_columns=True)
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0]), 2)
            self.assertIsInstance(dset_split[0]["col_2"], str)
            self.assertEqual(dset_split[0]["col_2"], "a")

        dset.set_format(type="numpy", columns=["col_1", "col_2"])
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0]), 2)
            self.assertIsInstance(dset_split[0]["col_2"], np.str_)
            self.assertEqual(dset_split[0]["col_2"].item(), "a")
        del dset

    @require_numpy1_on_windows
    @require_torch
    def test_set_format_torch(self):
        import torch

        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset.set_format(type="torch", columns=["col_1"])
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0]), 1)
            self.assertIsInstance(dset_split[0]["col_1"], torch.Tensor)
            self.assertListEqual(list(dset_split[0]["col_1"].shape), [])
            self.assertEqual(dset_split[0]["col_1"].item(), 3)

        dset.set_format(type="torch", columns=["col_1"], output_all_columns=True)
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0]), 2)
            self.assertIsInstance(dset_split[0]["col_2"], str)
            self.assertEqual(dset_split[0]["col_2"], "a")

        dset.set_format(type="torch")
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0]), 2)
            self.assertIsInstance(dset_split[0]["col_1"], torch.Tensor)
            self.assertListEqual(list(dset_split[0]["col_1"].shape), [])
            self.assertEqual(dset_split[0]["col_1"].item(), 3)
            self.assertIsInstance(dset_split[0]["col_2"], str)
            self.assertEqual(dset_split[0]["col_2"], "a")
        del dset

    @require_tf
    def test_set_format_tf(self):
        import tensorflow as tf

        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset.set_format(type="tensorflow", columns=["col_1"])
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0]), 1)
            self.assertIsInstance(dset_split[0]["col_1"], tf.Tensor)
            self.assertListEqual(list(dset_split[0]["col_1"].shape), [])
            self.assertEqual(dset_split[0]["col_1"].numpy().item(), 3)

        dset.set_format(type="tensorflow", columns=["col_1"], output_all_columns=True)
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0]), 2)
            self.assertIsInstance(dset_split[0]["col_2"], str)
            self.assertEqual(dset_split[0]["col_2"], "a")

        dset.set_format(type="tensorflow", columns=["col_1", "col_2"])
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0]), 2)
            self.assertEqual(dset_split[0]["col_2"].numpy().decode("utf-8"), "a")
        del dset

    def test_set_format_pandas(self):
        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset.set_format(type="pandas", columns=["col_1"])
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0].columns), 1)
            self.assertIsInstance(dset_split[0], pd.DataFrame)
            self.assertListEqual(list(dset_split[0].shape), [1, 1])
            self.assertEqual(dset_split[0]["col_1"].item(), 3)

        dset.set_format(type="pandas", columns=["col_1", "col_2"])
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0].columns), 2)
            self.assertEqual(dset_split[0]["col_2"].item(), "a")
        del dset

    @require_polars
    def test_set_format_polars(self):
        import polars as pl

        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset.set_format(type="polars", columns=["col_1"])
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0].columns), 1)
            self.assertIsInstance(dset_split[0], pl.DataFrame)
            self.assertEqual(dset_split[0].shape, (1, 1))
            self.assertEqual(dset_split[0]["col_1"].item(), 3)

        dset.set_format(type="polars", columns=["col_1", "col_2"])
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0].columns), 2)
            self.assertEqual(dset_split[0]["col_2"].item(), "a")
        del dset

    def test_set_transform(self):
        def transform(batch):
            return {k: [str(i).upper() for i in v] for k, v in batch.items()}

        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset.set_transform(transform=transform, columns=["col_1"])
        for dset_split in dset.values():
            self.assertEqual(dset_split.format["type"], "custom")
            self.assertEqual(len(dset_split[0].keys()), 1)
            self.assertEqual(dset_split[0]["col_1"], "3")
            self.assertEqual(dset_split[:2]["col_1"], ["3", "2"])
            self.assertEqual(dset_split["col_1"][:2], ["3", "2"])

        prev_format = dset[list(dset.keys())[0]].format
        for dset_split in dset.values():
            dset_split.set_format(**dset_split.format)
            self.assertEqual(prev_format, dset_split.format)

        dset.set_transform(transform=transform, columns=["col_1", "col_2"])
        for dset_split in dset.values():
            self.assertEqual(len(dset_split[0].keys()), 2)
            self.assertEqual(dset_split[0]["col_2"], "A")
        del dset

    def test_with_format(self):
        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset2 = dset.with_format("numpy", columns=["col_1"])
        dset.set_format("numpy", columns=["col_1"])
        for dset_split, dset_split2 in zip(dset.values(), dset2.values()):
            self.assertDictEqual(dset_split.format, dset_split2.format)
        del dset, dset2

    def test_with_transform(self):
        def transform(batch):
            return {k: [str(i).upper() for i in v] for k, v in batch.items()}

        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset2 = dset.with_transform(transform, columns=["col_1"])
        dset.set_transform(transform, columns=["col_1"])
        for dset_split, dset_split2 in zip(dset.values(), dset2.values()):
            self.assertDictEqual(dset_split.format, dset_split2.format)
        del dset, dset2

    def test_cast(self):
        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        features = dset["train"].features
        features["col_1"] = Value("float64")
        dset = dset.cast(features)
        for dset_split in dset.values():
            self.assertEqual(dset_split.num_columns, 2)
            self.assertEqual(dset_split.features["col_1"], Value("float64"))
            self.assertIsInstance(dset_split[0]["col_1"], float)
        del dset

    def test_remove_columns(self):
        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset = dset.remove_columns(column_names="col_1")
        for dset_split in dset.values():
            self.assertEqual(dset_split.num_columns, 1)
            self.assertListEqual(list(dset_split.column_names), ["col_2"])

        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset = dset.remove_columns(column_names=["col_1", "col_2"])
        for dset_split in dset.values():
            self.assertEqual(dset_split.num_columns, 0)

        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        for dset_split in dset.values():
            dset_split._format_columns = ["col_1", "col_2"]
        dset = dset.remove_columns(column_names=["col_1"])
        for dset_split in dset.values():
            self.assertListEqual(dset_split._format_columns, ["col_2"])
            self.assertEqual(dset_split.num_columns, 1)
            self.assertListEqual(list(dset_split.column_names), ["col_2"])
        del dset

    def test_rename_column(self):
        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset = dset.rename_column(original_column_name="col_1", new_column_name="new_name")
        for dset_split in dset.values():
            self.assertEqual(dset_split.num_columns, 2)
            self.assertListEqual(list(dset_split.column_names), ["new_name", "col_2"])
        del dset

    def test_select_columns(self):
        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset = dset.select_columns(column_names=[])
        for dset_split in dset.values():
            self.assertEqual(dset_split.num_columns, 0)

        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset = dset.select_columns(column_names="col_1")
        for dset_split in dset.values():
            self.assertEqual(dset_split.num_columns, 1)
            self.assertListEqual(list(dset_split.column_names), ["col_1"])

        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        dset = dset.select_columns(column_names=["col_1", "col_2"])
        for dset_split in dset.values():
            self.assertEqual(dset_split.num_columns, 2)

        dset = self._create_dummy_dataset_dict(multiple_columns=True)
        for dset_split in dset.values():
            dset_split._format_columns = ["col_1", "col_2"]
        dset = dset.select_columns(column_names=["col_1"])
        for dset_split in dset.values():
            self.assertEqual(dset_split.num_columns, 1)
            self.assertListEqual(list(dset_split.column_names), ["col_1"])
            self.assertListEqual(dset_split._format_columns, ["col_1"])

    def test_map(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            dsets = self._create_dummy_dataset_dict()

            mapped_dsets_1: DatasetDict = dsets.map(lambda ex: {"foo": ["bar"] * len(ex["filename"])}, batched=True)
            self.assertListEqual(list(dsets.keys()), list(mapped_dsets_1.keys()))
            self.assertListEqual(mapped_dsets_1["train"].column_names, ["filename", "foo"])

            cache_file_names = {
                "train": os.path.join(tmp_dir, "train.arrow"),
                "test": os.path.join(tmp_dir, "test.arrow"),
            }
            mapped_dsets_2: DatasetDict = mapped_dsets_1.map(
                lambda ex: {"bar": ["foo"] * len(ex["filename"])}, batched=True, cache_file_names=cache_file_names
            )
            self.assertListEqual(list(dsets.keys()), list(mapped_dsets_2.keys()))
            self.assertListEqual(sorted(mapped_dsets_2["train"].column_names), sorted(["filename", "foo", "bar"]))
            del dsets, mapped_dsets_1, mapped_dsets_2

        # casting int labels to float labels
        with tempfile.TemporaryDirectory() as tmp_dir:
            dset_dict = self._create_dummy_dataset_dict(int_to_float=True)

            def _preprocess(examples):
                result = {"labels": [list(map(float, labels)) for labels in examples["labels"]]}
                return result

            with dset_dict.map(
                _preprocess, remove_columns=["labels", "text"], batched=True, try_original_type=True
            ) as dset_test:
                for labels in dset_test["test"]["labels"]:
                    for label in labels:
                        self.assertIsInstance(label, int)

            with dset_dict.map(
                _preprocess, remove_columns=["labels", "text"], batched=True, try_original_type=False
            ) as dset_test:
                for labels in dset_test["test"]["labels"]:
                    for label in labels:
                        self.assertIsInstance(label, float)

    def test_iterable_map(self):
        dsets = self._create_dummy_iterable_dataset_dict()
        fn_kwargs = {"n": 3}
        mapped_dsets: IterableDatasetDict = dsets.map(
            lambda x, n: {"foo": [n] * len(x["filename"])},
            batched=True,
            fn_kwargs=fn_kwargs,
        )
        mapped_example = next(iter(mapped_dsets["train"]))
        self.assertListEqual(sorted(mapped_example.keys()), sorted(["filename", "foo"]))
        self.assertLessEqual(mapped_example["foo"], 3)
        del dsets, mapped_dsets

    def test_filter(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            dsets = self._create_dummy_dataset_dict()

            filtered_dsets_1: DatasetDict = dsets.filter(lambda ex: int(ex["filename"].split("_")[-1]) < 10)
            self.assertListEqual(list(dsets.keys()), list(filtered_dsets_1.keys()))
            self.assertEqual(len(filtered_dsets_1["train"]), 10)

            cache_file_names = {
                "train": os.path.join(tmp_dir, "train.arrow"),
                "test": os.path.join(tmp_dir, "test.arrow"),
            }
            filtered_dsets_2: DatasetDict = filtered_dsets_1.filter(
                lambda ex: int(ex["filename"].split("_")[-1]) < 5, cache_file_names=cache_file_names
            )
            self.assertListEqual(list(dsets.keys()), list(filtered_dsets_2.keys()))
            self.assertEqual(len(filtered_dsets_2["train"]), 5)

            filtered_dsets_3: DatasetDict = dsets.filter(
                lambda examples: [int(ex.split("_")[-1]) < 10 for ex in examples["filename"]], batched=True
            )
            self.assertListEqual(list(dsets.keys()), list(filtered_dsets_3.keys()))
            self.assertEqual(len(filtered_dsets_3["train"]), 10)
            del dsets, filtered_dsets_1, filtered_dsets_2, filtered_dsets_3

    def test_iterable_filter(self):
        dsets = self._create_dummy_iterable_dataset_dict()
        example = next(iter(dsets["train"]))
        fn_kwargs = {"n": 3}
        filtered_dsets: IterableDatasetDict = dsets.filter(
            lambda ex, n: n < int(ex["filename"].split("_")[-1]), fn_kwargs=fn_kwargs
        )
        filtered_example = next(iter(filtered_dsets["train"]))
        self.assertListEqual(list(example.keys()), list(filtered_example.keys()))
        self.assertEqual(int(filtered_example["filename"].split("_")[-1]), 4)  # id starts from 3
        del dsets, filtered_dsets

    def test_iterable_dataset_dict_push_to_hub_max_shard_size_and_num_shards_are_mutually_exclusive(self):
        dsets = self._create_dummy_iterable_dataset_dict()
        with pytest.raises(ValueError, match="either max_shard_size or num_shards"):
            dsets.push_to_hub("user/dataset", max_shard_size="1MB", num_shards={"train": 1, "test": 1})

    def test_iterable_dataset_dict_push_to_hub_forwards_max_shard_size_to_each_split(self):
        class DummyApi:
            def __init__(self, *args, **kwargs):
                pass

            def repo_info(self, repo_id, repo_type="dataset", revision=None):
                return SimpleNamespace(id=repo_id, sha="dummy-sha")

            def create_branch(self, *args, **kwargs):
                pass

            def list_repo_tree(self, *args, **kwargs):
                return []

            def create_commit(self, *args, **kwargs):
                return SimpleNamespace(commit_url="https://hf.co/commit/dummy")

        dummy_fs = MemoryFileSystem(skip_instance_cache=True)
        dummy_fs.touch("datasets/user/dataset@dummy-sha/README.md")

        forwarded_calls = []

        def mock_push_parquet_shards_to_hub(
            resolved_output_path,
            data_dir,
            split,
            token,
            create_pr,
            max_shard_size,
            num_shards,
            embed_external_files,
            num_proc,
        ):
            forwarded_calls.append(
                {
                    "split": split,
                    "max_shard_size": max_shard_size,
                    "num_shards": num_shards,
                }
            )
            return [], [], Features(), SplitInfo(name=split), 0

        dsets = self._create_dummy_iterable_dataset_dict()
        max_shard_size = sum(split_dataset.num_shards for split_dataset in dsets.values())
        with (
            patch("datasets.dataset_dict.HfApi", DummyApi),
            patch("datasets.dataset_dict.HfFileSystem", return_value=dummy_fs),
            patch.object(IterableDataset, "_push_parquet_shards_to_hub", side_effect=mock_push_parquet_shards_to_hub),
        ):
            dsets.push_to_hub("user/dataset", max_shard_size=max_shard_size)

        assert {call["split"] for call in forwarded_calls} == set(dsets.keys())
        assert all(call["max_shard_size"] == max_shard_size for call in forwarded_calls)
        assert all(call["num_shards"] is None for call in forwarded_calls)

    def test_sort(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            dsets = self._create_dummy_dataset_dict()

            sorted_dsets_1: DatasetDict = dsets.sort("filename")
            self.assertListEqual(list(dsets.keys()), list(sorted_dsets_1.keys()))
            self.assertListEqual(
                [f.split("_")[-1] for f in sorted_dsets_1["train"]["filename"]],
                sorted(f"{x:03d}" for x in range(30)),
            )

            indices_cache_file_names = {
                "train": os.path.join(tmp_dir, "train.arrow"),
                "test": os.path.join(tmp_dir, "test.arrow"),
            }
            sorted_dsets_2: DatasetDict = sorted_dsets_1.sort(
                "filename", indices_cache_file_names=indices_cache_file_names, reverse=True
            )
            self.assertListEqual(list(dsets.keys()), list(sorted_dsets_2.keys()))
            self.assertListEqual(
                [f.split("_")[-1] for f in sorted_dsets_2["train"]["filename"]],
                sorted((f"{x:03d}" for x in range(30)), reverse=True),
            )
            del dsets, sorted_dsets_1, sorted_dsets_2

    def test_shuffle(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            dsets = self._create_dummy_dataset_dict()

            indices_cache_file_names = {
                "train": os.path.join(tmp_dir, "train.arrow"),
                "test": os.path.join(tmp_dir, "test.arrow"),
            }
            seeds = {
                "train": 1234,
                "test": 1234,
            }
            dsets_shuffled = dsets.shuffle(
                seeds=seeds, indices_cache_file_names=indices_cache_file_names, load_from_cache_file=False
            )
            self.assertSequenceEqual(dsets_shuffled["train"]["filename"], dsets_shuffled["test"]["filename"])

            self.assertEqual(len(dsets_shuffled["train"]), 30)
            self.assertEqual(dsets_shuffled["train"][0]["filename"], "my_name-train_028")
            self.assertEqual(dsets_shuffled["train"][2]["filename"], "my_name-train_010")
            self.assertDictEqual(dsets["train"].features, Features({"filename": Value("string")}))
            self.assertDictEqual(dsets_shuffled["train"].features, Features({"filename": Value("string")}))

            # Reproducibility
            indices_cache_file_names_2 = {
                "train": os.path.join(tmp_dir, "train_2.arrow"),
                "test": os.path.join(tmp_dir, "test_2.arrow"),
            }
            dsets_shuffled_2 = dsets.shuffle(
                seeds=seeds, indices_cache_file_names=indices_cache_file_names_2, load_from_cache_file=False
            )
            self.assertSequenceEqual(dsets_shuffled["train"]["filename"], dsets_shuffled_2["train"]["filename"])

            seeds = {
                "train": 1234,
                "test": 1,
            }
            indices_cache_file_names_3 = {
                "train": os.path.join(tmp_dir, "train_3.arrow"),
                "test": os.path.join(tmp_dir, "test_3.arrow"),
            }
            dsets_shuffled_3 = dsets.shuffle(
                seeds=seeds, indices_cache_file_names=indices_cache_file_names_3, load_from_cache_file=False
            )
            self.assertNotEqual(dsets_shuffled_3["train"]["filename"], dsets_shuffled_3["test"]["filename"])

            # other input types
            dsets_shuffled_int = dsets.shuffle(42)
            dsets_shuffled_alias = dsets.shuffle(seed=42)
            dsets_shuffled_none = dsets.shuffle()
            self.assertEqual(len(dsets_shuffled_int["train"]), 30)
            self.assertEqual(len(dsets_shuffled_alias["train"]), 30)
            self.assertEqual(len(dsets_shuffled_none["train"]), 30)

            del dsets, dsets_shuffled, dsets_shuffled_2, dsets_shuffled_3
            del dsets_shuffled_int, dsets_shuffled_alias, dsets_shuffled_none

    def test_flatten_indices(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            dsets = self._create_dummy_dataset_dict()

            indices_cache_file_names = {
                "train": os.path.join(tmp_dir, "train.arrow"),
                "test": os.path.join(tmp_dir, "test.arrow"),
            }
            dsets_shuffled = dsets.shuffle(
                seed=42, indices_cache_file_names=indices_cache_file_names, load_from_cache_file=False
            )

            self.assertIsNotNone(dsets_shuffled["train"]._indices)
            self.assertIsNotNone(dsets_shuffled["test"]._indices)

            dsets_flat = dsets_shuffled.flatten_indices()

            self.assertIsNone(dsets_flat["train"]._indices)
            self.assertIsNone(dsets_flat["test"]._indices)

            del dsets, dsets_shuffled, dsets_flat

    def test_check_values_type(self):
        dsets = self._create_dummy_dataset_dict()
        dsets["bad_split"] = None
        self.assertRaises(TypeError, dsets.map, lambda x: x)
        self.assertRaises(TypeError, dsets.filter, lambda x: True)
        self.assertRaises(TypeError, dsets.shuffle)
        self.assertRaises(TypeError, dsets.sort, "filename")
        del dsets

    def test_serialization(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            dsets = self._create_dummy_dataset_dict()
            dsets.save_to_disk(tmp_dir)
            reloaded_dsets = DatasetDict.load_from_disk(tmp_dir)
            self.assertListEqual(sorted(reloaded_dsets), ["test", "train"])
            self.assertEqual(len(reloaded_dsets["train"]), 30)
            self.assertListEqual(reloaded_dsets["train"].column_names, ["filename"])
            self.assertEqual(len(reloaded_dsets["test"]), 30)
            self.assertListEqual(reloaded_dsets["test"].column_names, ["filename"])
            del reloaded_dsets

            del dsets["test"]
            dsets.save_to_disk(tmp_dir)
            reloaded_dsets = DatasetDict.load_from_disk(tmp_dir)
            self.assertListEqual(sorted(reloaded_dsets), ["train"])
            self.assertEqual(len(reloaded_dsets["train"]), 30)
            self.assertListEqual(reloaded_dsets["train"].column_names, ["filename"])
            del dsets, reloaded_dsets

            dsets = self._create_dummy_dataset_dict()
            dsets.save_to_disk(tmp_dir, num_shards={"train": 3, "test": 2})
            reloaded_dsets = DatasetDict.load_from_disk(tmp_dir)
            self.assertListEqual(sorted(reloaded_dsets), ["test", "train"])
            self.assertEqual(len(reloaded_dsets["train"]), 30)
            self.assertListEqual(reloaded_dsets["train"].column_names, ["filename"])
            self.assertEqual(len(reloaded_dsets["train"].cache_files), 3)
            self.assertEqual(len(reloaded_dsets["test"]), 30)
            self.assertListEqual(reloaded_dsets["test"].column_names, ["filename"])
            self.assertEqual(len(reloaded_dsets["test"].cache_files), 2)
            del reloaded_dsets

            dsets = self._create_dummy_dataset_dict()
            dsets.save_to_disk(tmp_dir, num_proc=2)
            reloaded_dsets = DatasetDict.load_from_disk(tmp_dir)
            self.assertListEqual(sorted(reloaded_dsets), ["test", "train"])
            self.assertEqual(len(reloaded_dsets["train"]), 30)
            self.assertListEqual(reloaded_dsets["train"].column_names, ["filename"])
            self.assertEqual(len(reloaded_dsets["train"].cache_files), 2)
            self.assertEqual(len(reloaded_dsets["test"]), 30)
            self.assertListEqual(reloaded_dsets["test"].column_names, ["filename"])
            self.assertEqual(len(reloaded_dsets["test"].cache_files), 2)
            del reloaded_dsets

    def test_load_from_disk(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            dsets = self._create_dummy_dataset_dict()
            dsets.save_to_disk(tmp_dir)
            del dsets
            dsets = load_from_disk(tmp_dir)
            self.assertListEqual(sorted(dsets), ["test", "train"])
            self.assertEqual(len(dsets["train"]), 30)
            self.assertListEqual(dsets["train"].column_names, ["filename"])
            self.assertEqual(len(dsets["test"]), 30)
            self.assertListEqual(dsets["test"].column_names, ["filename"])
            del dsets

    def test_align_labels_with_mapping(self):
        train_features = Features(
            {
                "input_text": Value("string"),
                "input_labels": ClassLabel(num_classes=3, names=["entailment", "neutral", "contradiction"]),
            }
        )
        test_features = Features(
            {
                "input_text": Value("string"),
                "input_labels": ClassLabel(num_classes=3, names=["entailment", "contradiction", "neutral"]),
            }
        )
        train_data = {"input_text": ["a", "a", "b", "b", "c", "c"], "input_labels": [0, 0, 1, 1, 2, 2]}
        test_data = {"input_text": ["a", "a", "c", "c", "b", "b"], "input_labels": [0, 0, 1, 1, 2, 2]}
        label2id = {"CONTRADICTION": 0, "ENTAILMENT": 2, "NEUTRAL": 1}
        id2label = {v: k for k, v in label2id.items()}
        train_expected_labels = [2, 2, 1, 1, 0, 0]
        test_expected_labels = [2, 2, 0, 0, 1, 1]
        train_expected_label_names = [id2label[idx] for idx in train_expected_labels]
        test_expected_label_names = [id2label[idx] for idx in test_expected_labels]
        dsets = DatasetDict(
            {
                "train": Dataset.from_dict(train_data, features=train_features),
                "test": Dataset.from_dict(test_data, features=test_features),
            }
        )
        dsets = dsets.align_labels_with_mapping(label2id, "input_labels")
        self.assertListEqual(train_expected_labels, dsets["train"]["input_labels"][:])
        self.assertListEqual(test_expected_labels, dsets["test"]["input_labels"][:])
        train_aligned_label_names = [
            dsets["train"].features["input_labels"].int2str(idx) for idx in dsets["train"]["input_labels"]
        ]
        test_aligned_label_names = [
            dsets["test"].features["input_labels"].int2str(idx) for idx in dsets["test"]["input_labels"]
        ]
        self.assertListEqual(train_expected_label_names, train_aligned_label_names)
        self.assertListEqual(test_expected_label_names, test_aligned_label_names)


def test_dummy_datasetdict_serialize_fs(mockfs):
    dataset_dict = DatasetDict(
        {
            "train": Dataset.from_dict({"a": range(30)}),
            "test": Dataset.from_dict({"a": range(10)}),
        }
    )
    dataset_path = "mock://my_dataset"
    dataset_dict.save_to_disk(dataset_path, storage_options=mockfs.storage_options)
    assert mockfs.isdir(dataset_path)
    assert mockfs.glob(dataset_path + "/*")
    reloaded = DatasetDict.load_from_disk(dataset_path, storage_options=mockfs.storage_options)
    assert list(reloaded) == list(dataset_dict)
    for k in dataset_dict:
        assert reloaded[k].features == dataset_dict[k].features
        assert reloaded[k].to_dict() == dataset_dict[k].to_dict()


def _check_csv_datasetdict(dataset_dict, expected_features, splits=("train",)):
    assert isinstance(dataset_dict, DatasetDict)
    for split in splits:
        dataset = dataset_dict[split]
        assert dataset.num_rows == 4
        assert dataset.num_columns == 3
        assert dataset.column_names == ["col_1", "col_2", "col_3"]
        for feature, expected_dtype in expected_features.items():
            assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("keep_in_memory", [False, True])
def test_datasetdict_from_csv_keep_in_memory(keep_in_memory, csv_path, tmp_path):
    cache_dir = tmp_path / "cache"
    expected_features = {"col_1": "int64", "col_2": "int64", "col_3": "float64"}
    with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
        dataset = DatasetDict.from_csv({"train": csv_path}, cache_dir=cache_dir, keep_in_memory=keep_in_memory)
    _check_csv_datasetdict(dataset, expected_features)


@pytest.mark.parametrize(
    "features",
    [
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ],
)
def test_datasetdict_from_csv_features(features, csv_path, tmp_path):
    cache_dir = tmp_path / "cache"
    # CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
    default_expected_features = {"col_1": "int64", "col_2": "int64", "col_3": "float64"}
    expected_features = features.copy() if features else default_expected_features
    features = (
        Features({feature: Value(dtype) for feature, dtype in features.items()}) if features is not None else None
    )
    dataset = DatasetDict.from_csv({"train": csv_path}, features=features, cache_dir=cache_dir)
    _check_csv_datasetdict(dataset, expected_features)


@pytest.mark.parametrize("split", [None, NamedSplit("train"), "train", "test"])
def test_datasetdict_from_csv_split(split, csv_path, tmp_path):
    if split:
        path = {split: csv_path}
    else:
        split = "train"
        path = {"train": csv_path, "test": csv_path}
    cache_dir = tmp_path / "cache"
    expected_features = {"col_1": "int64", "col_2": "int64", "col_3": "float64"}
    dataset = DatasetDict.from_csv(path, cache_dir=cache_dir)
    _check_csv_datasetdict(dataset, expected_features, splits=list(path.keys()))
    assert all(dataset[split].split == split for split in path.keys())


def _check_json_datasetdict(dataset_dict, expected_features, splits=("train",)):
    assert isinstance(dataset_dict, DatasetDict)
    for split in splits:
        dataset = dataset_dict[split]
        assert dataset.num_rows == 4
        assert dataset.num_columns == 3
        assert dataset.column_names == ["col_1", "col_2", "col_3"]
        for feature, expected_dtype in expected_features.items():
            assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("keep_in_memory", [False, True])
def test_datasetdict_from_json_keep_in_memory(keep_in_memory, jsonl_path, tmp_path):
    cache_dir = tmp_path / "cache"
    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
        dataset = DatasetDict.from_json({"train": jsonl_path}, cache_dir=cache_dir, keep_in_memory=keep_in_memory)
    _check_json_datasetdict(dataset, expected_features)


@pytest.mark.parametrize(
    "features",
    [
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ],
)
def test_datasetdict_from_json_features(features, jsonl_path, tmp_path):
    cache_dir = tmp_path / "cache"
    default_expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    expected_features = features.copy() if features else default_expected_features
    features = (
        Features({feature: Value(dtype) for feature, dtype in features.items()}) if features is not None else None
    )
    dataset = DatasetDict.from_json({"train": jsonl_path}, features=features, cache_dir=cache_dir)
    _check_json_datasetdict(dataset, expected_features)


@pytest.mark.parametrize("split", [None, NamedSplit("train"), "train", "test"])
def test_datasetdict_from_json_splits(split, jsonl_path, tmp_path):
    if split:
        path = {split: jsonl_path}
    else:
        split = "train"
        path = {"train": jsonl_path, "test": jsonl_path}
    cache_dir = tmp_path / "cache"
    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    dataset = DatasetDict.from_json(path, cache_dir=cache_dir)
    _check_json_datasetdict(dataset, expected_features, splits=list(path.keys()))
    assert all(dataset[split].split == split for split in path.keys())


def _check_parquet_datasetdict(dataset_dict, expected_features, splits=("train",)):
    assert isinstance(dataset_dict, DatasetDict)
    for split in splits:
        dataset = dataset_dict[split]
        assert dataset.num_rows == 4
        assert dataset.num_columns == 3
        assert dataset.column_names == ["col_1", "col_2", "col_3"]
        for feature, expected_dtype in expected_features.items():
            assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("keep_in_memory", [False, True])
def test_datasetdict_from_parquet_keep_in_memory(keep_in_memory, parquet_path, tmp_path):
    cache_dir = tmp_path / "cache"
    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
        dataset = DatasetDict.from_parquet({"train": parquet_path}, cache_dir=cache_dir, keep_in_memory=keep_in_memory)
    _check_parquet_datasetdict(dataset, expected_features)


@pytest.mark.parametrize(
    "features",
    [
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ],
)
def test_datasetdict_from_parquet_features(features, parquet_path, tmp_path):
    cache_dir = tmp_path / "cache"
    default_expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    expected_features = features.copy() if features else default_expected_features
    features = (
        Features({feature: Value(dtype) for feature, dtype in features.items()}) if features is not None else None
    )
    dataset = DatasetDict.from_parquet({"train": parquet_path}, features=features, cache_dir=cache_dir)
    _check_parquet_datasetdict(dataset, expected_features)


@pytest.mark.parametrize("split", [None, NamedSplit("train"), "train", "test"])
def test_datasetdict_from_parquet_split(split, parquet_path, tmp_path):
    if split:
        path = {split: parquet_path}
    else:
        split = "train"
        path = {"train": parquet_path, "test": parquet_path}
    cache_dir = tmp_path / "cache"
    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
    dataset = DatasetDict.from_parquet(path, cache_dir=cache_dir)
    _check_parquet_datasetdict(dataset, expected_features, splits=list(path.keys()))
    assert all(dataset[split].split == split for split in path.keys())


def _check_text_datasetdict(dataset_dict, expected_features, splits=("train",)):
    assert isinstance(dataset_dict, DatasetDict)
    for split in splits:
        dataset = dataset_dict[split]
        assert dataset.num_rows == 4
        assert dataset.num_columns == 1
        assert dataset.column_names == ["text"]
        for feature, expected_dtype in expected_features.items():
            assert dataset.features[feature].dtype == expected_dtype


@pytest.mark.parametrize("keep_in_memory", [False, True])
def test_datasetdict_from_text_keep_in_memory(keep_in_memory, text_path, tmp_path):
    cache_dir = tmp_path / "cache"
    expected_features = {"text": "string"}
    with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
        dataset = DatasetDict.from_text({"train": text_path}, cache_dir=cache_dir, keep_in_memory=keep_in_memory)
    _check_text_datasetdict(dataset, expected_features)


@pytest.mark.parametrize(
    "features",
    [
        None,
        {"text": "string"},
        {"text": "int32"},
        {"text": "float32"},
    ],
)
def test_datasetdict_from_text_features(features, text_path, tmp_path):
    cache_dir = tmp_path / "cache"
    default_expected_features = {"text": "string"}
    expected_features = features.copy() if features else default_expected_features
    features = (
        Features({feature: Value(dtype) for feature, dtype in features.items()}) if features is not None else None
    )
    dataset = DatasetDict.from_text({"train": text_path}, features=features, cache_dir=cache_dir)
    _check_text_datasetdict(dataset, expected_features)


@pytest.mark.parametrize("split", [None, NamedSplit("train"), "train", "test"])
def test_datasetdict_from_text_split(split, text_path, tmp_path):
    if split:
        path = {split: text_path}
    else:
        split = "train"
        path = {"train": text_path, "test": text_path}
    cache_dir = tmp_path / "cache"
    expected_features = {"text": "string"}
    dataset = DatasetDict.from_text(path, cache_dir=cache_dir)
    _check_text_datasetdict(dataset, expected_features, splits=list(path.keys()))
    assert all(dataset[split].split == split for split in path.keys())
