# Copyright (C) 2023, Advanced Micro Devices, Inc. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause

import torch
from torch import nn

from brevitas.fx import GraphModule
from brevitas.fx import immutable_dict
from brevitas.fx import immutable_list
import brevitas.nn as qnn
from brevitas.nn.utils import merge_bn

from .base import UntilFixedPointGraphTransform
from .utils import del_module
from .utils import get_module
from .utils import get_output_channel_dim
from .utils import get_output_channels
from .utils import matches_module_pattern
from .utils import replace_all_uses_except

__all__ = ['MoveSplitBatchNormBeforeCat', 'MergeBatchNorm', 'CollapseConsecutiveConcats']


class MoveSplitBatchNormBeforeCat(UntilFixedPointGraphTransform):

    DEFAULT_BEFORE_MODULES_TYPES = (
        nn.BatchNorm1d,
        nn.BatchNorm2d,
        nn.BatchNorm3d,
        nn.Linear,
        nn.Conv1d,
        nn.Conv2d,
        nn.Conv3d,
        nn.ConvTranspose1d,
        nn.ConvTranspose2d,
        qnn.QuantLinear,
        qnn.QuantConv1d,
        qnn.QuantConv2d,
        qnn.QuantConvTranspose1d,
        qnn.QuantConvTranspose2d)

    def __init__(self, before_modules_types=DEFAULT_BEFORE_MODULES_TYPES):
        super(MoveSplitBatchNormBeforeCat, self).__init__()
        self.before_modules_types = before_modules_types

    def is_converged(self, graph_model: GraphModule) -> bool:
        for cat_node in graph_model.graph.nodes:
            if (cat_node.target is torch.cat and len(cat_node.users) == 1 and
                    cat_node.kwargs['dim'] == 1 and list(cat_node.users)[0].op == 'call_module'):
                bn_node = list(cat_node.users)[0]
                module = get_module(graph_model, bn_node.target)
                if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d)):
                    inp_nodes = cat_node.all_input_nodes
                    if all([inp_node.op == 'call_module' and len(inp_node.users) == 1
                            for inp_node in inp_nodes]):
                        before_mods = [
                            get_module(graph_model, inp_node.target) for inp_node in inp_nodes]
                        if all(isinstance(mod, self.before_modules_types) for mod in before_mods):
                            assert inp_nodes == cat_node.kwargs['tensors']
                            num_features_list = [get_output_channels(mod) for mod in before_mods]
                            chunk_bn_list = [type(module)(n) for n in num_features_list]
                            for i, chunk_bn in enumerate(chunk_bn_list):
                                chunk_bn_name = f'{bn_node.name}_{i}'
                                graph_model.add_module(chunk_bn_name, chunk_bn)
                                start = sum(num_features_list[:i])
                                end = sum(num_features_list[:i + 1])
                                chunk_bn.weight.data = module.weight.data[start:end]
                                chunk_bn.bias.data = module.bias.data[start:end]
                                chunk_bn.running_mean = module.running_mean.data[start:end]
                                chunk_bn.running_var = module.running_var.data[start:end]
                                inp_node = cat_node.kwargs['tensors'][i]
                                with graph_model.graph.inserting_after(inp_node):
                                    chunk_bn_node = graph_model.graph.call_module(
                                        chunk_bn_name, args=(inp_node,))
                                replace_all_uses_except(inp_node, chunk_bn_node, [chunk_bn_node])
                            bn_node.replace_all_uses_with(cat_node)
                            graph_model.graph.erase_node(bn_node)
                            del_module(graph_model, bn_node.target)
                            graph_model.graph.lint()
                            graph_model.recompile()
                            return False
        return True


