# -*- coding: utf-8 -*-
# Owner(s): ["module: unknown"]

import itertools
import logging
import re

import torch
from torch import nn
from torch.ao.sparsity import BaseSparsifier, WeightNormSparsifier, FakeSparsity, NearlyDiagonalSparsifier
from torch.nn.utils.parametrize import is_parametrized

from torch.testing._internal.common_utils import TestCase

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

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.seq = nn.Sequential(
            nn.Linear(16, 16)
        )
        self.linear = nn.Linear(16, 16)
        self.head = nn.Linear(16, 4)

    def forward(self, x):
        x = self.seq(x)
        x = self.linear(x)
        x = self.head(x)
        return x


class ImplementedSparsifier(BaseSparsifier):
    def __init__(self, **kwargs):
        super().__init__(defaults=kwargs)

    def update_mask(self, module, **kwargs):
        module.parametrizations.weight[0].mask[0] = 0
        linear_state = self.state['linear.weight']
        linear_state['step_count'] = linear_state.get('step_count', 0) + 1


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

    def test_prepare_config(self):
        model = Model()
        sparsifier = ImplementedSparsifier(test=3)
        # Make sure there are no parametrizations before `prepare`
        assert not hasattr(model.seq[0], 'parametrizations')
        assert not hasattr(model.linear, 'parametrizations')
        assert not hasattr(model.head, 'parametrizations')
        sparsifier.prepare(model, config=[
            {'tensor_fqn': 'seq.0.weight', 'test': 42},
            # No 'linear' to make sure it will be skipped in the sparsification
            {'tensor_fqn': 'head.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'] == 'head.weight'
        assert sparsifier.groups[1]['module'] == model.head
        # Check if parameterizations are attached
        assert hasattr(model.seq[0], 'parametrizations')
        assert not hasattr(model.linear, 'parametrizations')
        assert hasattr(model.head, 'parametrizations')

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

    def test_state_dict(self):
        step_count = 3
        model0 = Model()
        sparsifier0 = ImplementedSparsifier(test=3)
        sparsifier0.prepare(model0, [{'tensor_fqn': 'linear.weight'}])
        mask = model0.linear.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']['linear.weight']
        assert state_dict['state']['linear.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'] == 'linear.weight'

        # Check loading static_dict creates an equivalent model
        model1 = Model()
        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'] == 'linear.weight':
                mask0 = mg['module'].parametrizations.weight[0].mask
        for mg in sparsifier1.groups:
            if mg['tensor_fqn'] == 'linear.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_mask_squash(self):
        model = Model()
        sparsifier = ImplementedSparsifier(test=3)
        sparsifier.prepare(model, [{'tensor_fqn': 'linear.weight'}])
        assert hasattr(model.linear.parametrizations.weight[0], 'mask')
        assert is_parametrized(model.linear, '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.linear, 'weight')

    def test_mask_squash_with_params1(self):
        model = Model()
        sparsifier = ImplementedSparsifier(foo=3, bar=2, baz=1)
        sparsifier.prepare(model, [{'tensor_fqn': 'linear.weight'}, {'tensor_fqn': 'seq.0.weight'}])
        sparsifier.squash_mask(
            params_to_keep_per_layer={
                'linear': ('foo', 'bar'),
                'seq.0': ('baz',)
            })
        assert not is_parametrized(model.seq[0], 'weight')
        assert not is_parametrized(model.linear, 'weight')
        assert hasattr(model.seq[0], 'sparse_params')
        assert hasattr(model.linear, '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.linear.sparse_params.get('foo', None) == 3
        assert model.linear.sparse_params.get('bar', None) == 2
        assert model.linear.sparse_params.get('baz', None) is None

    def test_mask_squash_with_params2(self):
        model = Model()
        sparsifier = ImplementedSparsifier(foo=3, bar=2, baz=1)
        sparsifier.prepare(model, [{'tensor_fqn': 'linear.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.linear, 'weight')
        assert hasattr(model.seq[0], 'sparse_params')
        assert hasattr(model.linear, '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.linear.sparse_params.get('foo', None) == 3
        assert model.linear.sparse_params.get('bar', None) == 2
        assert model.linear.sparse_params.get('baz', None) is None

    def test_mask_squash_with_params3(self):
        model = Model()
        sparsifier = ImplementedSparsifier(foo=3, bar=2, baz=1)
        sparsifier.prepare(model, [{'tensor_fqn': 'linear.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.linear, 'weight')
        assert hasattr(model.seq[0], 'sparse_params')
        assert hasattr(model.linear, '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.linear.sparse_params.get('foo', None) == 3
        assert model.linear.sparse_params.get('bar', None) == 2
        assert model.linear.sparse_params.get('baz', None) is None


class TestWeightNormSparsifier(TestCase):
    def test_constructor(self):
        model = Model()
        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', 'linear', 'head')

    def test_step(self):
        model = Model()
        sparsifier = WeightNormSparsifier(sparsity_level=0.5)
        sparsifier.prepare(model, config=[{'tensor_fqn': 'linear.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.linear.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.linear.weight.data = torch.randn(model.linear.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 = Model()
        sparsifier = WeightNormSparsifier(sparsity_level=1.0,
                                          sparse_block_shape=(1, 4),
                                          zeros_per_block=2)
        sparsifier.prepare(model, config=[{'tensor_fqn': 'linear.weight'}])
        sparsifier.step()
        # make sure the sparsity level is approximately 50%
        mask = model.linear.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 = Model()
        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 = Model()
        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 = Model()
        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', 'linear', 'head')

    def test_step(self):
        model = Model()
        sparsifier = NearlyDiagonalSparsifier(nearliness=1)
        sparsifier.prepare(model, config=[{'tensor_fqn': 'linear.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.linear.weight.data = torch.randn(model.linear.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 = Model()
        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 = Model()
        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(nearliness for nearliness in 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