import pytest
from pydantic import ValidationError

from scaling.core.topology import Topology, TopologyConfig
from tests.core.utils import assert_nested_dicts_equal


@pytest.mark.short
@pytest.mark.parametrize(
    "config_input,expected_config,raises_error",
    [
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 1,
                "model_parallel_size": 1,
                "pipe_parallel_size": 1,
                "global_batch_size": 16,
                "gradient_accumulation_steps": 1,
            },
            {
                "global_rank": 0,
                "world_size": 1,
                "model_parallel_size": 1,
                "pipe_parallel_size": 1,
                "data_parallel_size": 1,
                "global_batch_size": 16,
                "micro_batch_size": 16,
                "gradient_accumulation_steps": 1,
                "local_slot": None,
                "pipe_partition_method": "uniform",
                "pipe_partition_overwrite": None,
                "activation_checkpointing_type": "disabled",
                "sequence_parallel": False,
            },
            False,
            id="data_parallel_size not set, all other parallelization parameters set to 1",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "global_batch_size": 32,
                "gradient_accumulation_steps": 2,
            },
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "data_parallel_size": 2,
                "global_batch_size": 32,
                "micro_batch_size": 8,
                "gradient_accumulation_steps": 2,
                "local_slot": None,
                "pipe_partition_method": "uniform",
                "pipe_partition_overwrite": None,
                "activation_checkpointing_type": "disabled",
                "sequence_parallel": False,
            },
            False,
            id="data_parallel_size and micro_batch_size need to be inferred",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": None,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "global_batch_size": 32,
                "gradient_accumulation_steps": 2,
            },
            {},
            True,
            id="Both world_size and data_parallel_size not set leads to validation error",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "global_batch_size": 33,
                "gradient_accumulation_steps": 2,
            },
            {},
            True,
            id="global_batch_size not divisible by product of data_parallel_size and gradient_accumulation_steps "
            "raises error",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "micro_batch_size": 9,
                "gradient_accumulation_steps": 2,
            },
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "data_parallel_size": 2,
                "global_batch_size": 36,
                "micro_batch_size": 9,
                "gradient_accumulation_steps": 2,
                "local_slot": None,
                "pipe_partition_method": "uniform",
                "pipe_partition_overwrite": None,
                "activation_checkpointing_type": "disabled",
                "sequence_parallel": False,
            },
            False,
            id="data_parallel_size and global_batch_size need to be inferred",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "micro_batch_size": 9,
                "global_batch_size": 9,
                "gradient_accumulation_steps": 2,
            },
            {},
            True,
            id="incorrectly set batch size params lead to error",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "micro_batch_size": 9,
            },
            {},
            True,
            id="only 1 out of 3 batch parameters (micro_batch_size) given leads to error",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "global_batch_size": 9,
            },
            {},
            True,
            id="only 1 out of 3 batch parameters (global_batch_size) given leads to error",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "gradient_accumulation_steps": 2,
            },
            {},
            True,
            id="only 1 out of 3 batch parameters (gradient_accumulation_steps) given leads to error",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 24,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "data_parallel_size": 3,
                "gradient_accumulation_steps": 2,
                "global_batch_size": 16,
            },
            {},
            True,
            id="global_batch_size not modulo divisible by data_parallel_size and gradient_accumulation_steps",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 4,
                "data_parallel_size": 4,
                "micro_batch_size": 9,
                "global_batch_size": 9,
                "gradient_accumulation_steps": 2,
            },
            {},
            True,
            id="incorrectly set parallelization params lead to error",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "data_parallel_size": 4,
                "micro_batch_size": 9,
                "gradient_accumulation_steps": 2,
            },
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 2,
                "data_parallel_size": 4,
                "global_batch_size": 72,
                "micro_batch_size": 9,
                "gradient_accumulation_steps": 2,
                "local_slot": None,
                "pipe_partition_method": "uniform",
                "pipe_partition_overwrite": None,
                "activation_checkpointing_type": "disabled",
                "sequence_parallel": False,
            },
            False,
            id="pipe_parallel_size correctly inferred",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": 16,
                "pipe_parallel_size": 2,
                "data_parallel_size": 4,
                "micro_batch_size": 9,
                "gradient_accumulation_steps": 2,
            },
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 2,
                "data_parallel_size": 4,
                "global_batch_size": 72,
                "micro_batch_size": 9,
                "gradient_accumulation_steps": 2,
                "local_slot": None,
                "pipe_partition_method": "uniform",
                "pipe_partition_overwrite": None,
                "activation_checkpointing_type": "disabled",
                "sequence_parallel": False,
            },
            False,
            id="model_parallel_size correctly inferred",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "world_size": None,
                "model_parallel_size": 2,
                "pipe_parallel_size": 2,
                "data_parallel_size": 4,
                "micro_batch_size": 9,
                "gradient_accumulation_steps": 2,
            },
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 2,
                "data_parallel_size": 4,
                "global_batch_size": 72,
                "micro_batch_size": 9,
                "gradient_accumulation_steps": 2,
                "local_slot": None,
                "pipe_partition_method": "uniform",
                "pipe_partition_overwrite": None,
                "activation_checkpointing_type": "disabled",
                "sequence_parallel": False,
            },
            False,
            id="world_size correctly inferred when set as None",
        ),
        pytest.param(
            {
                "global_rank": 0,
                "model_parallel_size": 2,
                "pipe_parallel_size": 2,
                "data_parallel_size": 4,
                "micro_batch_size": 9,
                "gradient_accumulation_steps": 2,
            },
            {
                "global_rank": 0,
                "world_size": 16,
                "model_parallel_size": 2,
                "pipe_parallel_size": 2,
                "data_parallel_size": 4,
                "global_batch_size": 72,
                "micro_batch_size": 9,
                "gradient_accumulation_steps": 2,
                "local_slot": None,
                "pipe_partition_method": "uniform",
                "pipe_partition_overwrite": None,
                "activation_checkpointing_type": "disabled",
                "sequence_parallel": False,
            },
            False,
            id="world_size correctly inferred when not set",
        ),
    ],
)
def test_topology(config_input: dict, expected_config: dict, raises_error: bool):
    if raises_error:
        with pytest.raises(ValidationError):
            _ = TopologyConfig(**config_input)
    else:
        topology_config = TopologyConfig(**config_input)
        assert_nested_dicts_equal(topology_config.as_dict(), expected_config)

        _ = Topology(config=topology_config)
