# Copyright (c) 2020, Xilinx, Inc.
# Copyright (C) 2024, 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.

import json
import os
from qonnx.core.modelwrapper import ModelWrapper
from qonnx.custom_op.registry import getCustomOp
from qonnx.transformation.base import Transformation
from qonnx.transformation.general import (
    GiveReadableTensorNames,
    GiveUniqueNodeNames,
    RemoveUnusedTensors,
)
from subprocess import CalledProcessError

from finn.builder.build_dataflow_config import FpgaMemoryType, VitisOptStrategy
from finn.transformation.fpgadataflow.create_dataflow_partition import CreateDataflowPartition
from finn.transformation.fpgadataflow.create_stitched_ip import CreateStitchedIP
from finn.transformation.fpgadataflow.floorplan import Floorplan
from finn.transformation.fpgadataflow.hlssynth_ip import HLSSynthIP
from finn.transformation.fpgadataflow.insert_dwc import InsertDWC
from finn.transformation.fpgadataflow.insert_fifo import InsertFIFO
from finn.transformation.fpgadataflow.insert_iodma import InsertIODMA
from finn.transformation.fpgadataflow.prepare_ip import PrepareIP
from finn.transformation.fpgadataflow.specialize_layers import SpecializeLayers
from finn.util.basic import launch_process_helper, make_build_dir
from finn.util.exception import FINNError

from . import templates


def _check_vitis_envvars():
    assert "XILINX_VITIS" in os.environ, "XILINX_VITIS must be set for Vitis"
    assert "PLATFORM_REPO_PATHS" in os.environ, "PLATFORM_REPO_PATHS must be set for Vitis"
    assert (
        "XILINX_XRT" in os.environ
    ), "XILINX_XRT must be set for Vitis, ensure the XRT env is sourced"


class CreateVitisXO(Transformation):
    """Create a Vitis object file from a stitched FINN ip.

    Outcome if successful: sets the vitis_xo attribute in the ONNX
    ModelProto's metadata_props field with the name of the object file as value.
    The object file can be found under the ip subdirectory.
    """

    def __init__(self, ip_name="finn_design"):
        super().__init__()
        self.ip_name = ip_name

    def apply(self, model):
        _check_vitis_envvars()
        vivado_proj_dir = model.get_metadata_prop("vivado_stitch_proj")
        stitched_ip_dir = vivado_proj_dir + "/ip"
        interfaces = json.loads(model.get_metadata_prop("vivado_stitch_ifnames"))
        args_string = []
        arg_id = 0
        # NOTE: this assumes the graph is Vitis-compatible: max one axi lite interface
        # developed from instructions in UG1393 (v2019.2) and package_xo documentation
        # package_xo is responsible for generating the kernel xml
        assert len(interfaces["axilite"]) <= 1, "CreateVitisXO supports max 1 AXI lite interface"
        axilite_intf_name = None
        if len(interfaces["axilite"]) == 1:
            axilite_intf_name = interfaces["axilite"][0]
            if len(interfaces["aximm"]) > 0:
                args_string.append(
                    "{addr:1:%s:%s:0x8:0x10:ap_uint&lt;%s>*:0}"
                    % (
                        str(arg_id),
                        interfaces["aximm"][0][0],
                        str(interfaces["aximm"][0][1]),
                    )
                )
                arg_id += 1
                args_string.append(
                    "{numReps:0:%s:%s:0x4:0x1C:uint:0}" % (str(arg_id), axilite_intf_name)
                )
                arg_id += 1
            else:
                args_string.append(
                    "{numReps:0:%s:%s:0x4:0x10:uint:0}" % (str(arg_id), axilite_intf_name)
                )
                arg_id += 1
        for intf in interfaces["s_axis"] + interfaces["m_axis"]:
            stream_width = intf[1]
            stream_name = intf[0]
            args_string.append(
                "{%s:4:%s:%s:0x0:0x0:ap_uint&lt;%s>:0}"
                % (stream_name, str(arg_id), stream_name, str(stream_width))
            )
            arg_id += 1

        # save kernel xml then run package_xo
        xo_name = self.ip_name + ".xo"
        xo_path = vivado_proj_dir + "/" + xo_name
        model.set_metadata_prop("vitis_xo", xo_path)

        # generate the package_xo command in a tcl script
        package_xo_string = "package_xo -force -xo_path %s -kernel_name %s -ip_directory %s" % (
            xo_path,
            self.ip_name,
            stitched_ip_dir,
        )
        for arg in args_string:
            package_xo_string += " -kernel_xml_args " + arg
        with open(vivado_proj_dir + "/gen_xo.tcl", "w") as f:
            f.write(package_xo_string)

        # create a shell script and call Vivado
        package_xo_sh = vivado_proj_dir + "/gen_xo.sh"
        working_dir = os.getcwd()
        with open(package_xo_sh, "w") as f:
            f.write("#!/bin/bash \n")
            f.write("cd {}\n".format(vivado_proj_dir))
            f.write("vivado -mode batch -source gen_xo.tcl\n")
            f.write("cd {}\n".format(working_dir))
        bash_command = ["bash", package_xo_sh]
        try:
            launch_process_helper(bash_command, print_stdout=False)
        except CalledProcessError:
            # Check success manually by looking for .xo file
            pass
        if not os.path.isfile(xo_path):
            raise FINNError("Vitis .xo file not created, check logs under %s" % vivado_proj_dir)

        return (model, False)


