# Owner(s): ["module: unknown"]

import itertools
import logging
import re

import torch
from torch import nn
from torch.ao.pruning import (
    BaseSparsifier,
    FakeSparsity,
    NearlyDiagonalSparsifier,
    WeightNormSparsifier,
)
from torch.nn.utils.parametrize import is_parametrized
from torch.testing._internal.common_pruning import (
    ImplementedSparsifier,
    MockSparseLinear,
    SimpleLinear,
)
from torch.testing._internal.common_utils import TestCase


logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
)


class TestBaseSparsifier(TestCase):
    def test_constructor(self):
        # Cannot instantiate the abstract base
        self.assertRaises(TypeError, BaseSparsifier)
        # Can instantiate the model with no configs
        model = SimpleLinear()
        sparsifier = ImplementedSparsifier(test=3)
        sparsifier.prepare(model, config=None)
        assert len(sparsifier.groups) == 5
        sparsifier.step()
        # Can instantiate the model with configs
        sparsifier = ImplementedSparsifier(test=3)
        sparsifier.prepare(model, [{"tensor_fqn": "linear1.weight"}])
        assert len(sparsifier.groups) == 1
        assert sparsifier.groups[0]["tensor_fqn"] == "linear1.weight"
        assert "test" in sparsifier.groups[0]
        assert sparsifier.groups[0]["test"] == 3

    def test_prepare_config(self):
        model = SimpleLinear()
        sparsifier = ImplementedSparsifier(test=3)
        # Make sure there are no parametrizations before `prepare`
        assert not hasattr(model.seq[0], "parametrizations")
        assert not hasattr(model.linear1, "parametrizations")
        assert not hasattr(model.linear2, "parametrizations")
        sparsifier.prepare(
            model,
            config=[
                {"tensor_fqn": "seq.0.weight", "test": 42},
                # No 'linear1' to make sure it will be skipped in the sparsification
                {"tensor_fqn": "linear2.weight"},
            ],
        )
        assert len(sparsifier.groups) == 2
        # Check if default argument is not assigned if explicit
        assert sparsifier.groups[0]["tensor_fqn"] == "seq.0.weight"
        assert sparsifier.groups[0]["test"] == 42
        # Check if FQN and module are pointing to the same location
        assert sparsifier.groups[1]["tensor_fqn"] == "linear2.weight"
        assert sparsifier.groups[1]["module"] == model.linear2
        # Check if parameterizations are attached
        assert hasattr(model.seq[0], "parametrizations")
        assert not hasattr(model.linear1, "parametrizations")
        assert hasattr(model.linear2, "parametrizations")

    def test_step(self):
        model = SimpleLinear()
        sparsifier = ImplementedSparsifier(test=3)
        sparsifier.enable_mask_update = True
        sparsifier.prepare(model, [{"tensor_fqn": "linear1.weight"}])
        sparsifier.step()
        assert torch.all(model.linear1.parametrizations.weight[0].mask[0] == 0)

    def test_state_dict(self):
        step_count = 3
        model0 = SimpleLinear()
        sparsifier0 = ImplementedSparsifier(test=3)
        sparsifier0.prepare(model0, [{"tensor_fqn": "linear1.weight"}])
        mask = model0.linear1.parametrizations["weight"][0].mask
        mask.data = torch.arange(mask.shape[0] * mask.shape[1]).reshape(mask.shape)
        for step in range(step_count):
            sparsifier0.step()
        state_dict = sparsifier0.state_dict()

        # Check the expected keys in the state_dict
        assert "state" in state_dict
        assert "step_count" in state_dict["state"]["linear1.weight"]
        assert state_dict["state"]["linear1.weight"]["step_count"] == 3
        assert "groups" in state_dict
        assert "test" in state_dict["groups"][0]
        assert "tensor_fqn" in state_dict["groups"][0]
        assert state_dict["groups"][0]["tensor_fqn"] == "linear1.weight"

        # Check loading static_dict creates an equivalent model
        model1 = SimpleLinear()
        sparsifier1 = ImplementedSparsifier()
        sparsifier1.prepare(model1, None)

        assert sparsifier0.state != sparsifier1.state

        # Make sure the masks are different in the beginning
        for mg in sparsifier0.groups:
            if mg["tensor_fqn"] == "linear1.weight":
                mask0 = mg["module"].parametrizations.weight[0].mask
        for mg in sparsifier1.groups:
            if mg["tensor_fqn"] == "linear1.weight":
                mask1 = mg["module"].parametrizations.weight[0].mask
        self.assertNotEqual(mask0, mask1)

        sparsifier1.load_state_dict(state_dict)

        # Make sure the states are loaded, and are correct
        assert sparsifier0.state == sparsifier1.state

        # Make sure the masks (and all dicts) are the same after loading
        assert len(sparsifier0.groups) == len(sparsifier1.groups)
        for idx in range(len(sparsifier0.groups)):
            mg0 = sparsifier0.groups[idx]
            mg1 = sparsifier1.groups[idx]
            for key in mg0.keys():
                assert key in mg1
                if key == "module":
                    # We cannot compare modules as they are different
                    param0 = mg0[key].parametrizations.weight[0]
                    param1 = mg1[key].parametrizations.weight[0]
                    assert hasattr(param0, "mask")
                    assert hasattr(param1, "mask")
                    self.assertEqual(param0.__dict__, param1.__dict__)
                else:
                    assert mg0[key] == mg1[key]

    def test_convert(self):
        model = SimpleLinear()
        sparsifier = ImplementedSparsifier(test=3)
        sparsifier.prepare(model, [{"tensor_fqn": "linear1.weight"}])
        new_model = sparsifier.convert(
            model, mapping={nn.Linear: MockSparseLinear}, inplace=False
        )

        assert isinstance(new_model.linear1, MockSparseLinear)
        assert isinstance(new_model.seq[0], nn.Linear)
        assert isinstance(new_model.linear2, nn.Linear)

    def test_mask_squash(self):
        model = SimpleLinear()
        sparsifier = ImplementedSparsifier(test=3)
        sparsifier.prepare(model, [{"tensor_fqn": "linear1.weight"}])
        assert hasattr(model.linear1.parametrizations.weight[0], "mask")
        assert is_parametrized(model.linear1, "weight")
        assert not is_parametrized(model.seq[0], "weight")

        sparsifier.squash_mask()
        assert not is_parametrized(model.seq[0], "weight")
        assert not is_parametrized(model.linear1, "weight")

    def test_mask_squash_with_params1(self):
        model = SimpleLinear()
        sparsifier = ImplementedSparsifier(foo=3, bar=2, baz=1)
        sparsifier.prepare(
            model, [{"tensor_fqn": "linear1.weight"}, {"tensor_fqn": "seq.0.weight"}]
        )
        sparsifier.squash_mask(
            params_to_keep_per_layer={"linear1": ("foo", "bar"), "seq.0": ("baz",)}
        )
        assert not is_parametrized(model.seq[0], "weight")
        assert not is_parametrized(model.linear1, "weight")
        assert hasattr(model.seq[0], "sparse_params")
        assert hasattr(model.linear1, "sparse_params")
        assert model.seq[0].sparse_params.get("foo", None) is None
        assert model.seq[0].sparse_params.get("bar", None) is None
        assert model.seq[0].sparse_params.get("baz", None) == 1
        assert model.linear1.sparse_params.get("foo", None) == 3
        assert model.linear1.sparse_params.get("bar", None) == 2
        assert model.linear1.sparse_params.get("baz", None) is None

    def test_mask_squash_with_params2(self):
        model = SimpleLinear()
        sparsifier = ImplementedSparsifier(foo=3, bar=2, baz=1)
        sparsifier.prepare(
            model, [{"tensor_fqn": "linear1.weight"}, {"tensor_fqn": "seq.0.weight"}]
        )
        sparsifier.squash_mask(params_to_keep=("foo", "bar"))
        assert not is_parametrized(model.seq[0], "weight")
        assert not is_parametrized(model.linear1, "weight")
        assert hasattr(model.seq[0], "sparse_params")
        assert hasattr(model.linear1, "sparse_params")
        assert model.seq[0].sparse_params.get("foo", None) == 3
        assert model.seq[0].sparse_params.get("bar", None) == 2
        assert model.seq[0].sparse_params.get("baz", None) is None
        assert model.linear1.sparse_params.get("foo", None) == 3
        assert model.linear1.sparse_params.get("bar", None) == 2
        assert model.linear1.sparse_params.get("baz", None) is None

    def test_mask_squash_with_params3(self):
        model = SimpleLinear()
        sparsifier = ImplementedSparsifier(foo=3, bar=2, baz=1)
        sparsifier.prepare(
            model, [{"tensor_fqn": "linear1.weight"}, {"tensor_fqn": "seq.0.weight"}]
        )
        sparsifier.squash_mask(
            params_to_keep=("foo", "bar"), params_to_keep_per_layer={"seq.0": ("baz",)}
        )
        assert not is_parametrized(model.seq[0], "weight")
        assert not is_parametrized(model.linear1, "weight")
        assert hasattr(model.seq[0], "sparse_params")
        assert hasattr(model.linear1, "sparse_params")
        assert model.seq[0].sparse_params.get("foo", None) == 3
        assert model.seq[0].sparse_params.get("bar", None) == 2
        assert model.seq[0].sparse_params.get("baz", None) == 1
        assert model.linear1.sparse_params.get("foo", None) == 3
        assert model.linear1.sparse_params.get("bar", None) == 2
        assert model.linear1.sparse_params.get("baz", None) is None


