import random
from typing import Optional
from unittest.mock import MagicMock

import pytest

import ray
from ray import train
from ray.data import DataIterator
from ray.data._internal.execution.interfaces.execution_options import (
    ExecutionOptions,
    ExecutionResources,
)
from ray.tests.conftest import *  # noqa
from ray.train import DataConfig, ScalingConfig
from ray.train.data_parallel_trainer import DataParallelTrainer


@pytest.fixture
def ray_start_4_cpus():
    address_info = ray.init(num_cpus=4)
    yield address_info
    ray.shutdown()


class TestBasic(DataParallelTrainer):
    def __init__(
        self, num_workers: int, expect_ds: bool, expect_sizes: Optional[dict], **kwargs
    ):
        def train_loop_per_worker():
            data_shard = train.get_dataset_shard("train")
            assert isinstance(data_shard, DataIterator), data_shard
            for k, v in expect_sizes.items():
                shard = train.get_dataset_shard(k)
                if v == -1:
                    assert shard is None, shard
                else:
                    count = 0
                    for batch in shard.iter_batches():
                        for arr in batch.values():
                            count += arr.size
                    assert count == v, shard

        kwargs.pop("scaling_config", None)
        super().__init__(
            train_loop_per_worker=train_loop_per_worker,
            scaling_config=ScalingConfig(num_workers=num_workers),
            **kwargs,
        )


def test_basic(ray_start_4_cpus):
    ds = ray.data.range(10)

    # Single worker basic case.
    test = TestBasic(
        1,
        True,
        {"train": 10, "test": 10},
        datasets={"train": ds, "test": ds},
    )
    test.fit()

    # Single worker, no test ds.
    test = TestBasic(1, True, {"train": 10, "test": -1}, datasets={"train": ds})
    test.fit()

    # Two workers, train and test split.
    test = TestBasic(
        2, True, {"train": 5, "test": 5}, datasets={"train": ds, "test": ds}
    )
    test.fit()

    # Two workers, both split.
    test = TestBasic(
        2,
        True,
        {"train": 5, "test": 5},
        dataset_config=DataConfig(datasets_to_split=["train", "test"]),
        datasets={"train": ds, "test": ds},
    )
    # Test get config.
    assert isinstance(test.get_dataset_config(), DataConfig)
    test.fit()


def test_split(ray_start_4_cpus):
    ds = ray.data.range(10)

    # Split all by default
    test = TestBasic(
        2,
        True,
        {"train": 5, "test": 5, "val": 5},
        datasets={"train": ds, "test": ds, "val": ds},
    )
    test.fit()

    # Test flag "all"
    test = TestBasic(
        2,
        True,
        {"train": 5, "test": 5},
        datasets={"train": ds, "test": ds},
        dataset_config=DataConfig(datasets_to_split="all"),
    )

    # Test split train only.
    test = TestBasic(
        2,
        True,
        {"train": 5, "test": 10},
        datasets={"train": ds, "test": ds},
        dataset_config=DataConfig(datasets_to_split=["train"]),
    )
    test.fit()

    # Test invalid arguments
    for datasets_to_split in ["train", ("train"), {}]:
        with pytest.raises(TypeError, match="`datasets_to_split` should be.*"):
            test = TestBasic(
                2,
                True,
                {"train": 5, "test": 10},
                datasets={"train": ds, "test": ds},
                dataset_config=DataConfig(datasets_to_split=datasets_to_split),
            )

    # Test empty `datasets_to_split` list
    test = TestBasic(
        2,
        True,
        {"train": 10, "test": 10},
        datasets={"train": ds, "test": ds},
        dataset_config=DataConfig(datasets_to_split=[]),
    )
    test.fit()


