# fmt: off
# Disable formatter. This is deliberately formatted to stay within 80 characters
# per line. Black, however, formats some lines going beyond this.

# Numpy math and arrays
import numpy as np

# Operating system stuff, e.g. paths
import os

# Base class for specializing HW operators as implemented via HLS
from finn.custom_op.fpgadataflow.hlsbackend import HLSBackend
# The generic HW custom operator version of the operator as a base class
from finn.custom_op.fpgadataflow.attention_heads import (  # noqa
    MergeMultiHeads, SplitMultiHeads
)


# HLS Backend specialization of the multi-head attention splitting operator
class SplitMultiHeads_hls(  # noqa: Class name does not follow
    # CapWords convention
    SplitMultiHeads, HLSBackend
):
    # Node attributes matching the HLS operator
    def get_nodeattr_types(self):
        # Start from parent operator class attributes
        attrs = SplitMultiHeads.get_nodeattr_types(self)
        # Add the HLSBackend default attributes on top
        attrs.update(HLSBackend.get_nodeattr_types(self))
        # Add/Specialize implementation specific attributes here...
        # Return the updated attributes dictionary
        return attrs

    # Executes multi-head splitting in C++ simulation
    def _execute_node_cppsim(self, context, graph):  # noqa: graph unused
        # Get the node wrapped by this custom op  # noqa Duplicate
        node = self.onnx_node   # noqa Duplicate
        # Input data is stored in numpy files in the code generation dictionary
        code_gen_dir = self.get_nodeattr("code_gen_dir_cppsim")
        # Get the input out of the execution context
        #   Note: Shape must be either seq x 1 x dim or seq x dim
        inp = context[node.input[0]]
        # Validate the shape of the input
        assert inp.shape == self.get_normal_input_shape(ind=0), \
            f"Input shape mismatch for {node.input[0]}"
        # Reshape the input into folded form
        inp = inp.reshape(self.get_folded_input_shape(ind=0))
        # Save the folded inputs to file to be used by simulation
        np.save(os.path.join(code_gen_dir, "in.npy"), inp)

        # Execute the precompiled model
        super().exec_precompiled_singlenode_model()

        # Enumerate the node outputs
        for i, name in enumerate(node.output):
            # Load the output numpy file generated by the C++ simulation
            out = np.load(os.path.join(code_gen_dir, f"out{i}.npy"))
            # Reshape the folded output and insert into the execution context
            context[name] = out.reshape(self.get_normal_output_shape(ind=i))

    # Maximum width of any ap_int used in this operator
    def get_ap_int_max_w(self):
        # Find the widths of the widest input
        # Note: There is just one input.
        i_bits_max = self.get_instream_width(ind=0)
        # Find the widths of the widest output
        # Note: there is one output per head
        o_bits_max = max(
            (self.get_outstream_width(ind) for ind in range(self.heads))
        )
        # Find the biggest of the inputs/outputs
        return max([i_bits_max, o_bits_max])

    # Note: End of shape and datatype utilities

    # Generates list of C++ includes to be placed at the top of the generated
    # code
    def global_includes(self):
        # Currently nothing to include
        self.code_gen_dict["$GLOBALS$"] = []

    # Generates C++ code of type alias, global constant and macro definitions
    def defines(self, var):
        # Insert constants and type aliases into the dictionary
        self.code_gen_dict["$DEFINES$"] = [
            # Input and output element datatypes
            f"using IType = {self.dtype.get_hls_datatype_str()};",
            f"using OType = {self.dtype.get_hls_datatype_str()};",
            # Datatype of elements packed into the input stream
            f"using IPacked = ap_uint<{self.get_instream_width()}>;",
            # Datatype of elements packed into the output stream
            f"using OPacked = ap_uint<{self.get_outstream_width()}>;",
            # Input and output HLS stream datatypes
            "using IStream = hls::stream<"
            f"  ap_uint<{self.get_instream_width()}>"
            ">;",
            "using OStream = hls::stream<"
            f"  ap_uint<{self.get_outstream_width()}>"
            ">;",
        ]

    # Generates C++ code for reading data from .npy (numpy format) for testing
    # in C++ simulation
    def read_npy_data(self):
        # Input data is stored in numpy files in the code generation dictionary
        code_gen_dir = self.get_nodeattr("code_gen_dir_cppsim")
        # Generate function calls for reading the input files into the input
        # streams
        self.code_gen_dict["$READNPYDATA$"] = [
            # Generate function call reading from file into the input stream
            #   Note: Inputs are always represented as numpy floats
            'npy2apintstream<IPacked, IType, IType::width, float>(',
            f'"{code_gen_dir}/in.npy", in_{self.hls_sname()}, false',
            ');'
        ]

    # Generates C++ code for declaring all streams involved in C++ simulation
    # for testing
    def strm_decl(self):
        # Declare input and output streams
        # Note: Assumes stream type aliases to be set in defines
        self.code_gen_dict["$STREAMDECLARATIONS$"] = [
            # There is one input datastream
            f"IStream in_{self.hls_sname()};",
            # There is one output datastream per head
            *(f"OStream out{i}_{self.hls_sname()};" for i in range(self.heads))
        ]

    # Generates C++ code for calling the computation part of the operator
    def docompute(self):
        # Generates the bit-slicing indices string for the ith split of the
        # input
        def split(i):
            # Assemble a C++ indexing/bit-slicing string
            return f"({i + 1} * OPacked::width - 1, {i} * OPacked::width)"

        # Generates the name of the ith output stream
        def out(i):
            return f"out{i}_{self.hls_sname()}"

        # Write the body of the head-splitting top-level function
        self.code_gen_dict["$DOCOMPUTE$"] = [
            # Repeat for the number of inputs
            # Note: Repeat for all num_inputs dimensions
            f"for(std::size_t i = 0; i < {np.prod(self.num_inputs)}; ++i) {{",
            # Pipeline the steps of this loop
            "#pragma HLS pipeline II=1 style=flp",
            # Read the next input element from the stream
            f"const auto x = in_{self.hls_sname()}.read();",
            # Split the next element from the input stream into the number of
            # output elements per head and write into the corresponding stream
            *(f"{out(i)}.write(x{split(i)});" for i in range(self.heads)),
            # End of for-loop over repetitions body
            f"}}"  # noqa: f-string symmetry
        ]

    # Generates C++ code for reading the output stream and converting back to
    # numpy format for testing in C++ simulation
    def dataoutstrm(self):
        # Output data will be stored in numpy files in the  # noqa Duplicate
        # code generation dictionary
        code_gen_dir = self.get_nodeattr("code_gen_dir_cppsim")
        # Get the expected shape of the folded output array formatted as a C++
        # vector initializer
        # Note: Valid formatting relies on correct placement of curly braces
        # and line breaks: Open/close all three braces on the same line of code
        # to avoid '\n' to be inserted into the string
        shape = f"""{{{','.join((str(i) for i in self.get_folded_output_shape()))}}}"""
        # Start collecting function calls to write the output data stream
        self.code_gen_dict["$DATAOUTSTREAM$"] = []

        # Generates the name of the ith output stream
        def out(i):
            return f"out{i}_{self.hls_sname()}"

        # Generate code for each output stream
        for i in range(self.heads):
            # Append each reading/writing function call
            self.code_gen_dict["$DATAOUTSTREAM$"] += [
                # Generate function call reading from stream into the output
                # file
                #   Note: Outputs are always represented as numpy floats
                'apintstream2npy<OPacked, OType, OType::width, float>(',
                f'{out(i)}, {shape}, "{code_gen_dir}/out{i}.npy", false',
                ');'
            ]

    # Generates C++ code for saving the output of C++ simulation to a file in
    # numpy format
    def save_as_npy(self):
        # Note: This seems to be empty in ALL HLSCustomOps. Probably it was used
        # for something before, which is now integrated into dataoutstrm()?
        self.code_gen_dict["$SAVEASCNPY$"] = []

    # Generates essentially the head of the C++ function from which the IP block
    # will be generated during ipgen, i.e. actual synthesis
    def blackboxfunction(self):
        # Insert function head describing the top level interface of the head
        # splitting operator
        self.code_gen_dict["$BLACKBOXFUNCTION$"] = [
            # @formatter:off Prevent Python formatter from messing with C++
            # formatting
            # Note: Assumes stream type aliases to be set in defines
            f"void {self.onnx_node.name} (",
            # Input HLS stream
            f"IStream &in_{self.hls_sname()}, ", ",".join([
                # One output HLS stream per head  # noqa: Formatting
                f"OStream &out{i}_{self.hls_sname()}" for i in range(self.heads)
            ]),
            ")",
            # @formatter:off
        ]

    # Generates C++ pragmas to be inserted into the main function of the C++
    # simulation and the ipgen-blackboxfunction as well
    def pragmas(self):
        # Add HLS interface directives specifying how to create RTL ports for
        # the top-level function arguments
        self.code_gen_dict["$PRAGMAS$"] = [
            # Connect the input stream with an axi stream interface
            f"#pragma HLS INTERFACE axis port=in_{self.hls_sname()}"
        ]
        # Connect each output stream with an axi stream interface
        for i in range(self.heads):
            # Add new interface directive for the output stream
            self.code_gen_dict["$PRAGMAS$"] += [
                f"#pragma HLS INTERFACE axis port=out{i}_{self.hls_sname()}"
            ]
        # No block-level I/O protocol for the function return value
        self.code_gen_dict["$PRAGMAS$"].append(
            "#pragma HLS INTERFACE ap_ctrl_none port=return"
        )

    # Returns the names of input and output interfaces grouped by protocol
    def get_verilog_top_module_intf_names(self):
        # Start collecting interface names in a dictionary  # noqa Duplicate
        # starting with clock and reset
        intf_names = {"clk": ["ap_clk"], "rst": ["ap_rst_n"]}  # noqa
        # AXI stream input interfaces
        intf_names["s_axis"] = [
            # Just one input stream
            (f"in_{self.hls_sname()}", self.get_instream_width_padded(ind=0)),
        ]
        # AXI stream output interfaces
        intf_names["m_axis"] = [
            # One output stream per head
            (f"out{i}_{self.hls_sname()}",
             self.get_outstream_width_padded(ind=i)) for i in range(self.heads)
        ]
        # No AXI-MM, AXI-Lite or protocol-less interfaces
        intf_names["aximm"] = []
        intf_names["axilite"] = []
        intf_names["ap_none"] = []
        # Return the interface name dictionary
        return intf_names