class MergeBatchNorm(UntilFixedPointGraphTransform):

    DEFAULT_PATTERNS = ((nn.BatchNorm1d, nn.BatchNorm1d), (nn.BatchNorm2d, nn.BatchNorm2d),
                        (nn.BatchNorm3d, nn.BatchNorm3d), (nn.Linear, nn.BatchNorm1d),
                        (nn.Conv1d, nn.BatchNorm1d), (nn.Conv2d, nn.BatchNorm2d),
                        (nn.Conv3d, nn.BatchNorm3d), (nn.ConvTranspose1d, nn.BatchNorm1d),
                        (nn.ConvTranspose2d, nn.BatchNorm2d), (nn.ConvTranspose3d, nn.BatchNorm3d),
                        (qnn.BatchNorm1dToQuantScaleBias,
                         nn.BatchNorm1d), (qnn.BatchNorm2dToQuantScaleBias,
                                           nn.BatchNorm2d), (qnn.QuantLinear, nn.BatchNorm1d),
                        (qnn.QuantConv1d, nn.BatchNorm1d), (qnn.QuantConv2d, nn.BatchNorm2d),
                        (qnn.QuantConvTranspose1d,
                         nn.BatchNorm1d), (qnn.QuantConvTranspose2d, nn.BatchNorm2d))

    def __init__(self, patterns=DEFAULT_PATTERNS):
        super(MergeBatchNorm, self).__init__()
        self.patterns = list(patterns)

    def is_converged(self, graph_model: GraphModule):
        named_modules = dict(graph_model.named_modules())
        for node in graph_model.graph.nodes:
            for pattern in self.patterns:
                if matches_module_pattern(pattern, node, named_modules):
                    if len(node.args[0].users) > 1:
                        continue
                    layer = named_modules[node.args[0].target]
                    bn = named_modules[node.target]
                    merge_bn(layer, bn, get_output_channel_dim(layer))
                    node.replace_all_uses_with(node.args[0])
                    graph_model.graph.erase_node(node)
                    del_module(graph_model, node.target)
        graph_model.recompile()
        graph_model.graph.lint()
        return graph_model


class CollapseConsecutiveConcats(UntilFixedPointGraphTransform):

    def merge_tensor_kwargs(self, node_to_extract, node_to_merge_in):
        cat_tensors1 = list(node_to_extract.kwargs['tensors'])
        cat_tensors2 = node_to_merge_in.kwargs['tensors']
        if not isinstance(cat_tensors2, (list, tuple)):
            cat_tensors2 = [cat_tensors2]
        index_for_insertion = cat_tensors2.index(node_to_extract)
        cat_tensors2 = [t for t in cat_tensors2 if t is not node_to_extract]
        cat_tensors2[index_for_insertion:index_for_insertion] = cat_tensors1
        kwargs = dict(node_to_merge_in.kwargs)
        kwargs['tensors'] = cat_tensors2
        node_to_merge_in.kwargs = immutable_dict(kwargs)

    def is_converged(self, graph_model):
        for i, node in enumerate(graph_model.graph.nodes):
            if node.op == 'call_function' and node.target is torch.cat:
                for inp_node in node.all_input_nodes:
                    if (inp_node.op == 'call_function' and inp_node.target is torch.cat and
                            node.kwargs['dim'] == inp_node.kwargs['dim'] and
                            len(inp_node.users) == 1):
                        self.merge_tensor_kwargs(inp_node, node)
                        graph_model.graph.erase_node(inp_node)
                        graph_model.graph.lint()
                        graph_model.recompile()
                        return False
        return True

    def move_args_to_kwargs(self, graph_model):
        for node in graph_model.graph.nodes:
            if node.op == 'call_function' and node.target is torch.cat:
                if len(node.args) > 0:
                    if isinstance(node.args[-1], int):
                        kwargs = dict(node.kwargs)
                        kwargs['dim'] = node.args[-1]
                        node.kwargs = immutable_dict(kwargs)
                        node.args = node.args[:-1]
                    if isinstance(node.args[0], (tuple, list)):
                        kwargs = dict(node.kwargs)
                        kwargs['tensors'] = node.args[0]
                        node.kwargs = immutable_dict(kwargs)
                        node.args = node.args[1:]

    def apply(self, graph_model: GraphModule) -> GraphModule:
        self.move_args_to_kwargs(graph_model)
        return super(CollapseConsecutiveConcats, self).apply(graph_model)
