# 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:
# Internal representation of a Neural Network Operation.
from collections import namedtuple
from enum import Enum


class NpuBlockType(Enum):
    Default = 0
    ConvolutionMxN = 1
    VectorProduct = 2
    Pooling = 3
    ConvolutionDepthWise = 4
    ElementWise = 5
    ReduceSum = 6


# Classifies operators of type Custom
class CustomType(Enum):
    ThirdPartyOp = 0  # Third party custom op
    NpuOp = 1  # NPU op
    ExistingNpuOp = 2  # NPU op that was part of the input network


TensorIndices = namedtuple("TensorIndices", ["ifms", "weights", "biases"])

NO_INDICES = TensorIndices([], [], [])
IFM_INDICES = TensorIndices([0], [], [])
IFM_WEIGHTS_INDICES = TensorIndices([0], [1], [])
IFM_WEIGHTS_BIAS_INDICES = TensorIndices([0], [1], [2])
IFM_IFM2_INDICES = TensorIndices([0, 1], [], [])
CONV2D_BACKPROP_INDICES = TensorIndices([2], [1], [3])
TRANSPOSE_CONV_INDICES = TensorIndices([0], [1], [3])
CONCAT_INDICES = TensorIndices([1, 2], [], [])
SPLIT_IFM_INDICES = TensorIndices([1], [], [])
BLOCK_LSTM_INDICES = TensorIndices([3], [4], [])


# Static information related to operation codes
class OperatorInfo:
    __slots__ = ("id", "block_type", "indices", "is_unary")
    _id = 0

    def __init__(self, block_type=NpuBlockType.Default, indices=NO_INDICES, is_unary=False):
        OperatorInfo._id += 1
        self.id = OperatorInfo._id
        self.block_type = block_type
        self.indices = indices  # Indices of the different tensor purposes
        self.is_unary = is_unary  # Classifies elementwise operators