# HLS Backend specialization of the multi-head attention merging operator
class MergeMultiHeads_hls(  # noqa: Class name does not follow
    # CapWords convention
    MergeMultiHeads, HLSBackend
):
    # Node attributes matching the HLS operator
    def get_nodeattr_types(self):
        # Start from parent operator class attributes
        attrs = MergeMultiHeads.get_nodeattr_types(self)
        # Add the HLSBackend default attributes on top
        attrs.update(HLSBackend.get_nodeattr_types(self))
        # Add/Specialize implementation specific attributes here...
        # Return the updated attributes dictionary
        return attrs

    # Executes multi-head slicing in C++ simulation
    def _execute_node_cppsim(self, context, graph):  # noqa: graph unused
        # Get the node wrapped by this custom op
        node = self.onnx_node
        # Input data is stored in numpy files in the code generation dictionary
        code_gen_dir = self.get_nodeattr("code_gen_dir_cppsim")

        # Enumerate the node outputs
        for i, name in enumerate(node.input):
            # Get the input out of the execution context
            #   Note: Shape must be either 1 x seq x dim or seq x dim
            inp = context[name]
            # Validate the shape of the input
            assert inp.shape == self.get_normal_input_shape(ind=i), \
                f"Input shape mismatch for {name}"
            # Reshape the input into folded form
            inp = inp.reshape(self.get_folded_input_shape(ind=i))
            # Save the folded inputs to file to be used by simulation
            np.save(os.path.join(code_gen_dir, f"in{i}.npy"), inp)

        # Execute the precompiled model
        super().exec_precompiled_singlenode_model()

        # Load the output numpy file generated by the C++ simulation
        out = np.load(os.path.join(code_gen_dir, "out.npy"))
        # Reshape the folded output and insert into the execution context
        context[node.output[0]] = out.reshape(
            self.get_normal_output_shape(ind=0)
        )

    # Maximum width of any ap_int used in this operator
    def get_ap_int_max_w(self):
        # Find the widths of the widest input
        # Note: There is just one input.
        i_bits_max = self.get_instream_width(ind=0)
        # Find the widths of the widest output
        # Note: there is one output per head
        o_bits_max = max(
            (self.get_outstream_width(ind) for ind in range(self.heads))
        )
        # Find the biggest of the inputs/outputs
        return max([i_bits_max, o_bits_max])