def test_configure_execution_options_carryover_context(ray_start_4_cpus):
    """Tests that execution options in DataContext are carried over to DatConfig
    automatically."""

    ctx = ray.data.DataContext.get_current()
    ctx.execution_options.preserve_order = True
    ctx.execution_options.verbose_progress = True

    data_config = DataConfig()

    ingest_options = data_config.default_ingest_options()
    assert ingest_options.preserve_order is True
    assert ingest_options.verbose_progress is True


@pytest.mark.parametrize("enable_locality", [True, False])
def test_configure_locality(enable_locality):
    data_config = DataConfig(enable_shard_locality=enable_locality)

    mock_ds = MagicMock()
    mock_ds.streaming_split = MagicMock()
    mock_ds.copy = MagicMock(return_value=mock_ds)
    world_size = 2
    worker_handles = [MagicMock() for _ in range(world_size)]
    worker_node_ids = ["node" + str(i) for i in range(world_size)]
    data_config.configure(
        datasets={"train": mock_ds},
        world_size=world_size,
        worker_handles=worker_handles,
        worker_node_ids=worker_node_ids,
    )
    mock_ds.streaming_split.assert_called_once()
    mock_ds.streaming_split.assert_called_with(
        world_size,
        equal=True,
        locality_hints=worker_node_ids if enable_locality else None,
    )


class CustomConfig(DataConfig):
    def __init__(self):
        pass

    def configure(self, *args, **kwargs):
        ds = ray.data.range(10)
        return [
            {"train": ds.iterator()},
            {"train": ds.iterator()},
        ]


def test_custom_config_subclass(ray_start_4_cpus):
    test = TestBasic(
        1,
        True,
        {"train": 10},
        dataset_config=CustomConfig(),
    )
    test.fit()


class TestRandom(DataParallelTrainer):
    def __init__(self, num_workers: int, expect_random: bool, **kwargs):
        def train_loop_per_worker():
            data_shard = train.get_dataset_shard("train")
            assert isinstance(data_shard, DataIterator), data_shard
            epoch1 = list(data_shard.iter_rows())
            epoch2 = list(data_shard.iter_rows())
            print("Epochs", epoch1, "\n", epoch2)
            if expect_random:
                assert epoch1 != epoch2
            else:
                assert epoch1 == epoch2

        kwargs.pop("scaling_config", None)
        super().__init__(
            train_loop_per_worker=train_loop_per_worker,
            scaling_config=ScalingConfig(num_workers=num_workers),
            **kwargs,
        )


def test_per_epoch_preprocessing(ray_start_4_cpus):
    ds = ray.data.range(100, override_num_blocks=100).randomize_block_order()
    test = TestRandom(2, True, datasets={"train": ds})
    test.fit()

    ds = ray.data.range(100, override_num_blocks=100).random_shuffle()
    test = TestRandom(2, True, datasets={"train": ds})
    test.fit()

    ds = ray.data.range(100, override_num_blocks=100).map(
        lambda x: {"id": x["id"] * random.random()}
    )
    test = TestRandom(2, True, datasets={"train": ds})
    test.fit()


def test_materialized_preprocessing(ray_start_4_cpus):
    # TODO(ekl) we should test all these configs with splitting enabled, but this
    # requires implementing deterministic streaming split.
    ds = ray.data.range(100, override_num_blocks=100).randomize_block_order()
    ds = ds.materialize()
    test = TestRandom(
        2,
        False,
        datasets={"train": ds},
        dataset_config=DataConfig(datasets_to_split=[]),
    )
    test.fit()

    ds = ray.data.range(100, override_num_blocks=100).random_shuffle()
    ds = ds.materialize()
    test = TestRandom(
        2,
        False,
        datasets={"train": ds},
        dataset_config=DataConfig(datasets_to_split=[]),
    )
    test.fit()

    ds = ray.data.range(100, override_num_blocks=100).map(
        lambda x: {"id": x["id"] * random.random()}
    )
    ds = ds.materialize()
    test = TestRandom(
        2,
        False,
        datasets={"train": ds},
        dataset_config=DataConfig(datasets_to_split=[]),
    )
    test.fit()


