# 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 Xilinx 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.

"""
Dataflow build configuration module for FINN.

This module provides configuration classes and enums for building dataflow
accelerators using the FINN framework. The main class DataflowBuildConfig
contains all the configuration options needed to control the dataflow build
process, including target hardware, optimization settings, verification options,
and output generation preferences.

The module also defines various enumerations for build options such as:
- LogLevel: Logging verbosity levels
- ShellFlowType: Different shell flows for FPGA integration
- DataflowOutputType: Types of outputs that can be generated
- VerificationStepType: Steps where verification can be performed

This configuration system allows users to customize the entire FINN dataflow
build process through a single, serializable configuration object.
"""
from __future__ import annotations

import numpy as np
import os
from dataclasses import dataclass
from enum import Enum
from mashumaro.mixins.json import DataClassJSONMixin
from mashumaro.mixins.yaml import DataClassYAMLMixin
from typing import Any, Literal, Optional

from finn.util.basic import alveo_default_platform, part_map


class LogLevel(str, Enum):
    """Log levels printed on the commandline for the build process."""

    NONE = "NONE"
    DEBUG = "DEBUG"
    INFO = "INFO"
    WARNING = "WARNING"
    ERROR = "ERROR"
    CRITICAL = "CRITICAL"


class AutoFIFOSizingMethod(str, Enum):
    "Select the type of automatic FIFO sizing strategy."

    CHARACTERIZE = "characterize"
    LARGEFIFO_RTLSIM = "largefifo_rtlsim"


class ShellFlowType(str, Enum):
    """For builds that produce a bitfile, select the shell flow that will integrate
    the FINN-generated accelerator."""

    VIVADO_ZYNQ = "vivado_zynq"
    VITIS_ALVEO = "vitis_alveo"


class DataflowOutputType(str, Enum):
    "Output product types that can be generated by build_dataflow"

    STITCHED_IP = "stitched_ip"
    ESTIMATE_REPORTS = "estimate_reports"
    OOC_SYNTH = "out_of_context_synth"
    RTLSIM_PERFORMANCE = "rtlsim_performance"
    BITFILE = "bitfile"
    PYNQ_DRIVER = "pynq_driver"
    CPP_DRIVER = "cpp_driver"
    DEPLOYMENT_PACKAGE = "deployment_package"


class VitisOptStrategy(Enum):
    "Values applicable to VitisBuild optimization strategy."

    DEFAULT = "0"
    POWER = "1"
    PERFORMANCE = "2"
    PERFORMANCE_BEST = "3"
    SIZE = "s"
    BUILD_SPEED = "quick"


class FpgaMemoryType(str, Enum):
    "Memory Type used by the FPGA to store input/output data"

    DEFAULT = "default"
    HOST_MEM = "host_memory"


class LargeFIFOMemStyle(str, Enum):
    """Type of memory resource to use for large FIFOs."""

    AUTO = "auto"
    BRAM = "block"
    LUTRAM = "distributed"
    URAM = "ultra"


class VerificationStepType(str, Enum):
    "Steps at which FINN ONNX execution can be launched for verification."

    #: verify after step_qonnx_to_finn, using Python execution
    QONNX_TO_FINN_PYTHON = "finn_onnx_python"
    #: verify after step_tidy_up, using Python execution
    TIDY_UP_PYTHON = "initial_python"
    #: verify after step_streamline , using Python execution
    STREAMLINED_PYTHON = "streamlined_python"
    #: verify after step_apply_folding_config, using C++ for each HLS node
    FOLDED_HLS_CPPSIM = "folded_hls_cppsim"
    #: verify after step_hw_ipgen
    NODE_BY_NODE_RTLSIM = "node_by_node_rtlsim"
    #: verify after step_create_stitched_ip, using stitched-ip Verilog
    STITCHED_IP_RTLSIM = "stitched_ip_rtlsim"


#: List of steps that will be run as part of the standard dataflow build, in the
#: specified order. Use the `steps` as part of build config to restrict which
#: steps will be run.
default_build_dataflow_steps = [
    "step_qonnx_to_finn",
    "step_tidy_up",
    "step_streamline",
    "step_convert_to_hw",
    "step_create_dataflow_partition",
    "step_specialize_layers",
    "step_target_fps_parallelization",
    "step_apply_folding_config",
    "step_minimize_bit_width",
    "step_generate_estimate_reports",
    "step_set_fifo_depths",
    "step_hw_codegen",
    "step_hw_ipgen",
    "step_create_stitched_ip",
    "step_measure_rtlsim_performance",
    "step_out_of_context_synthesis",
    "step_synthesize_bitfile",
    "step_make_driver",
    "step_deployment_package",
]

