#
# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
#
"""
This file contains layer linting functions.
"""

from collections import OrderedDict
from typing import Dict

import pandas as pd

from .activations import create_activations
from .engine_plan import EnginePlan


class ConvLinter():
    """Convolution layer linter."""

    def __init__(self, plan: EnginePlan):
        self.plan = plan
        self.convs = plan.get_layers_by_type("Convolution")

    def tc_lint(self):
        """Search for Convolutions which are not accelerated by TensorCode"""

        def is_small_conv(conv):
            inputs, _ = create_activations(conv)
            if len(inputs[0].shape) != 4:
                return False
            n, c, h, w = inputs[0].shape
            return c < 32

        report = OrderedDict()

        # Look for kernels that are not scheduled for TensorCore acceleration
        tc_candidates = self.convs.query(f"precision != \"FP32\"").copy()

        # Identify acceleration from tactic name
        df = tc_candidates
        df = df[df["tactic"].str.contains("imma|hmma|xmma|i88|884", na=False) == False]
        for _, conv in df.iterrows():
            mitigation = ""
            if is_small_conv(conv):
                mitigation = "This Convolution has a small number " \
                    "of input channels so acceleration may not be possible."

            report[conv.Name] = OrderedDict({"name": conv.Name, "tactic": conv.tactic, "subtype": conv.subtype, "hazard": "Convolution is not accelerated.", "mitigation": mitigation, "help": "TensorCores accelerate large Convolution and GEMM operations."})
        return report

    def mixed_precision_lint(self):
        """Search for Convolutions with Int8 inputs and Float outputs"""
        report = OrderedDict()

        df = self.convs
        df = df.loc[df["precision"] == "INT8"].copy()
        for index, conv in df.iterrows():
            inputs, outputs = create_activations(conv)
            inf = inputs[0].format[:4]
            outf = outputs[0].format[:4]
            found = inf == "Int8" and outf != "Int8"
            if found:
                report[conv.Name] = OrderedDict({
                    "name": conv.Name,
                    "tactic": conv.tactic,
                    "subtype": conv.subtype,
                    "hazard": "Quantized Convolution has float outputs.",
                    "mitigation": "Consider adding quantization after the convolution.",
                    "help": "Quantized Convolution with float outputs is ill advised "
                    "for memory-limited convolutions."
                })
        return report

    def alignment_lint(self):

        def alignment_ok(conv):
            inputs, outputs = create_activations(conv)
            prec = conv["precision"]
            if len(inputs[0].shape) != 4 or len(outputs[0].shape) != 4:
                return True
            N, C, H, W = inputs[0].shape
            _, K, P, Q = outputs[0].shape
            if prec == "INT8":
                ok = (C % 16 == 0) and (K % 16 == 0)
            else:
                ok = (C % 8 == 0) and (K % 8 == 0)
            return ok

        report = OrderedDict()
        for _, conv in self.convs.iterrows():
            if not alignment_ok(conv):
                report[conv.Name] = {
                    "name": conv.Name,
                    "tactic": conv.tactic,
                    "subtype": conv.subtype,
                    "hazard": "Convolution channels are not optimally aligned.",
                    "mitigation": "Consider changing the alignment of the convolution's channels.",
                    "help": "For best performance, the input and outputs channels of a Tensor Core"
                    "accelerated convolution should be aligned to 8 (FP32/FP16) or 16 (INT8)"
                }
        return report

    def lint(self):
        report = self.tc_lint()
        report.update(self.mixed_precision_lint())
        report.update(self.alignment_lint())
        df = pd.DataFrame.from_dict(report, orient="index")
        return df

class ReformatLinter():
    """Reformat layer linter."""

    def __init__(self, plan: EnginePlan):
        self.plan = plan
        self.reformats = plan.get_layers_by_type("Reformat")

    def lint(self):
        """Search for conversions between types.

        Conversions between layouts are assumed to be optimized."""
        report = OrderedDict()

        for index, reformat in self.reformats.iterrows():
            inputs, outputs = create_activations(reformat)
            inf = inputs[0].format[:4]
            outf = outputs[0].format[:4]
            if inf != outf:
                mitigation = ""
                if "INT8" in [inf, outf]:
                    mitigation = "Consider adding quantization around float operations."
                report[reformat.Name] = OrderedDict({
                    "name": reformat.Name,
                    "origin": reformat["attr.origin"],
                    'type conversion': f"{inf} -> {outf}",
                    'shape conversion': f"{inputs[0].shape} -> {outputs[0].shape}",
                    "hazard": "Reformat layer is converting operand data type.",
                    "mitigation": mitigation,
                    "help": "Conversions between float32 and float16 are a red "
                    "flag, as are conversions between float32/16 and INT8."
                })

        df = pd.DataFrame.from_dict(report, orient="index")
        return df

class SliceLinter():
    """Slice layer linter."""

    def __init__(self, plan: EnginePlan):
        self.plan = plan
        self.slices = plan.get_layers_by_type("Slice")

    def lint(self):
        """Search for conversions between types.

        Conversions between layouts are assumed to be optimized."""
        report = OrderedDict()

        for index, slice in self.slices.iterrows():
            inputs, outputs = create_activations(slice)
            inf = inputs[0].format[:4]
            outf = outputs[0].format[:4]
            if inf != outf:
                mitigation = ""
                if "INT8" in [inf, outf]:
                    mitigation = "Consider adding quantization around float operations."
                report[slice.Name] = OrderedDict({
                    "name": slice.Name,
                    'type conversion': f"{inf} -> {outf}",
                    'shape conversion': f"{inputs[0].shape} -> {outputs[0].shape}",
                    "hazard": "Slice layer is converting operand data type.",
                    "mitigation": mitigation,
                    "help": "Conversions between float32 and float16 are a red "
                    "flag, as are conversions between float32/16 <=> INT8."
                })

        df = pd.DataFrame.from_dict(report, orient="index")
        return df

class QDQLinter():
    """Q/DQ layer linter."""

    def __init__(self, plan: EnginePlan):
        self.plan = plan
        self.scales = plan.get_layers_by_type("Scale")

    def lint(self):
        """Search for dangling Q/DQ layers."""
        report = OrderedDict()

        for index, scale in self.scales.iterrows():
            inputs, outputs = create_activations(scale)
            inf = inputs[0].format[:4]
            outf = outputs[0].format[:4]
            is_qdq = ("Int8" in inputs[0].format) ^ ("Int8" in outputs[0].format)
            if is_qdq:
                dq = "Int8" in inputs[0].format
                role = "Quantize" if not dq else "Dequanitize"
                report[scale.Name] = OrderedDict({
                    "name": scale.Name,
                    'type conversion': f"{inf} -> {outf}",
                    "hazard": f"Unfused {role} layer",
                    "mitigation": f"Check why the {role} layer is not fused",
                    "help": f"Unfused Quantize/Dequantize nodes are wasteful and "
                    "should be avoided. Quantize nodes may be necessary "
                    "for quantizing inputs."
                })

        df = pd.DataFrame.from_dict(report, orient="index")
        return df