def _run_data_config_resource_test(data_config):
    cluster_cpus, cluster_gpus = 20, 10
    num_workers = 2
    # Resources used by training workers.
    cpus_per_worker, gpus_per_worker = 2, 1

    original_execution_options = data_config._get_execution_options("train")

    ray.init(num_cpus=cluster_cpus, num_gpus=cluster_gpus)

    class MyTrainer(DataParallelTrainer):
        def __init__(self, **kwargs):
            def train_loop_fn():
                train_ds = train.get_dataset_shard("train")
                new_execution_options = train_ds.get_context().execution_options
                if original_execution_options.is_resource_limits_default():
                    # If the original resource limits are default, the new resource
                    # limits should be the default as well.
                    assert new_execution_options.is_resource_limits_default()
                    exclude_resources = new_execution_options.exclude_resources
                    assert (
                        exclude_resources.cpu
                        == original_execution_options.exclude_resources.cpu
                        + cpus_per_worker * num_workers
                        + 1  # trainer coordinator
                    )
                    assert (
                        exclude_resources.gpu
                        == original_execution_options.exclude_resources.gpu
                        + gpus_per_worker * num_workers
                    )
                else:
                    # If the original resource limits are not default, the new resource
                    # limits should be the same as the original ones.
                    # And the new exclude_resources should be zero.
                    assert (
                        new_execution_options.resource_limits
                        == original_execution_options.resource_limits
                    )
                    assert (
                        new_execution_options.exclude_resources
                        == ExecutionResources.zero()
                    )

            kwargs.pop("scaling_config", None)

            super().__init__(
                train_loop_per_worker=train_loop_fn,
                scaling_config=ScalingConfig(
                    num_workers=num_workers,
                    use_gpu=True,
                    resources_per_worker={
                        "CPU": cpus_per_worker,
                        "GPU": gpus_per_worker,
                    },
                ),
                datasets={"train": ray.data.range(10)},
                dataset_config=data_config,
                **kwargs,
            )

    trainer = MyTrainer()
    trainer.fit()


def test_data_config_default_resource_limits(shutdown_only):
    """Test that DataConfig preserves user-configured exclude_resources."""
    execution_options = ExecutionOptions()
    execution_options.exclude_resources = execution_options.exclude_resources.copy(
        cpu=2, gpu=1
    )
    data_config = DataConfig(execution_options=execution_options)

    _run_data_config_resource_test(data_config)


def test_data_config_manual_resource_limits(shutdown_only):
    """Test manually setting resource limits in DataConfig."""
    execution_options = ExecutionOptions()
    execution_options.resource_limits = execution_options.resource_limits.copy(
        cpu=10, gpu=5
    )
    data_config = DataConfig(execution_options=execution_options)

    _run_data_config_resource_test(data_config)


def test_v1_train_with_v2_data_autoscaler_sets_exclude_resources(
    shutdown_only, monkeypatch
):
    """Regression test for the Train V1 + V2 cluster autoscaler combination."""
    monkeypatch.setenv("RAY_DATA_CLUSTER_AUTOSCALER", "V2")

    ray.init(num_cpus=10, num_gpus=2)

    num_train_cpus, num_train_gpus = 4.0, 2.0
    data_config = DataConfig()
    data_config.set_train_total_resources(
        num_train_cpus=num_train_cpus, num_train_gpus=num_train_gpus
    )

    iterators = data_config.configure(
        datasets={"train": ray.data.range(10)},
        world_size=2,
        worker_handles=None,
        worker_node_ids=None,
    )

    exclude_resources = (
        iterators[0]["train"].get_context().execution_options.exclude_resources
    )
    assert exclude_resources.cpu == num_train_cpus
    assert exclude_resources.gpu == num_train_gpus


if __name__ == "__main__":
    import sys

    import pytest

    sys.exit(pytest.main(["-v", "-x", __file__]))