import deepspeed

from unit.common import DistributedTest
from unit.simple_model import *

import pytest


class TestSparseCheckpoint(DistributedTest):
    world_size = 2

    @pytest.mark.parametrize(["to_save_model_has_embedding",
                              "to_save_model_sparse"],
                             [
                                 [False,
                                  False],
                                 [True,
                                  False],
                                 [True,
                                  True],
                             ])
    @pytest.mark.parametrize(["destination_has_embedding",
                              "destination_sparse"],
                             [
                                 [False,
                                  False],
                                 [True,
                                  False],
                                 [True,
                                  True],
                             ])
    def test_non_strict_load_sparse(self,
                                    tmpdir,
                                    to_save_model_has_embedding,
                                    to_save_model_sparse,
                                    destination_has_embedding,
                                    destination_sparse):
        class ModelNoEmbedding(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.linear = torch.nn.Linear(3, 1)

            def forward(self, x):
                return self.linear(x)

        class ModelEmbedding(torch.nn.Module):
            def __init__(self):
                super().__init__()
                self.emb = torch.nn.Embedding(10, 3)
                self.linear = torch.nn.Linear(3, 1)

            def forward(self, x, offsets):
                return self.linear(self.emb(x, offsets))

        if to_save_model_has_embedding:
            model_to_save = ModelEmbedding()
        else:
            model_to_save = ModelNoEmbedding()
        if destination_has_embedding:
            model_destination = ModelEmbedding()
        else:
            model_destination = ModelNoEmbedding()

        engine_to_save, _, _, _ = deepspeed.initialize(
            model=model_to_save, config={"train_batch_size": 2, "sparse_gradients": to_save_model_sparse}
        )
        engine_destination, _, _, _ = deepspeed.initialize(
            model=model_destination, config={"train_batch_size": 2, "sparse_gradients": destination_sparse}
        )

        save_folder = os.path.join(tmpdir, 'saved_checkpoint')
        save_tag = '1'

        engine_to_save.save_checkpoint(save_folder, tag=save_tag)

        is_sparse_destination = isinstance(model_destination,
                                           ModelEmbedding) and destination_sparse
        if isinstance(model_destination,
                      ModelEmbedding) and model_destination.emb.sparse:
            assert "emb.weight" in engine_destination.sparse_tensor_module_names
        engine_destination.load_checkpoint(save_folder,
                                           tag=save_tag,
                                           load_module_strict=False,
                                           load_optimizer_states=False,
                                           load_lr_scheduler_states=False,
                                           load_module_only=False)
        if isinstance(model_destination,
                      ModelEmbedding) and isinstance(model_to_save,
                                                     ModelEmbedding):
            assert engine_destination.sparse_tensor_module_names == engine_to_save.sparse_tensor_module_names
        elif isinstance(model_destination, ModelEmbedding):
            assert not is_sparse_destination or "emb.weight" in engine_destination.sparse_tensor_module_names
        else:
            assert len(engine_destination.sparse_tensor_module_names) == 0
