"""
 Copyright (c) 2021 Intel Corporation
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
      http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""
from collections import namedtuple
from typing import NamedTuple

import networkx as nx
import torch
from torch import nn

from nncf.dynamic_graph.context import Scope
from nncf.graph.graph import NNCFGraph
from nncf.layers import NNCFConv2d
from nncf.module_operations import UpdatePaddingValue
from nncf.nncf_network import NNCFNetwork
from nncf.quantization.layers import BaseQuantizer
from nncf.quantization.layers import QuantizationMode
from nncf.quantization.layers import QuantizerConfig
from nncf.quantization.layers import SymmetricQuantizer


class AdjustPaddingArgs(NamedTuple):
    weight_bitwidth: int
    activation_quantizer: BaseQuantizer
    quantized_module: nn.Module
    module_scope: Scope


class CalculatePaddingAdjustment:
    """
    Calculates padding value to perform a workaround for U4 support on VPU.
    VPU supports only i4 for weights and activations with zero-point=0 and padding=0. This imposes some limitations on
    the quantization scheme we can apply. In case of unsigned input for a quantizer (e.g. output of ReLU) half of
    i4 range (8 values) is insufficient to preserve the accuracy. To overcome the problem it is proposed
    to transform u4 to i4 in the VPU plugin by shifting the input by half of the quantization range to left. Padding
    value should be shifted as well. And to make it zero after the shift (non-zero padding values are not
    supported), the model should be trained with padding value equal to the half of the quantization range.
    """

    def __init__(self, activation_quantizer: SymmetricQuantizer):
        if not isinstance(activation_quantizer, SymmetricQuantizer):
            raise RuntimeError('Padding adjustment is not supported for not symmetric quantization')
        self._activation_quantizer = activation_quantizer
        self._is_enabled = True

    @staticmethod
    def is_applicable(args: AdjustPaddingArgs):
        weight_bitwidth, activation_quantizer, module, _ = args
        result = False
        if isinstance(module, NNCFConv2d):
            padding_values = set(module.padding)
            padding_enabled = len(padding_values) >= 1 and padding_values.pop()
            if padding_enabled:
                symmetric = isinstance(activation_quantizer, SymmetricQuantizer)
                per_tensor = not activation_quantizer.per_channel
                a_int4 = activation_quantizer.num_bits == 4
                w_int24 = weight_bitwidth <= 4
                unsigned = not activation_quantizer.signed
                result = symmetric and per_tensor and a_int4 and w_int24 and unsigned
        return result

    def __call__(self, previous_padding_value) -> torch.Tensor:
        if self._is_enabled:
            scale = self._activation_quantizer.scale
            eps = self._activation_quantizer.eps
            safe_scale = abs(scale) + eps
            return safe_scale / 2
        return previous_padding_value

    @staticmethod
    def is_config_applicable(qconfig: QuantizerConfig):
        return not qconfig.per_channel and qconfig.num_bits == 4 and \
               not qconfig.signedness_to_force and qconfig.mode == QuantizationMode.SYMMETRIC


def add_adjust_padding_nodes(nncf_graph: NNCFGraph, model: NNCFNetwork) -> nx.DiGraph():
    # pylint:disable=protected-access
    nx_graph = nncf_graph._get_graph_for_structure_analysis()

    NewNodeArgs = namedtuple('NewNodeArgs', ('node_key', 'attr', 'parent_node_key'))

    args = []
    for node_key, nx_node in nx_graph.nodes.items():
        scope = Scope.from_str(nx_node['scope'])
        module = model.get_module_by_scope(scope)
        if isinstance(module, NNCFConv2d):
            adjust_padding_ops = filter(lambda x: isinstance(x, UpdatePaddingValue), module.pre_ops.values())
            for _ in adjust_padding_ops:
                new_node_key = f'{node_key}_apad'
                attr = dict(type='', label='adjust_padding_value', style='filled', color='yellow')
                args.append(NewNodeArgs(new_node_key, attr, node_key))

    for arg in args:
        nx_graph.add_node(arg.node_key, **arg.attr)
        nx_graph.add_edge(arg.node_key, arg.parent_node_key)
    return nx_graph