#: List of steps to run for an estimate-only (no synthesis) dataflow build
estimate_only_dataflow_steps = [
    "step_qonnx_to_finn",
    "step_tidy_up",
    "step_streamline",
    "step_convert_to_hw",
    "step_create_dataflow_partition",
    "step_specialize_layers",
    "step_target_fps_parallelization",
    "step_apply_folding_config",
    "step_minimize_bit_width",
    "step_generate_estimate_reports",
]

#: List of steps to run for a dataflow build including HW code generation, but
#: without any synthesis.
hw_codegen_dataflow_steps = estimate_only_dataflow_steps + ["step_hw_codegen"]


@dataclass
class DataflowBuildConfig(DataClassJSONMixin, DataClassYAMLMixin):
    """Build configuration to be passed to the build_dataflow function. Can be
    serialized into or de-serialized from JSON files for persistence.
    See list of attributes below for more information on the build configuration.
    """

    #: Directory where the final build outputs will be written into
    output_dir: str = "finn_build_output"

    #: Target clock frequency (in nanoseconds) for Vivado synthesis.
    #: e.g. synth_clk_period_ns=5.0 will target a 200 MHz clock.
    #: If hls_clk_period_ns is not specified it will default to this value.
    synth_clk_period_ns: float = 10.0

    #: Which output(s) to generate from the build flow.  See documentation of
    #: DataflowOutputType for available options.
    generate_outputs: Optional[list[DataflowOutputType]] = None

    #: (Optional) Path to configuration JSON file in which user can specify
    #: a preferred implementation style (HLS or RTL) for each node.
    #: The SpecializeLayers transformation picks up these settings and if possible
    #: fulfills the desired implementation style for each layer by converting the
    #: node into its HLS or RTL variant.
    #: Will be applied with :py:mod:`qonnx.transformation.general.ApplyConfig`
    specialize_layers_config_file: Optional[str] = None

    #: (Optional) Path to configuration JSON file. May include parallelization,
    #: FIFO sizes, RAM and implementation style attributes and so on.
    #: If the parallelization attributes (PE, SIMD) are part of the config,
    #: this will override the automatically generated parallelization
    #: attributes inferred from target_fps (if any)
    #: Will be applied with :py:mod:`qonnx.transformation.general.ApplyConfig`
    folding_config_file: Optional[str] = None

    #: (Optional) Target inference performance in frames per second.
    #: Note that target may not be achievable due to specific layer constraints,
    #: or due to resource limitations of the FPGA.
    #: If parallelization attributes are specified as part of folding_config_file
    #: that will override the target_fps setting here.
    target_fps: Optional[int] = None

    #: (Only relevant if target_fps is set) Use two-pass relaxation for folding.
    #: If enabled, parallelization will internally run a second time if the
    #: target cycles from the first pass could not be achieved, instead using the
    #: achievable target to obtain a balanced pipeline. If disabled, this can be
    #: useful for decreasing the latency (even though throughput won't increase).
    folding_two_pass_relaxation: bool = True

    #: (Optional) At which steps the generated intermediate output model
    #: will be verified. See documentation of VerificationStepType for
    #: available options.
    verify_steps: Optional[list[VerificationStepType]] = None

    #: (Only relevant if verify_steps is set)
    #: Name of .npy file that will be used as the input for verification.
    verify_input_npy: str = "input.npy"

    #: (Only relevant if verify_steps is set)
    #: Name of .npy file that will be used as the expected output for verification.
    verify_expected_output_npy: str = "expected_output.npy"

    #: (Only relevant if verify_steps is set)
    #: Save full execution context for each of the verify_steps.
    #: By default, only the top-level graph output is saved.
    #: This option also enables detailed activation tensor statistics logging.
    verify_save_full_context: bool = False

    #: (Only relevant if verify_steps is set or RTLSIM_PERFORMANCE output product is enabled)
    #: Save .vcd waveforms from rtlsim under reports.
    #: By default, waveforms won't be saved.
    verify_save_rtlsim_waveforms: bool = False

    #: Set verification tolerance: absolute error per output element.
    verification_atol: float = 1e-3

    #: Set verification tolerance: relative error per output element.
    #: This is added to the absolute tolerance (as computed by numpy.isclose).
    verification_rtol: float = 1e-5

    #: Set verification tolerance: mean absolute error over output elements of each frame.
    verification_mean_atol: float = 1e-3

    #: Set verification tolerance: mean relative error over output elements of each frame.
    verification_mean_rtol: float = 1e-1

    #: Run synthesis to generate a .dcp for the stitched-IP output product.
    #: This can make it easier to treat it as a standalone artifact without requiring
    #: the full list of layer IP build directories. By default, synthesis will not run.
    stitched_ip_gen_dcp: bool = False

    #: (Optional) Insert a signature node to the stitched-IP to read/write information
    #: to the design: e.g. Customer signature, application signature, version
    signature: Optional[list[int]] = None

    #: (Only relevant if target_fps is set) Control the maximum width of the per-PE MVAU stream
    #: while exploring the parallelization attributes to reach target_fps.
    #: Set this to a large value (e.g. 10000) if targeting full unfolding or
    #: very high performance.
    mvau_wwidth_max: int = 36

    #: Whether thresholding layers (which implement quantized
    #: activations in FINN) will be implemented as standalone HW layers,
    #: instead of being part of MatrixVectorActivation layer. This gives larger
    #: flexibility, and makes it possible to have runtime-writable thresholds.
    standalone_thresholds: bool = False

    #: Whether optimizations that minimize the bit width of the
    #: weights and accumulator will be applied. Because this optimization relies
    #: on the the values of the weights, it will only be applied if runtime-
    #: writeable weights is not enabled.
    minimize_bit_width: bool = True

    #: (Only needed for generating full bitfiles with shell integration)
    #: Target board, e.g. "Pynq-Z1" or "U250".
    board: Optional[str] = None

    #: (Only needed for generating full bitfiles with shell integration)
    #: Target shell flow, see documentation of ShellFlowType for options.
    shell_flow_type: Optional[ShellFlowType] = None

    #: (Only needed when board not specified) Target Xilinx FPGA part, e.g. "xc7z020clg400-1".
    fpga_part: Optional[str] = None

    #: Whether FIFO depths will be set automatically. Involves running stitched
    #: rtlsim and can take a long time.
    #: If set to False, the folding_config_file can be used to specify sizes
    #: for each FIFO.
    auto_fifo_depths: bool = True

    #: Enables experimental live FIFO sizing on the FPGA.
    live_fifo_sizing: bool = False

    #: Whether FIFO nodes with depth larger than 32768 will be split.
    #: Allow to configure very large FIFOs in the folding_config_file.
    split_large_fifos: bool = False

    #: (Only relevant when auto_fifo_depths is enabled)
    #: Select which method will be used for setting the FIFO sizes.
    auto_fifo_strategy: AutoFIFOSizingMethod = AutoFIFOSizingMethod.LARGEFIFO_RTLSIM

    #: (Only relevant when auto_fifo_depths is enabled)
    #: Memory resource type for large FIFOs.
    large_fifo_mem_style: LargeFIFOMemStyle = LargeFIFOMemStyle.AUTO

    #: (Only relevant if auto_fifo_strategy = LARGEFIFO_RTLSIM)
    #: Enable input throttling for simulation-based FIFO sizing.
    fifosim_input_throttle: bool = True

    #: (Only relevant if auto_fifo_strategy = LARGEFIFO_RTLSIM)
    #: Enable saving waveforms from simulation-based FIFO sizing.
    fifosim_save_waveform: bool = False

    #: (Only relevant if auto_fifo_strategy = LARGEFIFO_RTLSIM)
    #: Call CapConvolutionFIFODepths in InsertAndSetFIFODepths transform
    #: to make convolution FIFOs smaller where appropriate.
    default_swg_exception: bool = False

    #: (Optional) Target clock frequency (in nanoseconds) for Vitis HLS synthesis.
    #: e.g. `hls_clk_period_ns=5.0` will target a 200 MHz clock.
    #: If not specified it will default to synth_clk_period_ns
    hls_clk_period_ns: Optional[float] = None

    #: (Optional, only relevant when shell_flow_type = VITIS_ALVEO)
    #: Which Vitis platform will be used, e.g. "xilinx_u250_xdma_201830_2".
    #: If not specified but "board" is specified, will use the FINN
    #: default (if any) for that Alveo board.
    vitis_platform: Optional[str] = None

    #: (Optional, only relevant when shell_flow_type = VITIS_ALVEO)
    #: Path to JSON config file assigning each layer to an SLR.
    #: Will be applied with :py:mod:`qonnx.transformation.general.ApplyConfig`
    vitis_floorplan_file: Optional[str] = None

    #: (Only relevant when shell_flow_type = VITIS_ALVEO)
    #: Select Vitis optimization strategy.
    vitis_opt_strategy: VitisOptStrategy = VitisOptStrategy.DEFAULT

    #: (Only relevant when shell_flow_type = VITIS_ALVEO)
    #: Select FPGA memory type.
    #: Can be used to use host memory for input/output data instead of DDR or HBM memory
    fpga_memory: FpgaMemoryType = FpgaMemoryType.DEFAULT

    #: Whether intermediate ONNX files will be saved during the build process.
    #: These can be useful for debugging if the build fails.
    save_intermediate_models: bool = True

    #: When set to true, decorates every step in the build flow with a function
    #: that catches exceptions, snapshots the ONNX model, the config and the build log,
    #: saves them into the location of FINN_BUILD_DIR, and re-raises the exception for
    #: further error handling. By default this does _not_ save FINN itself; however
    #: any step can still be manually decorated to do so (see finn/utils/exception.py).
    enable_exception_snapshots: bool = False

    #: Whether hardware debugging will be enabled (e.g. ILA cores inserted to
    #: debug signals in the generated hardware).
    enable_hw_debug: bool = False

    #: Whether the accelerator will be simulated and synthesized with an
    #: instrumentation wrapper attached to accurately measure performance.
    enable_instrumentation: bool = False

    #: Whether pdb postmortem debugging will be launched when the build fails.
    enable_build_pdb_debug: bool = False

    #: When True, additional information (level = DEBUG) will be written to the log file.
    #: Otherwise, this additional information will be suppressed (level = INFO).
    verbose: bool = False

    #: Log level to be used on the command line for finn-plus internal logging.
    #: This is different from the log level used for build_dataflow.log,
    #: which is controlled using the verbose flag.
    console_log_level: LogLevel = LogLevel.ERROR

    #: (Optional) Only run the steps in the list. If not set, run default steps.
    #: See `default_build_dataflow_steps` for the default list of steps.
    #: When specified:
    #: Each item can either be a string, or a function (does not apply to json
    #: serialized configs) and does the following:
    #: - strings are resolved to functions from the default list
    #: - functions are called with (model, DataflowBuildConfig) as args
    steps: Optional[list[Any]] = None

    #: (Optional) Start build from this step, loading the intermediate model generated
    #: from the previous step (save_intermediate_models must be enabled).
    start_step: Optional[str] = None

    #: (Optional) Stop build at this step.
    stop_step: Optional[str] = None

    #: Control which Quant nodes of the QONNX format get converted to FINN's MultiThreshold nodes.
    #: This only affects Quant nodes in the activation path. Quant nodes, which define a
    #: bit width larger than `max_multithreshold_bit_width` are not converted to
    #: MultiThreshold nodes and a warning is raised instead.
    max_multithreshold_bit_width: int = 8

    #: Control the number of input frames for rtlsim performance measurement.
    rtlsim_batch_size: int = 1

    #: If set to True, FIFOs with impl_style=vivado will be kept during
    #: rtlsim, otherwise they will be replaced by RTL implementations.
    rtlsim_use_vivado_comps: bool = True

    #: (Only relevant if CPP_DRIVER output product is enabled) Selects C++ driver version.
    #: If set to "latest", newest version will be used.
    #: If set to commit hash, specified version will be used.
    cpp_driver_version: str = "latest"

    #: (Optional) Specify validation dataset to be used for deployment of the PYNQ driver.
    validation_dataset: Optional[str] = None

    #: (Only relevant if step_vivado_power_estimation is run)
    #: Whether to simulate the switching activity of the design for Vivado power estimation.
    #: If set to False, use only a fixed set of static toggle rates and static probabilities.
    #: If set to True, use simulated activity to generate an additional power report.
    vivado_power_simulate_activity: bool = True

    #: (Only relevant for step_vivado_power_estimation if vivado_power_simulate_activity is True)
    #: Whether to use "functional" or "timing" simulation for Vivado power estimation.
    vivado_power_simulation_type: Literal["timing", "functional"] = "functional"

    def _resolve_hls_clk_period(self) -> float:
        """
        Resolve the HLS clock period, falling back to synthesis clock period if not set.

        Returns:
            float: The HLS clock period in nanoseconds. If hls_clk_period_ns is not
                  specified, returns synth_clk_period_ns as the default.
        """
        if self.hls_clk_period_ns is None:
            # use same clk for synth and hls if not explicitly specified
            return self.synth_clk_period_ns
        else:
            return self.hls_clk_period_ns

    def _resolve_driver_platform(self) -> Literal["zynq-iodma", "alveo"]:
        """
        Resolve the driver platform based on the shell flow type.

        Returns:
            str: The driver platform identifier. Returns "zynq-iodma" for Vivado Zynq
                shell flows and "alveo" for Vitis Alveo shell flows.

        Raises:
            Exception: If the shell flow type is not recognized or supported.
        """
        if self.shell_flow_type == ShellFlowType.VIVADO_ZYNQ:
            return "zynq-iodma"
        elif self.shell_flow_type == ShellFlowType.VITIS_ALVEO:
            return "alveo"
        else:
            raise Exception("Couldn't resolve driver platform for " + str(self.shell_flow_type))

    def _resolve_fpga_part(self) -> str:
        """
        Resolve the FPGA part identifier.

        If fpga_part is explicitly specified in the configuration, it is returned as-is.
        If not specified, attempts to look up the part from the board name using the
        internal part_map dictionary.

        Returns:
            str: The FPGA part identifier (e.g., "xc7z020clg400-1").

        Raises:
            Exception: If the FPGA part cannot be resolved from the board name or
                      if the board is not found in the part map.
        """
        if self.fpga_part is None:
            # lookup from part map if not specified
            try:
                fpga_part = part_map[self.board]
                return fpga_part
            except KeyError:
                raise Exception("Couldn't resolve fpga_part for " + self.board)
        else:
            # return as-is when explicitly specified
            return self.fpga_part

    def _resolve_cycles_per_frame(self) -> None | int:
        """
        Calculate the number of clock cycles available per frame based on target FPS.

        Uses the target_fps and synth_clk_period_ns to compute how many clock cycles
        are available for processing each frame to achieve the target frame rate.

        Returns:
            int or None: The number of clock cycles per frame if target_fps is specified,
                        None if target_fps is not set.
        """
        if self.target_fps is None:
            return None
        else:
            n_clock_cycles_per_sec = 10**9 / self.synth_clk_period_ns
            n_cycles_per_frame = n_clock_cycles_per_sec / self.target_fps
            return int(n_cycles_per_frame)

    def _resolve_vitis_platform(self) -> str:
        """
        Resolve the Vitis platform identifier for Alveo board builds.

        If vitis_platform is explicitly specified, it is returned as-is. If not specified
        but a board is given, attempts to look up the default Vitis platform for that
        Alveo board from the internal alveo_default_platform mapping.

        Returns:
            str: The Vitis platform identifier (e.g., "xilinx_u250_xdma_201830_2").

        Raises:
            Exception: If neither vitis_platform nor board is specified, or if the
                      platform cannot be resolved from the given information.
        """
        if self.vitis_platform is not None:
            return self.vitis_platform
        elif (self.vitis_platform is None) and (self.board is not None):
            return alveo_default_platform[self.board]
        else:
            raise Exception(
                "Could not resolve Vitis platform:" " need either board or vitis_platform specified"
            )

    def _resolve_verification_steps(self) -> list[VerificationStepType]:
        """
        Resolve the list of verification steps to be performed during the build.

        Returns:
            List[VerificationStepType]: A list of verification steps to perform.
                                       Returns an empty list if verify_steps is None.
        """
        if self.verify_steps is None:
            return []
        else:
            return self.verify_steps

    def _resolve_verification_io_pair(self) -> None | tuple[Any, Any]:
        """
        Load and validate the input/output numpy arrays for verification.

        Loads the verification input and expected output arrays from the files
        specified in verify_input_npy and verify_expected_output_npy. Validates
        that both files exist before loading.

        Returns:
            tuple or None: A tuple containing (input_array, expected_output_array)
                          if verification is enabled, None if verify_steps is None.

        Raises:
            AssertionError: If either the input or expected output files cannot be found.
        """
        if self.verify_steps is None:
            return None
        else:
            assert os.path.isfile(self.verify_input_npy), (
                "verify_input_npy not found: " + self.verify_input_npy
            )
            verify_input_npy = np.load(self.verify_input_npy)
            assert os.path.isfile(self.verify_expected_output_npy), (
                "verify_expected_output_npy not found: " + self.verify_expected_output_npy
            )
            verify_expected_output_npy = np.load(self.verify_expected_output_npy)
            return (
                verify_input_npy,
                verify_expected_output_npy,
            )
