# Copyright (C) 2020 Arm Limited or its 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
#
# 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.


# Description:
# Functions used to write to a TensorFlow Lite format file. Supports adding in file identifiers.

import flatbuffers

from .tflite import Tensor
from .tflite import QuantizationParameters
from .tflite import Model
from .tflite import SubGraph
from .tflite import OperatorCode
from .tflite import Operator
from .tflite import Buffer
from .tflite import Metadata

import numpy as np

from .tflite_mapping import datatype_inv_map, builtin_operator_inv_map, custom_prefix, BuiltinOperator
from .nn_graph import PassPlacement
from .tensor import TensorPurpose, MemArea
from flatbuffers.builder import UOffsetTFlags

tflite_version = 3
tflite_file_identifier = "TFL" + str(tflite_version)


import flatbuffers.number_types as N
from flatbuffers import encode


def FinishWithFileIdentifier(self, rootTable, fid):
    if fid is None or len(fid) != 4:
        raise Exception("fid must be 4 chars")

    flags = N.Uint8Flags
    prepSize = 4
    self.Prep(self.minalign, prepSize + len(fid))
    for i in range(3, -1, -1):
        self.head = self.head - flags.bytewidth
        encode.Write(flags.packer_type, self.Bytes, self.Head(), ord(fid[i]))

    return self.Finish(rootTable)


flatbuffers.Builder.FinishWithFileIdentifier = FinishWithFileIdentifier


def make_vector(v):
    try:
        len(v)
        return v
    except TypeError:
        return [v]


