import torch.nn.utils.prune as prune
import torch
from src.vanilla_pytorch.utils import count_rem_weights
from src.vanilla_pytorch.models.linearnets import LeNet, init_weights
from src.vanilla_pytorch.models.resnets import Resnets


def remove_pruning(model):
    for i, (name, module) in enumerate(model.named_modules()):
        # name and val
        if any([isinstance(module, cl) for cl in [torch.nn.Conv2d, torch.nn.Linear]]):
            prune.remove(module, 'weight')


def get_masks(model, prune_amts=None):
    """
    prune the lowest p% weights by magnitude per layer

    :param model: model to prune
    :param p_rate: prune rate = 0.2 as per paper
    :param prune_amts: dictionary
    :return: the created mask. model has served it's purpose.
    """
    # TODO: Adjust pruning with output layer
    if prune_amts is None:  # ie dict is empty, use the default prune rate = 0.2
        prune_amts = {"linear": 0.2, "conv": 0.2, "last": 0.2}

    for i, (name, module) in enumerate(model.named_modules()):
        # prune 20% of connections in all 2D-conv layers
        if isinstance(module, torch.nn.Conv2d):
            module = prune.l1_unstructured(module, name='weight', amount=prune_amts['conv'])
        # prune 20% of connections in all linear layers
        elif isinstance(module, torch.nn.Linear):
            module = prune.l1_unstructured(module, name='weight', amount=prune_amts['linear'])

    masks = list(model.named_buffers())
    remove_pruning(model)
    return masks


def update_apply_masks(model, masks):
    # doesn't seem to be needed.
    # for key, val in masks.items():
    #     print(f"key {key}")
    #     layer = getattr(model, key.split('.')[0])
    #     layer.weight_mask = val
    for name, module in model.named_modules():
        if any([isinstance(module, cl) for cl in [torch.nn.Conv2d, torch.nn.Linear]]):
            module = prune.custom_from_mask(module, name='weight', mask=masks[name + ".weight_mask"])
    # remove_pruning(model)
    return model



def prune_random(model, prune_amts=None):
    if prune_amts is None:  # ie dict is empty, use the default prune rate =0.2
        prune_amts = {"linear": 0.2, "conv": 0.2, "last": 0.2}
    for name, module in model.named_modules():
        # prune 20% of connections in all 2D-conv layers
        if isinstance(module, torch.nn.Conv2d):
            module = prune.random_unstructured(module, name='weight', amount=prune_amts['conv'])
        # prune 20% of connections in all linear layers
        elif isinstance(module, torch.nn.Linear):
            module = prune.random_unstructured(module, name='weight', amount=prune_amts['linear'])

    remove_pruning(model)


if __name__ == '__main__':
    net = Resnets(in_channels=3)
    net.apply(init_weights)
    prune_rate = 0.8
    prune_custom = {"linear": 0.2, "conv": 0.2, "last": 0.1}
    for i in range(3):
        masks = get_masks(net, prune_amts=prune_custom)
        print(f"Count zero : {count_rem_weights(net)}")

[((70, 10, 70, 32), 'src.vanilla_pytorch.models.resnets.Resnets', 'Resnets', (), '', False, 'from src.vanilla_pytorch.models.resnets import Resnets\n'), ((12, 12, 12, 42), 'torch.nn.utils.prune.remove', 'prune.remove', ({(12, 25, 12, 31): 'module', (12, 33, 12, 41): '"""weight"""'}, {}), "(module, 'weight')", True, 'import torch.nn.utils.prune as prune\n'), ((31, 21, 31, 92), 'torch.nn.utils.prune.l1_unstructured', 'prune.l1_unstructured', (), '', True, 'import torch.nn.utils.prune as prune\n'), ((49, 21, 49, 101), 'torch.nn.utils.prune.custom_from_mask', 'prune.custom_from_mask', (), '', True, 'import torch.nn.utils.prune as prune\n'), ((61, 21, 61, 96), 'torch.nn.utils.prune.random_unstructured', 'prune.random_unstructured', (), '', True, 'import torch.nn.utils.prune as prune\n'), ((34, 21, 34, 94), 'torch.nn.utils.prune.l1_unstructured', 'prune.l1_unstructured', (), '', True, 'import torch.nn.utils.prune as prune\n'), ((64, 21, 64, 98), 'torch.nn.utils.prune.random_unstructured', 'prune.random_unstructured', (), '', True, 'import torch.nn.utils.prune as prune\n'), ((76, 30, 76, 52), 'src.vanilla_pytorch.utils.count_rem_weights', 'count_rem_weights', ({(76, 48, 76, 51): 'net'}, {}), '(net)', False, 'from src.vanilla_pytorch.utils import count_rem_weights\n')]