# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
#
# 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.

import paddle

from paddle.fluid import core, framework
from paddle.fluid.dygraph.parallel import ParallelEnv
from paddle.fluid.contrib.slim.quantization import utils
from paddle.fluid.contrib.slim.quantization import QuantizationTransformPassV2
from paddle.fluid.contrib.slim.quantization import AddQuantDequantPassV2
from paddle.fluid.contrib.slim.quantization import OutScaleForTrainingPass
from paddle.distributed.auto_parallel.dist_attribute import (
    OperatorDistributedAttribute,
    TensorDistributedAttribute,
)

from .pass_base import PassBase, register_pass

TRANSFORM_PASS_OP_TYPES = utils._weight_supported_quantizable_op_type
QUANT_DEQUANT_PASS_OP_TYPES = utils._act_supported_quantizable_op_type


def _node_id(node):
    return (node.node.graph_id(), node.node.id())


@register_pass("auto_parallel_quantization")
class QuantizationPass(PassBase):
    def __init__(self):
        super(QuantizationPass, self).__init__()
        self.set_attr("dist_context", None)
        self.set_attr("params_grads", None)

    def _check_self(self):
        if self.get_attr("dist_context") is None:
            return False
        if self.get_attr("params_grads") is None:
            return False
        return True

    def _check_conflict(self, other_pass):
        return True

    def _apply_single_impl(self, main_program, startup_program, context):

        dist_context = self.get_attr("dist_context")
        params_grads = self.get_attr("params_grads")

        # TODO: scope and place will be removed,
        # cause params should be initialized by engine module.
        scope = paddle.static.global_scope()
        place = paddle.fluid.CUDAPlace(ParallelEnv().dev_id)

        # 1. Program convert to Graph, and this pass is only for train mode
        main_graph = framework.IrGraph(
            core.Graph(main_program.desc), for_test=False
        )

        # 2. Prepare inputs
        transform_pass_ops = []
        quant_dequant_ops = []
        quantize_op_types = [
            'conv2d',
            'depthwise_conv2d',
            'mul',
            'matmul',
            'matmul_v2',
        ]
        for op_type in quantize_op_types:
            if op_type in TRANSFORM_PASS_OP_TYPES:
                transform_pass_ops.append(op_type)
            elif op_type in QUANT_DEQUANT_PASS_OP_TYPES:
                quant_dequant_ops.append(op_type)

        weight_quantize_type = (
            "channel_wise_abs_max"
            if self.get_attr('channel_wise_abs_max')
            else "abs_max"
        )

        # 3. Add quant op for ops which have parameters
        transform_pass = QuantizationTransformPassV2(
            scope=scope,
            place=place,
            weight_bits=self.get_attr('weight_bits'),
            activation_bits=self.get_attr('activation_bits'),
            skip_pattern=self.get_attr('not_quant_pattern'),
            activation_quantize_type="moving_average_abs_max",
            quantizable_op_type=transform_pass_ops,
            weight_quantize_type=weight_quantize_type,
            weight_quantize_func=None,
            act_quantize_func=None,
            weight_preprocess_func=None,
            act_preprocess_func=None,
            optimizer_func=None,
            executor=None,
        )
        transform_pass.apply(main_graph)

        # 4. Add quant op for ops which don't have parameter
        quant_dequant_pass = AddQuantDequantPassV2(
            scope=scope,
            place=place,
            quant_bits=self.get_attr('activation_bits'),
            skip_pattern=self.get_attr('not_quant_pattern'),
            quantizable_op_type=quant_dequant_ops,
        )
        quant_dequant_pass.apply(main_graph)

        # 5. Gather quantitative information for the output
        out_scale_training_pass = OutScaleForTrainingPass(
            scope=scope, place=place
        )
        out_scale_training_pass.apply(main_graph)

        # 6. Convert Graph back to Program
        quant_program = main_graph.to_program()

        # 7. get new prams_grads from quant_program
        new_params_grads = []
        for param, grad in params_grads:
            if param.name not in quant_program.global_block().vars:
                continue

            new_param = quant_program.global_block().vars[param.name]
            new_grad = quant_program.global_block().vars[grad.name]
            new_params_grads.append((new_param, new_grad))

        # 8. complete distributed attribution
        # NOTE: hack implement, upgrading soon
        for ib, block in enumerate(quant_program.blocks):
            # recover origin ops' dist_attr and set quant ops' dist_attr
            qat_offset = 0
            for ip, quant_op in enumerate(block.ops):
                quant_op_dist_attr = OperatorDistributedAttribute()

                if (
                    "quantize" in quant_op.type
                    or quant_op.type == "moving_average_abs_max_scale"
                ):

                    input_name = quant_op.desc.input('X')[0]
                    if "quantize" in input_name:
                        input_name = input_name[
                            : input_name.index(".quantized")
                        ]

                    if quant_op.type == "moving_average_abs_max_scale":
                        consume_op = main_program.blocks[ib].vars[input_name].op
                    else:
                        consume_op = main_program.blocks[ib].ops[
                            ip - qat_offset
                        ]
                    consume_op_dist_attr = dist_context.get_dist_op_for_program(
                        consume_op
                    ).dist_attr
                    ref_process_mesh = consume_op_dist_attr.process_mesh

                    if input_name in consume_op_dist_attr.outputs_dist_attrs:
                        consume_input_dist_attr = (
                            consume_op_dist_attr.outputs_dist_attrs[input_name]
                        )
                    else:
                        consume_input_dist_attr = (
                            consume_op_dist_attr.inputs_dist_attrs[input_name]
                        )

                    quant_op_dist_attr.impl_idx = 0
                    quant_op_dist_attr.impl_type = "default"
                    quant_op_dist_attr.process_mesh = ref_process_mesh
                    quant_op_dist_attr.set_input_dist_attr(
                        quant_op.desc.input('X')[0], consume_input_dist_attr
                    )

                    for slot_name in quant_op.desc.input_names():
                        if slot_name == "X":
                            continue
                        for in_name in quant_op.desc.input(slot_name):
                            input_var = block.vars[in_name]
                            tensor_dist_attr = TensorDistributedAttribute()
                            tensor_dist_attr.process_mesh = ref_process_mesh
                            tensor_dist_attr.dims_mapping = [-1]
                            dist_context.set_tensor_dist_attr_for_program(
                                input_var, tensor_dist_attr
                            )
                            quant_op_dist_attr.set_input_dist_attr(
                                in_name, tensor_dist_attr
                            )

                    for slot_name in quant_op.desc.output_names():
                        output_name = quant_op.desc.output(slot_name)[0]
                        output_var = block.vars[output_name]
                        if slot_name == "Y":
                            dist_context.set_tensor_dist_attr_for_program(
                                output_var, consume_input_dist_attr
                            )
                            quant_op_dist_attr.set_output_dist_attr(
                                output_name, consume_input_dist_attr
                            )
                        else:
                            tensor_dist_attr = TensorDistributedAttribute()
                            tensor_dist_attr.process_mesh = ref_process_mesh
                            tensor_dist_attr.dims_mapping = [-1]
                            dist_context.set_tensor_dist_attr_for_program(
                                output_var, tensor_dist_attr
                            )
                            quant_op_dist_attr.set_output_dist_attr(
                                output_name, tensor_dist_attr
                            )

                    quant_op._set_attr("op_device", "")
                    qat_offset += 1

                else:

                    origin_op = main_program.blocks[ib].ops[ip - qat_offset]
                    quant_op.desc.set_original_id(origin_op.desc.original_id())
                    dist_origin_op = dist_context.get_dist_op_for_program(
                        origin_op
                    )
                    assert (
                        dist_origin_op is not None
                    ), "origin op must have dist attr."

                    origin_op_dist_attr = dist_origin_op.dist_attr
                    quant_op_dist_attr.impl_idx = origin_op_dist_attr.impl_idx
                    quant_op_dist_attr.impl_type = origin_op_dist_attr.impl_type
                    quant_op_dist_attr.process_mesh = (
                        origin_op_dist_attr.process_mesh
                    )
                    for idx, input_name in enumerate(quant_op.input_arg_names):
                        origin_input_name = origin_op.input_arg_names[idx]
                        origin_input_dist_attr = (
                            origin_op_dist_attr.inputs_dist_attrs[
                                origin_input_name
                            ]
                        )
                        quant_op_dist_attr.set_input_dist_attr(
                            input_name, origin_input_dist_attr
                        )

                        if input_name not in main_program.blocks[ib].vars:
                            origin_input_var = main_program.blocks[ib].vars[
                                origin_input_name
                            ]
                            origin_in_tensor_dist_attr = (
                                dist_context.get_dist_tensor_for_program(
                                    origin_input_var
                                ).dist_attr
                            )
                            quant_input_var = block.vars[input_name]
                            dist_context.set_tensor_dist_attr_for_program(
                                quant_input_var, origin_in_tensor_dist_attr
                            )

                    for idx, output_name in enumerate(
                        quant_op.output_arg_names
                    ):
                        origin_output_name = origin_op.output_arg_names[idx]
                        origin_output_dist_attr = (
                            origin_op_dist_attr.outputs_dist_attrs[
                                origin_output_name
                            ]
                        )
                        quant_op_dist_attr.set_output_dist_attr(
                            output_name, origin_output_dist_attr
                        )

                        if output_name not in main_program.blocks[ib].vars:
                            origin_output_var = main_program.blocks[ib].vars[
                                origin_output_name
                            ]
                            origin_out_tensor_dist_attr = (
                                dist_context.get_dist_tensor_for_program(
                                    origin_output_var
                                ).dist_attr
                            )
                            quant_output_var = block.vars[output_name]
                            dist_context.set_tensor_dist_attr_for_program(
                                quant_output_var, origin_out_tensor_dist_attr
                            )

                dist_context.set_op_dist_attr_for_program(
                    quant_op, quant_op_dist_attr
                )

            # recover vars' dist_attr
            for name, dst_var in block.vars.items():
                if name in main_program.blocks[ib].vars:
                    src_var = main_program.blocks[ib].vars[name]
                    dist_tensor = dist_context.get_dist_tensor_for_program(
                        src_var
                    )
                    if not dist_tensor:
                        continue
                    dist_context.set_tensor_dist_attr_for_program(
                        dst_var, dist_tensor.dist_attr
                    )

        context.set_attr("main_program", quant_program)
        context.set_attr("startup_program", startup_program)
        context.set_attr("params_grads", new_params_grads)