# Owner(s): ["oncall: distributed checkpointing"]
import os
import shutil
import tempfile
from typing import Any

import torch
from torch.distributed.checkpoint._experimental.checkpoint_reader import (
    CheckpointReader,
)
from torch.distributed.checkpoint._experimental.types import RankInfo
from torch.testing._internal.common_utils import run_tests, TestCase


class TestCheckpointReader(TestCase):
    def setUp(self):
        super().setUp()
        # Create a temporary directory for test checkpoints
        self.temp_dir = tempfile.mkdtemp()

        # Create test objects
        self.rank_info = RankInfo(
            global_rank=0,
            global_world_size=1,
        )

        # Create the checkpoint reader
        self.reader = CheckpointReader(
            rank_info=self.rank_info,
        )

        # Create a test state dictionary
        self.state_dict = {
            "model": {
                "weight": torch.randn(10, 5),
                "bias": torch.randn(5),
                "test_list": [torch.randn(2), torch.randn(2)],
            },
            "optimizer": {
                "param_groups": [
                    {"lr": 0.01, "test_list": [torch.randn(2), torch.randn(2)]}
                ]
            },
            "epoch": 5,
            "step": 1000,
        }

        # Create a test checkpoint file
        self.checkpoint_path = os.path.join(self.temp_dir, "checkpoint")
        os.makedirs(self.checkpoint_path, exist_ok=True)
        checkpoint_file = os.path.join(
            self.checkpoint_path, f"checkpoint_{self.rank_info.global_rank}.pt"
        )
        torch.save(self.state_dict, checkpoint_file)

    def move_tensors_to_device(self, state_dict: Any, device: str) -> Any:
        """
        Recursively move all tensors in a nested dictionary to CUDA.

        Args:
            state_dict (dict): A dictionary potentially containing nested dictionaries and tensors.

        Returns:
            dict: A new dictionary with all tensors moved to CUDA.
        """
        if isinstance(state_dict, dict):
            return {
                key: self.move_tensors_to_device(value, device)
                for key, value in state_dict.items()
            }
        elif isinstance(state_dict, list):
            return [self.move_tensors_to_device(item, device) for item in state_dict]
        elif isinstance(state_dict, torch.Tensor):
            return state_dict.cuda() if device == "cpu" else state_dict.cpu()
        else:
            return state_dict

    def deep_compare(self, obj1: Any, obj2: Any) -> bool:
        if isinstance(obj1, dict) and isinstance(obj2, dict):
            if obj1.keys() != obj2.keys():
                return False
            return all(self.deep_compare(obj1[key], obj2[key]) for key in obj1)
        elif isinstance(obj1, (list, tuple)) and isinstance(obj2, (list, tuple)):
            if len(obj1) != len(obj2):
                return False
            return all(
                self.deep_compare(item1, item2) for item1, item2 in zip(obj1, obj2)
            )
        elif isinstance(obj1, torch.Tensor) and isinstance(obj2, torch.Tensor):
            return torch.equal(obj1, obj2)
        else:
            return obj1 == obj2

    def tearDown(self):
        # Clean up the temporary directory
        shutil.rmtree(self.temp_dir)

    def test_read_checkpoint(self):
        """Test that read correctly reads a checkpoint file."""

        # Call read
        read_state_dict, missing_keys = self.reader.read(self.checkpoint_path)
        self.assertEqual(missing_keys, [])

        # Verify that the read state dictionary contains the expected values
        self.assertIn("model", read_state_dict)
        self.assertIn("optimizer", read_state_dict)
        self.assertTrue(self.deep_compare(read_state_dict, self.state_dict))

        # No hooks to verify since we removed them

    def test_read_with_map_location(self):
        """Test that read correctly uses the map_location parameter."""
        # Call read with map_location='cpu'
        map_location = "cuda" if torch.cuda.is_available() else "cpu"
        read_state_dict, _ = self.reader.read(
            self.checkpoint_path, map_location=map_location
        )

        # Verify that the read state dictionary contains the expected values
        self.assertIn("model", read_state_dict)
        self.assertIn("optimizer", read_state_dict)
        self.assertEqual(read_state_dict["epoch"], 5)
        self.assertEqual(read_state_dict["step"], 1000)
        self.assertEqual(read_state_dict["model"]["weight"].device.type, map_location)
        read_state_dict, _ = self.reader.read(
            self.checkpoint_path, map_location=map_location
        )

        # Verify that the read state dictionary contains the expected values
        self.assertIn("model", read_state_dict)
        self.assertIn("optimizer", read_state_dict)
        self.assertEqual(read_state_dict["epoch"], 5)
        self.assertEqual(read_state_dict["step"], 1000)
        self.assertEqual(read_state_dict["model"]["weight"].device.type, map_location)

    def test_read_nonexistent_checkpoint(self):
        """Test that read raises FileNotFoundError for a nonexistent checkpoint."""
        # Set up a path to a nonexistent checkpoint
        nonexistent_path = os.path.join(self.temp_dir, "nonexistent_checkpoint")

        # Call read and expect a FileNotFoundError
        with self.assertRaises(FileNotFoundError):
            self.reader.read(nonexistent_path)

    def test_read_with_kwargs(self):
        """Test that read correctly passes kwargs."""
        # Call read with additional kwargs
        kwargs = {"extra": "value"}
        self.reader.read(self.checkpoint_path, **kwargs)

    def test_partial_read(self):
        """Test that read with state_dict correctly loads only the requested keys."""
        # Create a partial state dictionary with only some keys
        partial_state_dict = {}
        partial_state_dict["optimizer"] = None
        partial_state_dict["model"] = {"weight": torch.randn(10, 5)}
        partial_state_dict["epoch"] = None
        # Call read with state_dict
        updated_state_dict, _ = self.reader.read(
            self.checkpoint_path,
            partial_state_dict,
        )

        # Verify that the updated state dictionary contains values from both dictionaries
        self.assertIn("model", updated_state_dict)
        self.assertIn("epoch", updated_state_dict)
        self.assertTrue(
            torch.equal(
                updated_state_dict["model"]["weight"],
                self.state_dict["model"]["weight"],
            )
        )

        self.assertTrue(
            self.deep_compare(
                updated_state_dict["optimizer"], self.state_dict["optimizer"]
            )
        )
        self.assertEqual(updated_state_dict["epoch"], 5)  # From checkpoint

        self.assertNotIn("bias", updated_state_dict["model"])
        self.assertNotIn("step", updated_state_dict)

    def test_partial_read_missing_keys(self):
        """Test that partial_read correctly reports missing keys."""
        # Create a partial state dictionary with keys that don't exist in the checkpoint
        partial_state_dict = {
            "model": None,
            "nonexistent_key": None,  # This key doesn't exist in the checkpoint
            "another_missing_key": {"nested": None},  # This key also doesn't exist
        }

        # Call read with state_dict
        _, missing_keys = self.reader.read(
            self.checkpoint_path,
            partial_state_dict,
        )

        # Verify that missing keys are correctly reported
        self.assertIn("nonexistent_key", missing_keys)
        self.assertIn("another_missing_key", missing_keys)

        # Verify that keys that exist in the checkpoint are not in missing_keys
        self.assertNotIn("model", missing_keys)

    def test_partial_read_different_dtypes(self):
        """Test that partial_read correctly handles different tensor dtypes."""
        # Create a state dictionary with tensors of different dtypes
        dtype_state_dict = {
            "float32": torch.randn(10, 10, dtype=torch.float32),
            "float64": torch.randn(10, 10, dtype=torch.float64),
            "int32": torch.randint(-100, 100, (10, 10), dtype=torch.int32),
            "int64": torch.randint(-100, 100, (10, 10), dtype=torch.int64),
            "bool": torch.randint(0, 2, (10, 10), dtype=torch.bool),
        }

        # Save the state dictionary
        dtype_checkpoint_path = os.path.join(self.temp_dir, "dtype_checkpoint")
        os.makedirs(dtype_checkpoint_path, exist_ok=True)
        checkpoint_file = os.path.join(
            dtype_checkpoint_path, f"checkpoint_{self.rank_info.global_rank}.pt"
        )
        torch.save(dtype_state_dict, checkpoint_file)

        # Create a partial state dictionary requesting tensors of each dtype
        partial_state_dict = {
            "float32": torch.randn(10, 10, dtype=torch.float32),
            "float64": None,
            "int32": None,
            "int64": None,
            "bool": None,
        }

        # Load the partial state dictionary
        updated_state_dict, _ = self.reader.read(
            os.path.dirname(checkpoint_file),
            partial_state_dict,
        )

        # Verify that tensors of each dtype were loaded correctly
        for key in dtype_state_dict:
            self.assertIn(key, updated_state_dict)
            self.assertEqual(updated_state_dict[key].dtype, dtype_state_dict[key].dtype)
            self.assertTrue(
                torch.allclose(updated_state_dict[key], dtype_state_dict[key])
            )


if __name__ == "__main__":
    run_tests()
