# Copyright (C) 2023, Advanced Micro Devices, Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of FINN nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Transformations for specializing FINN layers to HLS or RTL implementations.

This module provides functionality to automatically select and specialize FINN dataflow
layers to their optimal hardware implementation variants (HLS or RTL) based on FPGA
target, layer constraints, and user preferences.
"""

import numpy as np
from onnx import helper
from qonnx.custom_op.registry import getCustomOp
from qonnx.transformation.base import Transformation
from qonnx.transformation.general import GiveUniqueNodeNames

from finn.custom_op.fpgadataflow.hls import custom_op as hls_variants
from finn.custom_op.fpgadataflow.rtl import custom_op as rtl_variants
from finn.util.basic import get_dsp_block, is_versal
from finn.util.logging import log


def _determine_impl_style(node, fpgapart, model):
    """Determine the optimal implementation style (HLS or RTL) for a given node.

    Analyzes node constraints, FPGA capabilities, and user preferences to select
    the best hardware implementation variant for a dataflow layer.
    """
    optype = node.op_type

    # check if there is an HLS or RTL variant or both
    hls_variant = optype + "_hls" in hls_variants.keys()
    rtl_variant = optype + "_rtl" in rtl_variants.keys()

    # check if user has specified a preferred_impl_style
    node_inst = getCustomOp(node)
    impl_style = node_inst.get_nodeattr("preferred_impl_style")

    # if impl_style not set, for "simple" layers always try
    # to use rtl variant if available
    if impl_style == "":
        if optype == "StreamingDataWidthConverter":
            return _dwc_determine_impl_style(node)
        if rtl_variant:
            if optype == "MVAU":
                idt = node_inst.get_input_datatype(0)
                wdt = node_inst.get_input_datatype(1)
                inp_width_fit = idt.bitwidth() >= 4
                weight_width_fit = wdt.bitwidth() >= 4
                if inp_width_fit and weight_width_fit and _mvu_rtl_possible(node, fpgapart, model):
                    return "rtl"
                else:
                    return "hls"
            elif optype == "VVAU":
                idt = node_inst.get_input_datatype(0)
                wdt = node_inst.get_input_datatype(1)
                inp_width_fit = idt.bitwidth() >= 4
                weight_width_fit = wdt.bitwidth() >= 4
                if inp_width_fit and weight_width_fit and _vvu_rtl_possible(node, fpgapart):
                    return "rtl"
                else:
                    return "hls"
            return "rtl"
        # but if no rtl variant, set impl_style to hls
        elif hls_variant:
            return "hls"
        # if there is neither an rtl nor hls variant
        # throw error
        else:
            raise Exception(
                """Node {} with optype {} has no hw implementation variant)""".format(
                    node.name, optype
                )
            )

    # check if user setting can be fulfilled
    # otherwise change impl_style
    elif impl_style == "hls":
        if hls_variant:
            return "hls"
        elif rtl_variant:
            warn_str = """There is no HLS variant of %s. Node %s will automatically be
                        set to RTL variant.""" % (
                node.op_type,
                node.name,
            )
            log.warning(warn_str)
            return "rtl"
        else:
            raise Exception(
                """Node {} with optype {} has no hw implementation variant)""".format(
                    node.name, optype
                )
            )
    elif impl_style == "rtl":
        # rtl dwc does not support every inWidth to outWidth ratio
        if optype == "StreamingDataWidthConverter":
            if _dwc_determine_impl_style(node) != "rtl":
                warn_str = """RTL implementation of DWC requires
                            stream widths that are integer width ratios
                            from each other. Node %s will automatically be
                            set to HLS variant.""" % (
                    node.name,
                )
                log.warning(warn_str)
                return "hls"
            else:
                # user setting can be fulfilled
                return "rtl"
        elif optype == "MVAU":
            if _mvu_rtl_possible(node, fpgapart, model):
                return "rtl"
            else:
                warn_str = """There is no RTL variant for %s. The node will automatically be
                        set to HLS variant. Please check the bit-widths to be <= 8 and ensure the
                        thresholds are implemented as standalone layer""" % (
                    node.name,
                )
                log.warning(warn_str)
                return "hls"
        elif optype == "VVAU":
            if _vvu_rtl_possible(node, fpgapart):
                return "rtl"
            else:
                warn_str = """There is no RTL variant for %s. The node will automatically be
                        set to HLS variant. Please check the bit-widths to be <= 8 and ensure the
                        thresholds are implemented as standalone layer. Note that the RTL-variant
                        of this layer is only supported on Versal boards""" % (
                    node.name,
                )
                log.warning(warn_str)
                return "hls"

        if rtl_variant:
            return "rtl"
        elif hls_variant:
            warn_str = """There is no RTL variant of %s. Node %s will automatically be
                        set to HLS variant.""" % (
                node.op_type,
                node.name,
            )
            log.warning(warn_str)
            return "hls"
        else:
            raise Exception(
                """Node {} with optype {} has no hw implementation variant)""".format(
                    node.name, optype
                )
            )
    else:
        raise Exception(
            """Invalid value for attribute preferred_impl_style! Is currently set to: {}
            has to be set to one of the following value ("hls", "rtl")""".format(
                impl_style
            )
        )


def _dwc_determine_impl_style(node):
    """Determine implementation style for StreamingDataWidthConverter nodes.

    When possible, uses RTL variant based on width ratio compatibility.
    """
    dwc = getCustomOp(node)
    dwc_in_width = dwc.get_nodeattr("inWidth")
    dwc_out_width = dwc.get_nodeattr("outWidth")
    # check if rtl variant can be used
    iwidth_d = dwc_in_width % dwc_out_width == 0
    owidth_d = dwc_out_width % dwc_in_width == 0
    if iwidth_d or owidth_d:
        return "rtl"
    else:
        return "hls"


def _mvu_rtl_possible(n, fpgapart, model):
    """Check whether RTL-based MVU implementation is supported for given node.

    RTL-MVU constraints:
    - DSP48: supports up to 8sx8u (8-bit signed weights x 8-bit activations)
    - DSP58: supports up to 8sx9s
    - DSP48E1: only supports narrow range weights
    - No embedded thresholding or binaryXnor mode supported
    """
    node_inst = getCustomOp(n)
    # first check if no Activation or binary xnor mode and return False
    # immediately if one of them is True
    no_activation = node_inst.get_nodeattr("noActivation") == 0
    not_binaryxnor_mode = node_inst.get_nodeattr("binaryXnorMode") == 1
    if no_activation or not_binaryxnor_mode:
        return False

    # check if weights are signed, if not return False
    wdt = node_inst.get_input_datatype(1)
    if not wdt.signed():
        return False

    # check which dsp block is available on fpga
    dsp_block = get_dsp_block(fpgapart)
    # check if weights are narrow
    weights = model.get_initializer(n.input[1])
    # if dynamic input, set minimum of weights to wdt.min()
    # otherwise set it to the minimum value in the weight matrix
    if weights is None:
        weights_min = wdt.min()
    else:
        weights_min = np.min(weights)
    narrow_weights = False if weights_min == wdt.min() else True
    # if non narrow weights and only DSP48E1 available return False
    if not narrow_weights and dsp_block == "DSP48E1":
        return False

    # if none of the above constraints have been triggered
    # we now check if input and weight data types are in range
    # we only use rtl mvau if the dtypes are at least 2 bit
    idt = node_inst.get_input_datatype()
    inp_width_in_range = (2 <= idt.bitwidth() <= 8) or (idt.bitwidth() == 9 and idt.signed())
    weight_width_in_range = 2 <= wdt.bitwidth() <= 8

    return inp_width_in_range and weight_width_in_range


def _vvu_rtl_possible(n, fpgapart):
    """Check whether RTL-based VVU implementation is supported for given node.

    RTL-VVU constraints:
    - Only supported on Versal DSP58
    - Supports up to 8sx9s inputs (8-bit signed weights x 9-bit signed or 8-bit activations)
    - No embedded thresholding functionality supported
    """
    node_inst = getCustomOp(n)
    if not node_inst.get_nodeattr("noActivation"):
        return False
    if not is_versal(fpgapart):
        return False

    idt = node_inst.get_input_datatype(0)
    wdt = node_inst.get_input_datatype(1)
    in_width_in_range = (idt.bitwidth() <= 8) or (idt.bitwidth() == 9 and idt.min() < 0)
    weight_width_in_range = wdt.bitwidth() <= 8
    signed_weights = wdt.min() < 0

    return in_width_in_range and weight_width_in_range and signed_weights


class SpecializeLayers(Transformation):
    """Specialize all layers to either HLS or RTL variants"""

    def __init__(self, fpgapart):
        """Initialize the SpecializeLayers transformation.

        Args:
            fpgapart: Target FPGA part string for implementation selection
        """
        super().__init__()
        self.fpgapart = fpgapart

    def apply(self, model):
        """Apply layer specialization transformation to model.

        Converts all dataflow layers to their optimal HLS or RTL implementation
        variants based on target FPGA and layer constraints.
        """
        graph = model.graph
        node_ind = 0
        graph_modified = False
        for node in graph.node:
            # Skip nodes that are not hw layers
            if not node.domain == "finn.custom_op.fpgadataflow":
                continue
            node_ind += 1
            impl_style = _determine_impl_style(node, self.fpgapart, model)
            optype = node.op_type + "_" + impl_style

            new_node = helper.make_node(
                optype,
                node.input,
                node.output,
                domain="finn.custom_op.fpgadataflow." + impl_style,
            )
            # add all attributes
            for attribute in node.attribute:
                if attribute.name != "preferred_impl_style":
                    new_node.attribute.append(attribute)
            graph.node.insert(node_ind, new_node)
            # remove old nodes
            graph.node.remove(node)
            # update node names to reflect new op types
            model = model.transform(GiveUniqueNodeNames())
            graph_modified = True
        return (model, graph_modified)