class TFLiteSerialiser:
    def __init__(self, nng):
        self.builder = flatbuffers.Builder(0)
        self.nng = nng

        self.scratch_buf_id = 0  # Always assign scratch to buffer 0
        self.buffer_offsets_map = {}
        self.buffers_to_write = []  # have an empty array there

        self.input_tensors = []
        self.ops_to_ignore = set(("Const", "Placeholder", "SubgraphInput"))

        self.tensors_to_reshape = {}

        self.subgraphs_to_write = [sg for sg in self.nng.subgraphs if sg.placement == PassPlacement.Cpu]

        all_ops = []
        for sg in self.subgraphs_to_write:
            for ps in sg.passes:
                for op in ps.ops:
                    if op.type not in self.ops_to_ignore:
                        all_ops.append(op)
                    if op.type.startswith("Conv2D") or op.type.startswith("DepthwiseConv2d"):
                        self.tensors_to_reshape[op.inputs[1]] = (3, 0, 1, 2)
                    if op.type.startswith("FullyConnected"):
                        self.tensors_to_reshape[op.inputs[1]] = (1, 0)

        self.operator_codes = list(sorted(set(op.type for op in all_ops)))
        self.operator_code_map = {}

    def write_byte_vector(self, v, alignment=1):
        builder = self.builder
        builder.StartVector(1, len(v), alignment)
        for e in v[::-1]:
            builder.PrependByte(e)
        return builder.EndVector(len(v))

    def write_int_vector(self, v):
        builder = self.builder
        builder.StartVector(4, len(v), 4)
        for e in v[::-1]:
            builder.PrependInt32(e)
        return builder.EndVector(len(v))

    def write_long_vector(self, v):
        builder = self.builder
        builder.StartVector(8, len(v), 8)
        for e in v[::-1]:
            builder.PrependInt64(e)
        return builder.EndVector(len(v))

    def write_float_vector(self, v):
        builder = self.builder
        builder.StartVector(4, len(v), 4)
        for e in v[::-1]:
            builder.PrependFloat32(e)
        return builder.EndVector(len(v))

    def write_offset_vector(self, v):
        builder = self.builder
        builder.StartVector(4, len(v), 4)
        for e in v[::-1]:
            builder.PrependUOffsetTRelative(e)
        return builder.EndVector(len(v))

    def assign_buffers_to_tensors(self, tensors):
        buffer_map = {}
        scratch_tensor = [tens for tens in tensors if tens.purpose == TensorPurpose.Scratch][0]
        buf_idx = 1

        for tens in tensors:
            if tens.mem_area == scratch_tensor.mem_area:
                buffer_map[tens] = self.scratch_buf_id
            else:
                buffer_map[tens] = buf_idx
                buf_idx += 1

        # Initialize buffers_to_write to a length equal to numer of buffers so
        # they can be appended at the correct index during tensor serialization
        self.buffers_to_write = [None] * (buf_idx)

        return buffer_map

    def serialise_operator_code(self, idx, code):
        builder = self.builder
        custom_code_offset = None
        if code.startswith(custom_prefix):
            tf_code, opt_serializer = builtin_operator_inv_map[custom_prefix]
            custom_code_offset = builder.CreateString(code[len(custom_prefix) :])
        else:
            try:
                tf_code, opt_serializer = builtin_operator_inv_map[code]
            except KeyError:
                print(
                    "Warning: Writing operation %s, which does not have a direct TensorFlow Lite mapping, as a custom operation"
                    % (code,)
                )
                tf_code, opt_serializer = builtin_operator_inv_map[custom_prefix]

            if tf_code == BuiltinOperator.CUSTOM:
                assert code == "NpuOp"  # Currently only support serialising NPU operators as a custom op
                custom_code_offset = builder.CreateString("ethos-u")

            self.operator_code_map[code] = (idx, tf_code, opt_serializer)

        OperatorCode.OperatorCodeStart(builder)
        OperatorCode.OperatorCodeAddBuiltinCode(builder, tf_code)
        if custom_code_offset is not None:
            OperatorCode.OperatorCodeAddCustomCode(builder, custom_code_offset)

        return OperatorCode.OperatorCodeEnd(builder)

    def serialise_quantization_parameters(self, quant):
        builder = self.builder

        min = None
        max = None
        scale = None
        zero_point = None
        if quant is not None:
            if quant.min is not None:
                min = self.write_float_vector(make_vector(quant.min))
            if quant.max is not None:
                max = self.write_float_vector(make_vector(quant.max))
            if quant.scale_f32 is not None:
                scale = self.write_float_vector(make_vector(quant.scale_f32))
            if quant.zero_point is not None:
                zero_point = self.write_long_vector(make_vector(quant.zero_point))

        QuantizationParameters.QuantizationParametersStart(builder)
        if min is not None:
            QuantizationParameters.QuantizationParametersAddMin(builder, min)
        if max is not None:
            QuantizationParameters.QuantizationParametersAddMax(builder, max)
        if scale is not None:
            QuantizationParameters.QuantizationParametersAddScale(builder, scale)
        if zero_point is not None:
            QuantizationParameters.QuantizationParametersAddZeroPoint(builder, zero_point)
        return QuantizationParameters.QuantizationParametersEnd(builder)

    def serialise_tensor(self, tens):
        builder = self.builder
        tens_shape = tens.shape
        values = tens.quant_values
        if values is None:
            values = tens.values

        if values is None:
            values = np.empty(shape=(0), dtype=np.uint8)

        if tens in self.tensors_to_reshape:
            reorder = self.tensors_to_reshape[tens]
            tens_shape = [tens_shape[idx] for idx in reorder]
            values = values.transpose(reorder)

        if tens.purpose == TensorPurpose.Scratch:
            tens_shape = [0]
            self.buffers_to_write[self.scratch_buf_id] = values.flatten().view(np.uint8)

        buf_id = self.buffer_map[tens]
        if buf_id != self.scratch_buf_id:
            self.buffers_to_write[buf_id] = values.flatten().view(np.uint8)

        shape = self.write_int_vector(tens_shape)

        name = builder.CreateString(tens.name)
        quant = self.serialise_quantization_parameters(tens.quantization)

        Tensor.TensorStart(builder)
        Tensor.TensorAddShape(builder, shape)
        Tensor.TensorAddType(builder, datatype_inv_map[tens.dtype])
        # All tensors must have a valid backing buffer, even if it is empty.
        # Empty buffers should be kept unique for TensorFlow Lite Micro
        Tensor.TensorAddBuffer(builder, buf_id)
        Tensor.TensorAddName(builder, name)
        Tensor.TensorAddQuantization(builder, quant)

        res = Tensor.TensorEnd(builder)
        return res

    def serialise_operator(self, op):
        builder = self.builder

        inputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in op.inputs])
        outputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in op.outputs])

        op_idx, tflop, opt_serializer = self.operator_code_map[op.type]

        builtin_opt_offset = None
        custom_opt_offset = None
        if opt_serializer is not None:
            attrs = dict(op.attrs)
            if "strides" in attrs:
                attrs["stride_h"] = attrs["strides"][1]
                attrs["stride_w"] = attrs["strides"][2]
            if "ksize" in attrs:
                attrs["filter_height"] = attrs["ksize"][1]
                attrs["filter_width"] = attrs["ksize"][2]
            if "dilation" in attrs:
                attrs["dilation_h_factor"] = attrs["dilation"][1]
                attrs["dilation_w_factor"] = attrs["dilation"][2]
            if "channel_multiplier" in attrs:
                attrs["depth_multiplier"] = attrs["channel_multiplier"]

            builtin_opt_offset, custom_opt_offset = opt_serializer.serialize(builder, attrs)

        mutating_variable_inputs_offset = self.write_byte_vector([])
        Operator.OperatorStart(builder)
        Operator.OperatorAddOpcodeIndex(builder, op_idx)
        Operator.OperatorAddInputs(builder, inputs_offset)
        Operator.OperatorAddOutputs(builder, outputs_offset)

        if builtin_opt_offset is not None:
            Operator.OperatorAddBuiltinOptionsType(builder, opt_serializer.builtin_opt_type)
            Operator.OperatorAddBuiltinOptions(builder, builtin_opt_offset)
        if custom_opt_offset is not None:
            Operator.OperatorAddCustomOptions(builder, custom_opt_offset)
            Operator.OperatorAddCustomOptionsFormat(builder, opt_serializer.custom_opt_format)

        Operator.OperatorAddMutatingVariableInputs(builder, mutating_variable_inputs_offset)
        return Operator.OperatorEnd(builder)

    def serialise_subgraph(self, sg):
        builder = self.builder
        tensor_set = set()

        all_ops = []
        for ps in sg.passes:
            for op in ps.ops:
                if op.type not in self.ops_to_ignore:
                    all_ops.append(op)

        for op in all_ops:
            for tens in op.inputs + op.outputs:
                tensor_set.add(tens)

        all_tensors = [tens for nm, idx, tens in sorted((tens.name, idx, tens) for idx, tens in enumerate(tensor_set))]

        self.tensor_map = {tens: idx for idx, tens in enumerate(all_tensors)}
        self.buffer_map = self.assign_buffers_to_tensors(all_tensors)

        tensors_offset = self.write_offset_vector([self.serialise_tensor(tens) for tens in all_tensors])

        # Add the Scratch Tensor as input to the NPU subgraph to get it allocated by TensorFlow Lite Micro
        scratch_tensor_idx = [v for k, v in self.tensor_map.items() if k.name.endswith("scratch")]

        # Make sure the input_tensors haven't been modified
        assert all(inp in sg.original_inputs for inp in sg.input_tensors)
        inputs_offset = self.write_int_vector(
            [self.tensor_map[tens] for tens in sg.original_inputs] + scratch_tensor_idx
        )
        outputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in sg.output_tensors])

        operators_offset = self.write_offset_vector([self.serialise_operator(op) for op in all_ops])

        SubGraph.SubGraphStart(builder)
        SubGraph.SubGraphAddTensors(builder, tensors_offset)
        SubGraph.SubGraphAddInputs(builder, inputs_offset)
        SubGraph.SubGraphAddOutputs(builder, outputs_offset)

        SubGraph.SubGraphAddOperators(builder, operators_offset)

        return SubGraph.SubGraphEnd(builder)

    def write_aligned_bytes(self, buf):
        builder = self.builder
        builder.nested = True
        data = bytes(buf)
        length_bytes = UOffsetTFlags.py_type(len(data))
        builder.Prep(16, length_bytes)  # Reserve aligned storage
        builder.head = UOffsetTFlags.py_type(builder.Head() - length_bytes)  # Update FlatBuffer internal pointer
        builder.Bytes[builder.Head() : builder.Head() + length_bytes] = data  # Assign bytes to aligned area
        return builder.EndVector(length_bytes)

    def serialise_buffer(self, buf):
        builder = self.builder
        data = None
        if buf is not None:
            data = self.write_aligned_bytes(buf)
        Buffer.BufferStart(builder)
        if data is not None:
            Buffer.BufferAddData(builder, data)
        return Buffer.BufferEnd(builder)

    def serialise_metadata(self, metadata):
        builder = self.builder
        name = builder.CreateString(metadata[0])

        Metadata.MetadataStart(builder)
        Metadata.MetadataAddName(builder, name)
        Metadata.MetadataAddBuffer(builder, metadata[1])

        return Metadata.MetadataEnd(builder)

    def serialise_model(self):
        builder = self.builder
        operator_code_offset = self.write_offset_vector(
            [self.serialise_operator_code(idx, code) for idx, code in enumerate(self.operator_codes)]
        )

        description = builder.CreateString("Vela Optimised")

        subgraph_offset = self.write_offset_vector([self.serialise_subgraph(sg) for sg in self.subgraphs_to_write])

        # Fill the metadata buffer
        version = np.int32(0)
        subgraph_idx = np.int32(len(self.subgraphs_to_write))  # Only 1 supported currently
        nbr_tensors = np.int32(len(self.tensor_map))

        # An offset of -1 indicates that the tensor will be allocated online by Tensorflow Lite Micro
        offsets = [np.int32(-1)] * nbr_tensors

        # Ensure that the order of the offsets match the order of the tensors
        for tens, idx in self.tensor_map.items():
            if tens.mem_area == MemArea.Sram:
                offsets[idx] = np.int32(tens.address)

        metadata_buffer = np.array([version, subgraph_idx, nbr_tensors] + offsets)
        self.buffers_to_write.append(metadata_buffer)

        buffers_offset = self.write_offset_vector([self.serialise_buffer(buf) for buf in self.buffers_to_write])

        metadata_list = [("OfflineMemoryAllocation", len(self.buffers_to_write) - 1)]
        metadata_offset = self.write_offset_vector([self.serialise_metadata(metadata) for metadata in metadata_list])

        Model.ModelStart(builder)
        Model.ModelAddVersion(builder, tflite_version)
        Model.ModelAddOperatorCodes(builder, operator_code_offset)
        Model.ModelAddSubgraphs(builder, subgraph_offset)
        Model.ModelAddDescription(builder, description)
        Model.ModelAddBuffers(builder, buffers_offset)
        Model.ModelAddMetadata(builder, metadata_offset)
        return Model.ModelEnd(builder)

    def serialise(self):

        model = self.serialise_model()

        self.builder.FinishWithFileIdentifier(model, tflite_file_identifier)

        return self.builder.Output()

    def write(self, filename):
        with open(self.filename, "wb") as f:
            f.write(self.serialised_buf)


def write_tflite(nng, filename):
    writer = TFLiteSerialiser(nng)
    buf = writer.serialise()

    with open(filename, "wb") as f:
        f.write(buf)