# Internally used operation codes
class Op(Enum):
    Abs = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=IFM_INDICES, is_unary=True)
    Add = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=IFM_IFM2_INDICES)
    AddN = OperatorInfo()
    Any = OperatorInfo()
    ArgMax = OperatorInfo()
    ArgMin = OperatorInfo()
    AvgPool = OperatorInfo(block_type=NpuBlockType.Pooling, indices=IFM_INDICES)
    BatchMatMul = OperatorInfo()
    BatchToSpaceND = OperatorInfo()
    BidirectionalSequenceLstm = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=IFM_WEIGHTS_INDICES)
    BidirectionalSequenceRnn = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=IFM_WEIGHTS_INDICES)
    BlockLSTM = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=BLOCK_LSTM_INDICES)

    CLZ = OperatorInfo(
        block_type=NpuBlockType.ElementWise, indices=IFM_INDICES, is_unary=True
    )  # NPU specific operation
    Call = OperatorInfo()
    Cast = OperatorInfo()
    Ceil = OperatorInfo()
    Concat = OperatorInfo(indices=CONCAT_INDICES)
    ConcatEmbeddings = OperatorInfo()
    ConcatSliceWrite = OperatorInfo(indices=IFM_INDICES)
    ConcatTFLite = OperatorInfo()
    Const = OperatorInfo()  # Constant tensor, only used in CPU subgraphs
    Conv2D = OperatorInfo(block_type=NpuBlockType.ConvolutionMxN, indices=IFM_WEIGHTS_INDICES)
    Conv2DBackpropInput = OperatorInfo(block_type=NpuBlockType.ConvolutionMxN, indices=CONV2D_BACKPROP_INDICES)
    Conv2DBackpropInputSwitchedBias = OperatorInfo(
        block_type=NpuBlockType.ConvolutionMxN, indices=TRANSPOSE_CONV_INDICES
    )
    Conv2DBias = OperatorInfo(block_type=NpuBlockType.ConvolutionMxN, indices=IFM_WEIGHTS_BIAS_INDICES)
    Cos = OperatorInfo()
    Custom = OperatorInfo()  # Custom 3rd party operator, only used in CPU subgraphs
    CustomNpuOp = OperatorInfo()  # NPU custom operator, only used in CPU subgraphs
    DMA = OperatorInfo()
    Delegate = OperatorInfo()
    Densify = OperatorInfo()
    DepthToSpace = OperatorInfo()
    DepthwiseConv2DBias = OperatorInfo(block_type=NpuBlockType.ConvolutionDepthWise, indices=IFM_WEIGHTS_BIAS_INDICES)
    Dequantize = OperatorInfo(indices=IFM_INDICES)
    Div = OperatorInfo()
    Elu = OperatorInfo()
    EmbeddingLookup = OperatorInfo()
    EmbeddingLookupSparse = OperatorInfo()
    Equal = OperatorInfo()
    Exp = OperatorInfo()
    ExpandDims = OperatorInfo(indices=IFM_INDICES)
    FakeQuantWithMinMaxArgs = OperatorInfo()
    Fill = OperatorInfo()
    Floor = OperatorInfo()
    FloorDiv = OperatorInfo()
    FloorMod = OperatorInfo()
    FullyConnected = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=IFM_WEIGHTS_BIAS_INDICES)
    GatherNd = OperatorInfo()
    GatherV2 = OperatorInfo()
    Greater = OperatorInfo()
    GreaterEqual = OperatorInfo()
    HardSwish = OperatorInfo()
    HashtableLookup = OperatorInfo()
    Identity = OperatorInfo()
    If = OperatorInfo()
    L2Norm = OperatorInfo()
    L2Pool2D = OperatorInfo()
    LRN = OperatorInfo()
    LSHProjection = OperatorInfo()
    LeakyRelu = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=IFM_INDICES, is_unary=True)
    Less = OperatorInfo()
    LessEqual = OperatorInfo()
    Log = OperatorInfo()
    LogSoftmax = OperatorInfo()
    LogicalAnd = OperatorInfo()
    LogicalNot = OperatorInfo()
    LogicalOr = OperatorInfo()
    Lstm = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=IFM_WEIGHTS_INDICES)
    LUT = OperatorInfo()  # NPU specific, operator has LUT, only used in fused activation functions
    MatMul = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=IFM_WEIGHTS_INDICES)
    MatrixDiag = OperatorInfo()
    MatrixSetDiag = OperatorInfo()
    Max = OperatorInfo()
    MaxPool = OperatorInfo(block_type=NpuBlockType.Pooling, indices=IFM_INDICES)
    Maximum = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=IFM_IFM2_INDICES)
    Mean = OperatorInfo()
    Min = OperatorInfo()
    Minimum = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=IFM_IFM2_INDICES)
    MirrorPad = OperatorInfo()
    Mul = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=IFM_IFM2_INDICES)
    Neg = OperatorInfo()
    NonMaxSuppressionV4 = OperatorInfo()
    NonMaxSuppressionV5 = OperatorInfo()
    NotEqual = OperatorInfo()
    OneHot = OperatorInfo()
    Pack = OperatorInfo()
    PackReshaped = OperatorInfo(indices=IFM_INDICES)
    Pad = OperatorInfo()
    PadV2 = OperatorInfo()
    Placeholder = OperatorInfo()  # Only used in CPU subgraphs
    Pow = OperatorInfo()
    Prelu = OperatorInfo()
    Prod = OperatorInfo()
    Quantize = OperatorInfo(indices=IFM_INDICES)
    QuantizedAvgPool = OperatorInfo(block_type=NpuBlockType.Pooling, indices=IFM_INDICES)
    QuantizedConv2D = OperatorInfo(block_type=NpuBlockType.ConvolutionMxN, indices=IFM_WEIGHTS_INDICES)
    QuantizedMatMul = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=IFM_WEIGHTS_INDICES)
    QuantizedMaxPool = OperatorInfo(block_type=NpuBlockType.Pooling, indices=IFM_INDICES)
    QuantizedReshape = OperatorInfo(indices=IFM_INDICES)
    Range = OperatorInfo()
    Rank = OperatorInfo()
    ReduceSum = OperatorInfo(block_type=NpuBlockType.ReduceSum, indices=IFM_INDICES)
    Relu = OperatorInfo(indices=IFM_INDICES)
    Relu6 = OperatorInfo(indices=IFM_INDICES)
    ReluN1To1 = OperatorInfo(indices=IFM_INDICES)
    Reshape = OperatorInfo(indices=IFM_INDICES)
    ResizeBilinear = OperatorInfo(block_type=NpuBlockType.Pooling, indices=IFM_INDICES)
    ResizeNearestNeighbor = OperatorInfo()
    ReverseSequence = OperatorInfo()
    ReverseV2 = OperatorInfo()
    Rnn = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=IFM_WEIGHTS_INDICES)
    Round = OperatorInfo()
    Rsqrt = OperatorInfo()
    SHL = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=IFM_IFM2_INDICES)  # NPU specific operation
    SHR = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=IFM_IFM2_INDICES)  # NPU specific operation
    ScatterNd = OperatorInfo()
    SegmentSum = OperatorInfo()
    Select = OperatorInfo()
    SelectV2 = OperatorInfo()
    Shape = OperatorInfo()
    Sigmoid = OperatorInfo(indices=IFM_INDICES)
    SignBit = OperatorInfo()
    Sin = OperatorInfo()
    SkipGram = OperatorInfo()
    Slice = OperatorInfo(indices=IFM_INDICES)
    Softmax = OperatorInfo()
    SpaceToBatchND = OperatorInfo()
    SpaceToDepth = OperatorInfo()
    SparseToDense = OperatorInfo()
    Split = OperatorInfo(indices=SPLIT_IFM_INDICES)
    SplitSliceRead = OperatorInfo(indices=IFM_INDICES)
    SplitV = OperatorInfo(indices=IFM_INDICES)
    Sqrt = OperatorInfo()
    Square = OperatorInfo()
    SquaredDifference = OperatorInfo()
    Squeeze = OperatorInfo(indices=IFM_INDICES)
    StridedSlice = OperatorInfo(indices=IFM_INDICES)
    Sub = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=IFM_IFM2_INDICES)
    SubgraphInput = OperatorInfo()  # Only used in CPU subgraphs
    Sum = OperatorInfo()
    Svdf = OperatorInfo()
    Tanh = OperatorInfo(indices=IFM_INDICES)
    Tile = OperatorInfo()
    TopKV2 = OperatorInfo()
    Transpose = OperatorInfo()
    UnidirectionalSequenceLstm = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=IFM_WEIGHTS_INDICES)
    UnidirectionalSequenceRnn = OperatorInfo(block_type=NpuBlockType.VectorProduct, indices=IFM_WEIGHTS_INDICES)
    Unique = OperatorInfo()
    Unpack = OperatorInfo()
    UnpackReshaped = OperatorInfo(indices=IFM_INDICES)
    Where = OperatorInfo()
    While = OperatorInfo()
    ZerosLike = OperatorInfo()

    @property
    def info(self):
        return self.value

    @property
    def npu_block_type(self):
        return self.info.block_type

    def is_conv2d_op(self):
        return self.info.block_type == NpuBlockType.ConvolutionMxN

    def is_depthwise_conv2d_op(self):
        return self.info.block_type == NpuBlockType.ConvolutionDepthWise

    def is_pool_op(self):
        return self.info.block_type == NpuBlockType.Pooling

    def is_maxpool_op(self):
        return self in (Op.MaxPool, Op.QuantizedMaxPool)

    def is_avgpool_op(self):
        return self in (Op.QuantizedAvgPool, Op.AvgPool)

    def is_elementwise_op(self):
        return self.info.block_type == NpuBlockType.ElementWise

    def is_unary_elementwise_op(self):
        return self.info.block_type == NpuBlockType.ElementWise and self.info.is_unary

    def is_binary_elementwise_op(self):
        return self.info.block_type == NpuBlockType.ElementWise and not self.info.is_unary

    def is_relu_op(self):
        return self in (Op.Relu, Op.Relu6, Op.ReluN1To1)

    def is_activation_op(self):
        return self.is_relu_op() or self in (Op.Tanh, Op.Sigmoid, Op.Softmax, Op.LUT)

    def is_split_op(self):
        return self in (Op.Split, Op.SplitV, Op.StridedSlice, Op.Slice, Op.UnpackReshaped)

    def is_concat_op(self):
        return self in (Op.Concat, Op.ConcatTFLite, Op.PackReshaped)

    def needs_bias(self):
        return bool(self.info.indices.biases)

    @classmethod
    def op_set(cls, predicate):
        # Returns the set of all operator codes that fulfill the given predicate
        return {op_type for op_type in Op if predicate(op_type)}

    def __str__(self):
        return self.name

    __repr__ = __str__

    def __lt__(self, other):
        return self.value.id < other.value.id


