# %%
import torch
import torch.nn as nn

from nni.compression.pruning import L1NormPruner
from nni.compression.utils import auto_set_denpendency_group_ids
from nni.compression.speedup import ModelSpeedup

# %%
class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 40, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(40)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(40, 80, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(80)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv3 = nn.Conv2d(80, 1, kernel_size=3, padding=1)
        self.bn3 = nn.BatchNorm2d(1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.conv3(x)
        x = self.bn3(x)
        return x
    
model = ConvNet()
num_params_unpruned = sum(p.numel() for p in model.parameters())
dummy_input = torch.randn(1, 3, 32, 32)
dummy_output = model(dummy_input)
print(dummy_output.shape)

# %%
sparsity_ratio = 0.5
config_list = [{
    'op_types': ['Conv2d'],
    'sparse_ratio': sparsity_ratio,
}]
config_list = auto_set_denpendency_group_ids(model, config_list, [dummy_input])
pruner = L1NormPruner(model, config_list)
_, masks = pruner.compress()
pruner.unwrap_model()
model = ModelSpeedup(model, [dummy_input], masks, garbage_collect_values=False).speedup_model()

# %%
num_params_pruned = sum(p.numel() for p in model.parameters())
print(f'Number of parameters before pruning: {num_params_unpruned}')
print(f'Number of parameters after pruning: {num_params_pruned}')

num_params_diff = num_params_unpruned - num_params_pruned
prune_ratio = num_params_diff / num_params_unpruned
print(f'Number of parameters pruned: {num_params_diff}')
print(f'Parameter ratio: {(1-prune_ratio)*100:.2f}%')