class VitisLink(Transformation):
    """Create an XCLBIN with Vitis.

    Outcome if successful: sets the bitfile attribute in the ONNX
    ModelProto's metadata_props field with the XCLBIN full path as value.
    """

    def __init__(
        self,
        platform,
        f_mhz=200,
        strategy=VitisOptStrategy.PERFORMANCE,
        enable_debug=False,
        fpga_memory_type="default",
    ):
        super().__init__()
        self.platform = platform
        self.f_mhz = f_mhz
        self.strategy = strategy
        self.enable_debug = enable_debug
        self.fpga_memory_type = fpga_memory_type

    def apply(self, model):
        _check_vitis_envvars()
        # create a config file and empty list of xo files
        config = ["[connectivity]"]
        object_files = []
        idma_idx = 0
        odma_idx = 0
        mem_idx = 0
        instance_names = {}
        for node in model.graph.node:
            assert node.op_type == "StreamingDataflowPartition", "Invalid link graph"
            sdp_node = getCustomOp(node)
            dataflow_model_filename = sdp_node.get_nodeattr("model")
            kernel_model = ModelWrapper(dataflow_model_filename)
            kernel_xo = kernel_model.get_metadata_prop("vitis_xo")
            object_files.append(kernel_xo)
            # gather info on connectivity
            # assume each node connected to outputs/inputs is DMA:
            # has axis, aximm and axilite
            # everything else is axis-only
            # assume only one connection from each ip to the next
            if len(node.input) == 0:
                producer = None
            else:
                producer = model.find_producer(node.input[0])
            consumer = model.find_consumers(node.output[0])
            # define kernel instances
            # name kernels connected to graph inputs as idmaxx
            # name kernels connected to graph inputs as odmaxx
            # TODO not a good way of checking for external in/out
            # check top-level in/out list instead
            if producer is None:
                instance_names[node.name] = "idma" + str(idma_idx)
                config.append("nk=%s:1:%s" % (node.name, instance_names[node.name]))
                idma_idx += 1
            elif consumer == []:
                instance_names[node.name] = "odma" + str(odma_idx)
                config.append("nk=%s:1:%s" % (node.name, instance_names[node.name]))
                odma_idx += 1
            else:
                instance_names[node.name] = node.name
                config.append("nk=%s:1:%s" % (node.name, instance_names[node.name]))
            sdp_node.set_nodeattr("instance_name", instance_names[node.name])
            # explicitly assign SLRs if the slr attribute is not -1
            node_slr = sdp_node.get_nodeattr("slr")
            if node_slr != -1:
                config.append("slr=%s:SLR%d" % (instance_names[node.name], node_slr))
            # assign memory banks
            if producer is None or consumer is None or consumer == []:
                node_mem_port = sdp_node.get_nodeattr("mem_port")
                if node_mem_port == "":
                    if self.fpga_memory_type == FpgaMemoryType.DEFAULT:
                        # configure good defaults based on board
                        if (
                            "u50" in self.platform
                            or "u280" in self.platform
                            or "u55c" in self.platform
                        ):
                            # Use HBM where available (also U50 does not have DDR)
                            mem_type = "HBM"
                            mem_idx = 0
                        elif "u200" in self.platform:
                            # Use DDR controller in static region of U200
                            mem_type = "DDR"
                            mem_idx = 1
                        elif "u250" in self.platform:
                            # Use DDR controller on the node's SLR if set, otherwise 0
                            mem_type = "DDR"
                            if node_slr == -1:
                                mem_idx = 0
                            else:
                                mem_idx = node_slr
                        else:
                            mem_type = "DDR"
                            mem_idx = 1
                    elif self.fpga_memory_type == FpgaMemoryType.HOST_MEM:
                        mem_type = "HOST"
                        mem_idx = 0
                    else:
                        raise RuntimeError(
                            "Unknown fpga memory type: "
                            + str(self.fpga_memory_type)
                            + ". Aborting!"
                        )
                    node_mem_port = "%s[%d]" % (mem_type, mem_idx)
                config.append("sp=%s.m_axi_gmem0:%s" % (instance_names[node.name], node_mem_port))
            # connect streams
            if producer is not None:
                for i in range(len(node.input)):
                    producer = model.find_producer(node.input[i])
                    if producer is not None:
                        j = list(producer.output).index(node.input[i])
                        config.append(
                            "stream_connect=%s.m_axis_%d:%s.s_axis_%d"
                            % (
                                instance_names[producer.name],
                                j,
                                instance_names[node.name],
                                i,
                            )
                        )

        # create a temporary folder for the project
        link_dir = make_build_dir(prefix="vitis_link_proj_")
        model.set_metadata_prop("vitis_link_proj", link_dir)

        # add Vivado physopt directives if desired
        if self.strategy == VitisOptStrategy.PERFORMANCE_BEST:
            config.append("[vivado]")
            config.append("prop=run.impl_1.STEPS.OPT_DESIGN.ARGS.DIRECTIVE=ExploreWithRemap")
            config.append("prop=run.impl_1.STEPS.PLACE_DESIGN.ARGS.DIRECTIVE=Explore")
            config.append("prop=run.impl_1.STEPS.PHYS_OPT_DESIGN.IS_ENABLED=true")
            config.append("prop=run.impl_1.STEPS.PHYS_OPT_DESIGN.ARGS.DIRECTIVE=Explore")
            config.append("prop=run.impl_1.STEPS.ROUTE_DESIGN.ARGS.DIRECTIVE=Explore")

        config = "\n".join(config) + "\n"
        with open(link_dir + "/config.txt", "w") as f:
            f.write(config)

        # create tcl script to generate resource report in XML format
        gen_rep_xml = templates.vitis_gen_xml_report_tcl_template
        gen_rep_xml = gen_rep_xml.replace("$VITIS_PROJ_PATH$", link_dir)
        with open(link_dir + "/gen_report_xml.tcl", "w") as f:
            f.write(gen_rep_xml)

        debug_commands = []
        if self.enable_debug:
            for inst in list(instance_names.values()):
                debug_commands.append("--dk chipscope:%s" % inst)

        # create a shell script and call Vitis
        script = link_dir + "/run_vitis_link.sh"
        working_dir = os.getcwd()
        with open(script, "w") as f:
            f.write("#!/bin/bash \n")
            f.write("cd {}\n".format(link_dir))
            f.write(
                "v++ -t hw --platform %s --link %s"
                " --kernel_frequency %d --config config.txt --optimize %s"
                " --save-temps -R2 %s\n"
                % (
                    self.platform,
                    " ".join(object_files),
                    self.f_mhz,
                    self.strategy.value,
                    " ".join(debug_commands),
                )
            )
            f.write("cd {}\n".format(working_dir))
        bash_command = ["bash", script]

        try:
            launch_process_helper(bash_command, print_stdout=False)
        except CalledProcessError:
            # Check success manually by looking for .xo file
            pass
        xclbin = link_dir + "/a.xclbin"
        if not os.path.isfile(xclbin):
            raise FINNError("Vitis .xclbin file not created, check logs under %s" % link_dir)

        # TODO rename xclbin appropriately here?
        model.set_metadata_prop("bitfile", xclbin)

        # run Vivado to gen xml report
        gen_rep_xml_sh = link_dir + "/gen_report_xml.sh"
        working_dir = os.getcwd()
        with open(gen_rep_xml_sh, "w") as f:
            f.write("#!/bin/bash \n")
            f.write("cd {}\n".format(link_dir))
            f.write("vivado -mode batch -source %s\n" % (link_dir + "/gen_report_xml.tcl"))
            f.write("cd {}\n".format(working_dir))
        bash_command = ["bash", gen_rep_xml_sh]
        launch_process_helper(bash_command, print_stdout=False)
        # filename for the synth utilization report
        synth_report_filename = link_dir + "/synth_report.xml"
        model.set_metadata_prop("vivado_synth_rpt", synth_report_filename)
        return (model, False)