class TestWeightNormSparsifier(TestCase):
    def test_constructor(self):
        model = SimpleLinear()
        sparsifier = WeightNormSparsifier()
        sparsifier.prepare(model, config=None)
        for g in sparsifier.groups:
            assert isinstance(g["module"], nn.Linear)
            # The groups are unordered
            assert g["module_fqn"] in ("seq.0", "seq.1", "seq.2", "linear1", "linear2")

    def test_step(self):
        model = SimpleLinear()
        sparsifier = WeightNormSparsifier(sparsity_level=0.5)
        sparsifier.prepare(model, config=[{"tensor_fqn": "linear1.weight"}])
        for g in sparsifier.groups:
            # Before step
            module = g["module"]
            assert (
                1.0 - module.parametrizations["weight"][0].mask.mean()
            ) == 0  # checking sparsity level is 0
        sparsifier.enable_mask_update = True
        sparsifier.step()
        self.assertAlmostEqual(
            model.linear1.parametrizations["weight"][0].mask.mean().item(),
            0.5,
            places=2,
        )
        for g in sparsifier.groups:
            # After step
            module = g["module"]
            assert (
                1.0 - module.parametrizations["weight"][0].mask.mean()
            ) > 0  # checking sparsity level has increased
        # Test if the mask collapses to all zeros if the weights are randomized
        iters_before_collapse = 1000
        for _ in range(iters_before_collapse):
            model.linear1.weight.data = torch.randn(model.linear1.weight.shape)
            sparsifier.step()
        for g in sparsifier.groups:
            # After step
            module = g["module"]
            assert (
                1.0 - module.parametrizations["weight"][0].mask.mean()
            ) > 0  # checking sparsity level did not collapse

    def test_step_2_of_4(self):
        model = SimpleLinear()
        sparsifier = WeightNormSparsifier(
            sparsity_level=1.0, sparse_block_shape=(1, 4), zeros_per_block=2
        )
        sparsifier.prepare(model, config=[{"tensor_fqn": "linear1.weight"}])
        sparsifier.step()
        # make sure the sparsity level is approximately 50%
        mask = model.linear1.parametrizations["weight"][0].mask.to(
            torch.float
        )  # mean works on float only
        self.assertAlmostEqual(mask.mean().item(), 0.5, places=2)
        # Make sure each block has exactly 50% zeros
        module = sparsifier.groups[0]["module"]
        mask = module.parametrizations["weight"][0].mask
        for row in mask:
            for idx in range(0, len(row), 4):
                block = row[idx : idx + 4]
                block, _ = block.sort()
                assert (block[:2] == 0).all()
                assert (block[2:] != 0).all()

    def test_prepare(self):
        model = SimpleLinear()
        sparsifier = WeightNormSparsifier()
        sparsifier.prepare(model, config=None)
        for g in sparsifier.groups:
            module = g["module"]
            # Check mask exists
            assert hasattr(module.parametrizations["weight"][0], "mask")
            # Check parametrization exists and is correct
            assert is_parametrized(module, "weight")
            assert type(module.parametrizations.weight[0]) == FakeSparsity

    def test_mask_squash(self):
        model = SimpleLinear()
        sparsifier = WeightNormSparsifier()
        sparsifier.prepare(model, config=None)
        sparsifier.squash_mask()
        for g in sparsifier.groups:
            module = g["module"]
            assert not is_parametrized(module, "weight")
            assert not hasattr(module, "mask")

    def test_sparsity_levels(self):
        sparsity_levels = [-1.0, 0.0, 0.5, 1.0, 2.0]
        sparse_block_shapes = [(1, 1), (1, 4), (2, 2), (4, 1)]
        zeros_per_blocks = [0, 1, 2, 3, 4]

        testcases = itertools.tee(
            itertools.product(sparsity_levels, sparse_block_shapes, zeros_per_blocks)
        )
        # Create a config and model with all the testcases
        model = nn.Sequential()
        sparsifier = WeightNormSparsifier()

        sparsity_per_layer_config = []
        p = re.compile(r"[-\.\s]")
        for sl, sbs, zpb in testcases[0]:
            # Make sure the number of zeros is not > values in a block
            if zpb > sbs[0] * sbs[1]:
                continue
            layer_name = f"{sl}_{sbs}_{zpb}"
            layer_name = p.sub("_", layer_name)

            layer = nn.Linear(12, 12, bias=False)
            layer.weight = nn.Parameter(torch.ones(12, 12))
            model.add_module(layer_name, layer)
            config = {
                "tensor_fqn": layer_name + ".weight",
                "sparsity_level": sl,
                "sparse_block_shape": sbs,
                "zeros_per_block": zpb,
            }
            sparsity_per_layer_config.append(config)

        sparsifier.prepare(model, sparsity_per_layer_config)
        sparsifier.step()
        sparsifier.squash_mask()
        model.eval()

        for sl, sbs, zpb in testcases[1]:
            if zpb > sbs[0] * sbs[1]:
                continue
            layer_name = f"{sl}_{sbs}_{zpb}"
            layer_name = p.sub("_", layer_name)
            layer = getattr(model, layer_name)

            # Level of sparsity is achieved
            sparse_mask = (layer.weight == 0).float()
            if zpb == 0:
                assert sparse_mask.mean() == 0
            else:
                # Ratio of individual zeros in the tensor
                true_sl = min(max(sl, 0.0), 1.0)
                true_sl = true_sl * zpb / sbs[0] / sbs[1]
                assert sparse_mask.mean() == true_sl