def create_avgpool_nop(name):
    op = Operation(Op.AvgPool, name)
    op.attrs["padding"] = b"VALID"
    op.attrs["stride_w"] = 1
    op.attrs["stride_h"] = 1
    op.attrs["filter_width"] = 1
    op.attrs["filter_height"] = 1
    op.attrs["strides"] = [1, 1, 1, 1]
    op.attrs["ksize"] = [1, 1, 1, 1]
    op.attrs["skirt"] = [0, 0, 0, 0]
    op.attrs["explicit_padding"] = [0, 0, 0, 0]
    return op


def get_slice_offsets(input_shape, offset_tens, offset_mask, is_begin=True):
    # For strided slice operator: get start or end offsets
    offsets = len(input_shape) * [0] if is_begin else input_shape[:]
    for idx in range(len(input_shape)):
        # If the i:th bit in the mask is set then the value on offset_tens[i] should be ignored
        if (offset_mask & (1 << idx)) == 0:
            offsets[idx] = offset_tens.values[idx]
            if offsets[idx] < 0:
                # Convert offset to positive value
                offsets[idx] += input_shape[idx]
    return offsets


class Operation:
    """Class representing a Neural Network operation. Has a name, a type,
    input and output tensors, as well as an attribute dictionary."""

    __slots__ = (
        "type",
        "name",
        "op_index",
        "attrs",
        "inputs",
        "outputs",
        "flops",
        "scheduled_pass",
        "run_on_npu",
        "activation",
        "memory_function",
        "forced_output_quantization",
        "activation_lut",
    )

    def __init__(self, op_type, name):
        self.type = op_type
        self.name = name
        self.attrs = {}
        self.inputs = []
        self.outputs = []
        self.flops = 0
        self.run_on_npu = True
        # Fused activation function. If not none: operator code.
        self.activation = None
        # Fused memory function, if not None: operator code
        self.memory_function = None
        # If not none: contains QuantizationParameters to be used as output quantization
        # (which overrides the ofm tensor's quantization), used in LUT
        self.forced_output_quantization = None
        self.scheduled_pass = None
        self.op_index = None  # input network operator index
        self.activation_lut = None

    def clone(self, suffix="_clone"):
        res = Operation(self.type, self.name + suffix)

        res.attrs = dict(self.attrs)
        res.inputs = list(self.inputs)
        res.outputs = list(self.outputs)
        res.flops = self.flops
        res.run_on_npu = self.run_on_npu
        res.activation = self.activation
        res.memory_function = self.memory_function
        res.forced_output_quantization = self.forced_output_quantization
        res.scheduled_pass = self.scheduled_pass
        res.op_index = None  # not relevant as not part of input network

        return res

    def __str__(self):
        return "<nng.Operation '{}' type={}>".format(self.name, self.type)

    __repr__ = __str__

    def get_ifm_ifm2_weights_ofm(self):
        return self.ifm, self.ifm2, self.weights, self.ofm

    def get_ifm_weights_biases_ofm(self):
        return self.ifm, self.weights, self.bias, self.ofm

    def get_ifm_ifm2_weights_biases_ofm(self):
        return self.ifm, self.ifm2, self.weights, self.bias, self.ofm

    def get_ifm_ofm(self):
        return self.ifm, self.ofm

    @property
    def ifm(self):
        # Gets the IFM tensor, or None if not applicable
        return self.get_input(self.type.info.indices.ifms, 0)

    @property
    def ifm2(self):
        # Gets the IFM2 tensor, or None if not applicable
        return self.get_input(self.type.info.indices.ifms, 1)

    @property
    def bias(self):
        # Gets the bias tensor, or None if not applicable
        return self.get_input(self.type.info.indices.biases, 0)

    @property
    def weights(self):
        # Gets the weight tensor, or None if not applicable
        return self.get_input(self.type.info.indices.weights, 0)

    def get_ifm_tensors(self):
        # Gets the IFM tensors, or empty list if not applicable
        return self._index_list_to_tensors(self.type.info.indices.ifms)

    def get_weight_tensors(self):
        # Gets the weight tensors, or empty list if not applicable
        return self._index_list_to_tensors(self.type.info.indices.weights)

    def get_bias_tensors(self):
        # Gets the bias tensors, or empty list if not applicable
        return self._index_list_to_tensors(self.type.info.indices.biases)

    def _index_list_to_tensors(self, index_list):
        return [self.inputs[ix] for ix in index_list if ix < len(self.inputs)]

    def get_input(self, index_list, ix):
        if ix >= len(index_list):
            return None
        if index_list[ix] >= len(self.inputs):
            return None
        return self.inputs[index_list[ix]]

    @property
    def ofm(self):
        # Gets the OFM tensor, or None if not applicable
        return self.outputs[0] if self.outputs else None

    def get_concat_inputs_axis(self):
        assert self.type.is_concat_op()

        if self.type == Op.Concat:
            axis_tensor = self.inputs[0]
            inputs = self.inputs[1:]
        elif self.type == Op.ConcatTFLite:
            inputs = self.inputs
            axis = self.attrs["axis"]
        elif self.type == Op.PackReshaped:
            # Requires fixup_pack_input to be called before this point
            inputs = self.inputs
            axis = self.attrs["axis"]
            assert len(self.inputs) == self.attrs["values_count"]
        else:
            assert len(axis_tensor.ops) == 1 and axis_tensor.ops[0].type == Op.Const
            axis = int(axis_tensor.values)

        return inputs, axis

    def get_dilation_h_w(self):
        _, dilation_h, dilation_w, _ = self.attrs.get("dilation", (1, 1, 1, 1))
        return dilation_h, dilation_w

    def get_split_inputs_axis(self):
        assert self.type.is_split_op()

        offset_start = None
        offset_end = None
        axis = None
        if self.type == Op.Split:
            num_splits = self.attrs.get("num_splits")
            axis_tens = self.inputs[0]
            assert len(axis_tens.ops) == 1 and axis_tens.ops[0].type == Op.Const
            axis = int(axis_tens.values)
            input_tens = self.inputs[1]
            outputs = self.outputs
            assert num_splits == len(outputs)

        elif self.type == Op.SplitV:
            num_splits = self.attrs.get("num_splits")
            input_tens = self.inputs[0]
            size_tens = self.inputs[1]
            assert len(size_tens.ops) == 1 and size_tens.ops[0].type == Op.Const
            sizes = size_tens.values

            axis_tens = self.inputs[2]
            assert len(axis_tens.ops) == 1 and axis_tens.ops[0].type == Op.Const
            axis = int(axis_tens.values)

            for idx, size in enumerate(sizes):
                # One but only one size might be set to -1, indicating that size should be inferred
                if size == -1:
                    sizes[idx] = input_tens.shape[axis] - (sum(sizes) + 1)
                    break

            outputs = self.outputs
            assert num_splits == len(outputs)
            assert sum(sizes) == input_tens.shape[axis]

        elif self.type == Op.Slice:
            input_tens, begin_tens, size_tens = self.inputs
            outputs = self.outputs
            offset_start = [0] * len(input_tens.shape)
            offset_end = [0] * len(input_tens.shape)

            for idx in range(len(begin_tens.values)):
                # Check if the op should slice in dimension idx
                if size_tens.values[idx] != input_tens.shape[idx]:
                    offset_start[idx] = begin_tens.values[idx]
                    offset_end[idx] = size_tens.values[idx] + offset_start[idx]

        elif self.type == Op.StridedSlice:
            input_tens, begin_tens, end_tens, strides_tens = self.inputs
            outputs = self.outputs
            out_tens = outputs[0]

            # Extract masks
            begin_mask = self.attrs["begin_mask"]
            ellipsis_mask = self.attrs["ellipsis_mask"]
            end_mask = self.attrs["end_mask"]
            new_axis_mask = self.attrs["new_axis_mask"]
            shrink_axis_mask = self.attrs["shrink_axis_mask"]

            # shrink_axis_mask/new_axis_mask/ellipsis_mask is not supported by the Operation class but the operation
            # may have the attribute modified and handled in the graph optimization phase.
            assert shrink_axis_mask == new_axis_mask == ellipsis_mask == 0
            assert len(input_tens.shape) == len(out_tens.shape)
            offset_start = get_slice_offsets(input_tens.shape, begin_tens, begin_mask, is_begin=True)
            offset_end = get_slice_offsets(input_tens.shape, end_tens, end_mask, is_begin=False)
        elif self.type == Op.UnpackReshaped:
            # Requires fixup_unpack_output to be called before this point
            input_tens = self.inputs[0]
            outputs = self.outputs
            axis = self.attrs["axis"]
            num_splits = self.attrs["num"]
            # Number of outputs have to equal the value of the dimension to unpack
            assert num_splits == len(outputs) == input_tens.shape[axis]
        else:
            assert False

        return input_tens, outputs, axis, offset_start, offset_end

    def set_activation_lut(self, lut_tensor):
        self.activation = Op.LUT
        self.activation_lut = lut_tensor
        self.add_input_tensor(lut_tensor)

    def add_input_tensor(self, tens):
        self.inputs.append(tens)
        if self not in tens.consumer_list:
            tens.consumer_list.append(self)

    def set_input_tensor(self, tens, idx):
        tens_to_remove = self.inputs[idx]
        if tens_to_remove in tens.consumer_list:
            tens.consumer_list.remove(tens_to_remove)

        self.inputs[idx] = tens
        if self not in tens.consumer_list:
            tens.consumer_list.append(self)

    def set_output_tensor(self, tens):
        tens.ops = [self]
        self.outputs = [tens]

    def get_output_quantization(self):
        if self.forced_output_quantization is not None:
            return self.forced_output_quantization
        return self.ofm.quantization