class VitisBuild(Transformation):
    """Best-effort attempt at building the accelerator with Vitis.
    It assumes the model has only fpgadataflow nodes

    :parameter fpga_part: string identifying the target FPGA
    :parameter period_ns: target clock period
    :parameter platform: target Alveo platform, one of ["U50", "U200", "U250", "U280"]
    :parameter strategy: Vitis optimization strategy
    :parameter enable_debug: add Chipscope to all AXI interfaces
    :parameter floorplan_file: path to a JSON containing a dictionary with
        SLR assignments for each node in the ONNX graph.
        Must be parse-able by the ApplyConfig transform.
    :parameter enable_link: enable linking kernels (.xo files),
        otherwise just synthesize them independently.
    :parameter fpga_memory_type: Specify whether Host or FPGA memory such as DDR/HBM should be used
    """

    def __init__(
        self,
        fpga_part,
        period_ns,
        platform,
        strategy=VitisOptStrategy.PERFORMANCE,
        enable_debug=False,
        floorplan_file=None,
        enable_link=True,
        partition_model_dir=None,
        fpga_memory_type=FpgaMemoryType.DEFAULT,
    ):
        super().__init__()
        self.fpga_part = fpga_part
        self.period_ns = period_ns
        self.platform = platform
        self.strategy = strategy
        self.enable_debug = enable_debug
        self.floorplan_file = floorplan_file
        self.enable_link = enable_link
        self.partition_model_dir = partition_model_dir
        self.fpga_memory_type = fpga_memory_type

    def apply(self, model):
        _check_vitis_envvars()
        # prepare at global level, then break up into kernels
        prep_transforms = [InsertIODMA(512), InsertDWC(), SpecializeLayers(self.fpga_part)]
        for trn in prep_transforms:
            model = model.transform(trn)
            model = model.transform(GiveUniqueNodeNames())
            model = model.transform(GiveReadableTensorNames())

        model = model.transform(Floorplan(floorplan=self.floorplan_file))

        model = model.transform(
            CreateDataflowPartition(partition_model_dir=self.partition_model_dir)
        )
        model = model.transform(GiveUniqueNodeNames())
        model = model.transform(GiveReadableTensorNames())

        # Build each kernel individually
        sdp_nodes = model.get_nodes_by_op_type("StreamingDataflowPartition")
        for sdp_node in sdp_nodes:
            prefix = sdp_node.name + "_"
            sdp_node = getCustomOp(sdp_node)
            dataflow_model_filename = sdp_node.get_nodeattr("model")
            kernel_model = ModelWrapper(dataflow_model_filename)
            kernel_model = kernel_model.transform(InsertFIFO())
            kernel_model = kernel_model.transform(SpecializeLayers(self.fpga_part))
            kernel_model = kernel_model.transform(RemoveUnusedTensors())
            kernel_model = kernel_model.transform(GiveUniqueNodeNames(prefix))
            kernel_model.save(dataflow_model_filename)
            kernel_model = kernel_model.transform(PrepareIP(self.fpga_part, self.period_ns))
            kernel_model = kernel_model.transform(HLSSynthIP())
            kernel_model = kernel_model.transform(
                CreateStitchedIP(self.fpga_part, self.period_ns, sdp_node.onnx_node.name, True)
            )
            kernel_model = kernel_model.transform(CreateVitisXO(sdp_node.onnx_node.name))
            kernel_model.set_metadata_prop("platform", "alveo")
            kernel_model.save(dataflow_model_filename)
        # Assemble design from kernels
        if self.enable_link:
            model = model.transform(
                VitisLink(
                    self.platform,
                    round(1000 / self.period_ns),
                    strategy=self.strategy,
                    enable_debug=self.enable_debug,
                    fpga_memory_type=self.fpga_memory_type,
                )
            )
        # set platform attribute for correct remote execution
        model.set_metadata_prop("platform", "alveo")

        return (model, False)