class TestNearlyDiagonalSparsifier(TestCase):
    def test_constructor(self):
        model = SimpleLinear()
        sparsifier = NearlyDiagonalSparsifier(nearliness=1)
        sparsifier.prepare(model, config=None)
        for g in sparsifier.groups:
            assert isinstance(g["module"], nn.Linear)
            # The groups are unordered
            assert g["module_fqn"] in ("seq.0", "seq.1", "seq.2", "linear1", "linear2")

    def test_step(self):
        model = SimpleLinear()
        sparsifier = NearlyDiagonalSparsifier(nearliness=1)
        sparsifier.prepare(model, config=[{"tensor_fqn": "linear1.weight"}])

        for g in sparsifier.groups:
            # Before step
            module = g["module"]
            assert (
                1.0 - module.parametrizations["weight"][0].mask.mean()
            ) == 0  # checking sparsity level is 0

        sparsifier.enable_mask_update = True
        sparsifier.step()
        mask = module.parametrizations["weight"][0].mask
        height, width = mask.shape
        assert torch.all(mask == torch.eye(height, width))

        for g in sparsifier.groups:
            # After step
            module = g["module"]
            assert (
                1.0 - module.parametrizations["weight"][0].mask.mean()
            ) > 0  # checking sparsity level has increased

        # Test if the mask collapses to all zeros if the weights are randomized
        iters_before_collapse = 1000
        for _ in range(iters_before_collapse):
            model.linear1.weight.data = torch.randn(model.linear1.weight.shape)
            sparsifier.step()
        for g in sparsifier.groups:
            # After step
            module = g["module"]
            assert (
                1.0 - module.parametrizations["weight"][0].mask.mean()
            ) > 0  # checking sparsity level did not collapse

    def test_prepare(self):
        model = SimpleLinear()
        sparsifier = NearlyDiagonalSparsifier(nearliness=1)
        sparsifier.prepare(model, config=None)
        for g in sparsifier.groups:
            module = g["module"]
            # Check mask exists
            assert hasattr(module.parametrizations["weight"][0], "mask")
            # Check parametrization exists and is correct
            assert is_parametrized(module, "weight")
            assert type(module.parametrizations.weight[0]) == FakeSparsity

    def test_mask_squash(self):
        model = SimpleLinear()
        sparsifier = NearlyDiagonalSparsifier(nearliness=1)
        sparsifier.prepare(model, config=None)
        sparsifier.step()
        sparsifier.squash_mask()
        for g in sparsifier.groups:
            module = g["module"]
            assert not is_parametrized(module, "weight")
            assert not hasattr(module, "mask")
            weights = module.weight
            height, width = weights.shape
            assert torch.all(
                weights == torch.eye(height, width) * weights
            )  # only diagonal to be present

    def test_sparsity_levels(self):
        nearliness_levels = list(range(-1, 100))
        model = nn.Sequential()

        p = re.compile(r"[-\.\s]")
        for nearliness in nearliness_levels:
            sparsifier = NearlyDiagonalSparsifier(nearliness=1)
            layer_name = f"{nearliness}"
            layer_name = p.sub("_", layer_name)

            layer = nn.Linear(32, 32, bias=False)
            layer.weight = nn.Parameter(torch.ones(32, 32))
            width, height = layer.weight.shape
            model.add_module(layer_name, layer)
            config = {"tensor_fqn": layer_name + ".weight", "nearliness": nearliness}

            sparsifier.prepare(model, [config])
            # should raise a ValueError when nearliness arg is illegal
            if (nearliness > 0 and nearliness % 2 == 0) or (
                nearliness // 2 >= min(width, height)
            ):
                with self.assertRaises(ValueError):
                    sparsifier.step()
            else:
                sparsifier.step()
                sparsifier.squash_mask()
                model.eval()

                layer = getattr(model, layer_name)
                # verify that mask created corresponds to the nearliness
                self._verify_nearliness(layer.weight, nearliness)

    # helper function to verify nearliness of a mask
    def _verify_nearliness(self, mask: torch.Tensor, nearliness: int):
        if nearliness <= 0:
            assert torch.all(mask == torch.zeros(mask.shape[0], mask.shape[1]))
        else:
            height, width = mask.shape
            dist_to_diagonal = nearliness // 2
            for row in range(0, height):
                for col in range(0, width):
                    if abs(row - col) <= dist_to_diagonal:
                        assert mask[row, col] == 1
                    else:
                        assert mask[row, col] == 0