# Note: End of shape and datatype utilities

    # Generates list of C++ includes to be placed at the top of the generated
    # code
    def global_includes(self):
        # Currently nothing to include
        self.code_gen_dict["$GLOBALS$"] = []

    # Generates C++ code of type alias, global constant and macro definitions
    def defines(self, var):
        # Insert constants and type aliases into the dictionary
        self.code_gen_dict["$DEFINES$"] = [
            # Input and output element datatypes
            f"using IType = {self.dtype.get_hls_datatype_str()};",
            f"using OType = {self.dtype.get_hls_datatype_str()};",
            # Datatype of elements packed into the input stream
            f"using IPacked = ap_uint<{self.get_instream_width()}>;",
            # Datatype of elements packed into the output stream
            f"using OPacked = ap_uint<{self.get_outstream_width()}>;",
            # Input and output HLS stream datatypes
            "using IStream = hls::stream<"
            f"  ap_uint<{self.get_instream_width()}>"
            ">;",
            "using OStream = hls::stream<"
            f"  ap_uint<{self.get_outstream_width()}>"
            ">;",
        ]

    # Generates C++ code for reading data from .npy (numpy format) for testing
    # in C++ simulation
    def read_npy_data(self):
        # Input data is stored in numpy files in the code generation dictionary
        code_gen_dir = self.get_nodeattr("code_gen_dir_cppsim")
        # Generate function calls for reading the input files into the input
        # streams
        self.code_gen_dict["$READNPYDATA$"] = []
        # Generate code for each input stream
        for i in range(self.heads):
            # Append each reading/writing function call
            self.code_gen_dict["$READNPYDATA$"] += [
                # Generate function call reading from file into the input stream
                #   Note: Inputs are always represented as numpy floats
                'npy2apintstream<IPacked, IType, IType::width, float>(',
                f'"{code_gen_dir}/in{i}.npy", in{i}_{self.hls_sname()}, false',
                ');'
            ]

    # Generates C++ code for declaring all streams involved in C++ simulation
    # for testing
    def strm_decl(self):
        # Declare input and output streams
        # Note: Assumes stream type aliases to be set in defines
        self.code_gen_dict["$STREAMDECLARATIONS$"] = [
            # There is one output stream
            f"OStream out_{self.hls_sname()};",
            # There is one input stream per head
            *(f"IStream in{i}_{self.hls_sname()};" for i in range(self.heads))
        ]

    # Generates C++ code for calling the computation part of the operator
    def docompute(self):
        reversed_reads = ", ".join([
            f"in{i}_{self.hls_sname()}.read()"
            for i in reversed(range(self.heads))
        ])

        # Write the body of the head-splitting top-level function
        self.code_gen_dict["$DOCOMPUTE$"] = [
            # Repeat for the number of inputs
            # Note: Repeat for all num_inputs dimensions
            f"for(std::size_t i = 0; i < {np.prod(self.num_inputs)}; ++i) {{",
            # Pipeline the steps of this loop
            "#pragma HLS pipeline II=1 style=flp",
            # Read the next input element from each input stream and concatenate
            # using the comma operator overload of ap_uint, writing into the
            # output stream
            f"out_{self.hls_sname()}.write(({reversed_reads}));"
            # End of for-loop over repetitions body
            f"}}"  # noqa: f-string symmetry
        ]

    # Generates C++ code for reading the output stream and converting back to
    # numpy format for testing in C** simulation
    def dataoutstrm(self):
        # Output data will be stored in numpy files in the code generation
        # dictionary
        code_gen_dir = self.get_nodeattr("code_gen_dir_cppsim")
        # Get the expected shape of the folded output array formatted as a C++
        # vector initializer
        # Note: Valid formatting relies on correct placement of curly braces
        # and line breaks: Open/close all three braces on the same line of code
        # to avoid '\n' to be inserted into the string
        shape = f"""{{{','.join((str(i) for i in self.get_folded_output_shape()))}}}"""
        # Generate function call for reading from the output stream into the
        # output file
        self.code_gen_dict["$DATAOUTSTREAM$"] = [
            # Generate function call reading from stream into the output file
            #   Note: Outputs are always represented as numpy floats
            'apintstream2npy<OPacked, OType, OType::width, float>(',
            f'out_{self.hls_sname()}, {shape}, "{code_gen_dir}/out.npy", false',
            ');',
        ]

    # Generates C++ code for saving the output of C++ simulation to a file in
    # numpy format
    def save_as_npy(self):
        # Note: This seems to be empty in ALL HLSCustomOps. Probably it was used
        # for something before, which is now integrated into dataoutstrm()?
        self.code_gen_dict["$SAVEASCNPY$"] = []

    # Generates essentially the head of the C++ function from which the IP block
    # will be generated during ipgen, i.e. actual synthesis
    def blackboxfunction(self):
        # Insert function head describing the top level interface of the head
        # splitting operator
        self.code_gen_dict["$BLACKBOXFUNCTION$"] = [
            # @formatter:off Prevent Python formatter from messing with C++
            # formatting
            # Note: Assumes stream type aliases to be set in defines
            f"void {self.onnx_node.name} (",
            # Output HLS stream
            f"OStream &out_{self.hls_sname()}, ", ",".join([
                # One input HLS stream per head  # noqa: Formatting
                f"IStream &in{i}_{self.hls_sname()}" for i in range(self.heads)
            ]),
            ")",
            # @formatter:off
        ]

    # Generates C++ pragmas to be inserted into the main function of the C++
    # simulation and the ipgen-blackboxfunction as well
    def pragmas(self):
        # Add HLS interface directives specifying how to create RTL ports for
        # the top-level function arguments
        self.code_gen_dict["$PRAGMAS$"] = [
            # Connect the output stream with an axi stream interface
            f"#pragma HLS INTERFACE axis port=out_{self.hls_sname()}"
        ]
        # Connect each input stream with an axi stream interface
        for i in range(self.heads):
            # Add new interface directive for the input stream
            self.code_gen_dict["$PRAGMAS$"] += [
                f"#pragma HLS INTERFACE axis port=in{i}_{self.hls_sname()}"
            ]
        # No block-level I/O protocol for the function return value
        self.code_gen_dict["$PRAGMAS$"].append(
            "#pragma HLS INTERFACE ap_ctrl_none port=return"
        )

    # Returns the names of input and output interfaces grouped by protocol
    def get_verilog_top_module_intf_names(self):
        # Start collecting interface names in a dictionary starting with clock
        # and reset
        intf_names = {"clk": ["ap_clk"], "rst": ["ap_rst_n"]}  # noqa
        # AXI stream input interfaces
        intf_names["s_axis"] = [
            # One input stream per head
            (f"in{i}_{self.hls_sname()}",
             self.get_instream_width_padded(ind=i)) for i in range(self.heads)
        ]
        # AXI stream output interfaces
        intf_names["m_axis"] = [
            # Just one output stream
            (f"out_{self.hls_sname()}", self.get_outstream_width_padded(ind=0)),
        ]
        # No AXI-MM, AXI-Lite or protocol-less interfaces
        intf_names["aximm"] = []
        intf_names["axilite"] = []
        intf_names["ap_none"] = []
        # Return the interface name dictionary
        return intf_names
