# 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:
# Mark purpose and select formats for Tensors. Also compresses the weights.
from . import rewrite_graph
from . import weight_compressor
from .architecture_features import Block
from .operation import NpuBlockType
from .tensor import TensorFormat
from .tensor import TensorPurpose


def purpose_from_list(lst):
    def purpose(op, idx):
        return lst[idx]

    return purpose


def all_fm(op, idx):
    return TensorPurpose.FeatureMap


def all_parameter(op, idx):
    return TensorPurpose.FeatureMap


def input0_from_output_rest_parameter(op, idx):
    if idx == 0:
        res = op.outputs[0].purpose
        if res == TensorPurpose.Unknown:
            print("Warning: Propagating unknown tensor purpose", op)
        return res
    return TensorPurpose.FeatureMap


def inputs_from_output(op, idx):
    res = op.outputs[0].purpose
    if res == TensorPurpose.Unknown:
        print("Warning: Propagating unknown tensor purpose", op)
    return res


tensor_purposes = [  # ops, input_purpose
    (
        set(
            (
                "Relu",
                "Relu6",
                "Mul",
                "Add",
                "Sub",
                "Rsqrt",
                "Abs",
                "Cast",
                "Exp",
                "Floor",
                "FloorDiv",
                "FloorMod",
                "SquaredDifference",
                "AddN",
                "BiasAdd",
                "RealDiv",
                "Maximum",
                "Minimum",
                "Sigmoid",
                "Tanh",
                "FusedBatchNorm",
                "AvgPool",
                "MaxPool",
                "Squeeze",
                "Softmax",
                "LRN",
                "Assign",
                "BatchMatMul",
                "ZerosLike",
                "ExtractImagePatches",
                "MulAct",
                "AddAct",
                "SubAct",
                "DivAct",
                "AvgPoolAct",
                "MaxPoolAct",
                "LeakyRelu",
            )
        ),
        all_fm,
    ),
    (
        set(
            (
                "Conv2D",
                "DepthwiseConv2dNative",
                "MatMul",
                "Conv2DBiasAct",
                "DepthwiseConv2dBiasAct",
                "FullyConnectedAct",
            )
        ),
        purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.Weights, TensorPurpose.FeatureMap]),
    ),
    (
        set(("Conv2DBackpropInputSwitched",)),
        purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.Weights, TensorPurpose.FeatureMap]),
    ),
    (
        set(("QuantizedConv2D", "QuantizedMatMul")),
        purpose_from_list(
            [
                TensorPurpose.FeatureMap,
                TensorPurpose.Weights,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
            ]
        ),
    ),
    (
        set(
            (
                "Reshape",
                "Min",
                "Max",
                "Mean",
                "Pad",
                "MirrorPad",
                "ArgMax",
                "ArgMin",
                "ExpandDims",
                "ResizeNearestNeighbor",
                "ResizeBilinear",
                "Tile",
                "Transpose",
                "Mfcc",
            )
        ),
        purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]),
    ),
    (
        set(("QuantizedReshape", "QuantizedResizeBilinear")),
        purpose_from_list(
            [TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]
        ),
    ),
    (
        set(("QuantizedBiasAdd", "QuantizedAdd", "QuantizedMul")),
        purpose_from_list(
            [
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
            ]
        ),
    ),
    (
        set(
            (
                "Dequantize",
                "Quantize",
                "QuantizeV2",
                "QuantizedRelu",
                "QuantizedRelu1",
                "QuantizedRelu6",
                "QuantizedAvgPool",
                "QuantizedMaxPool",
                "Slice",
                "SplitV",
            )
        ),
        purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]),
    ),
    (
        set(("BatchToSpaceND", "SpaceToBatchND", "DepthToSpaceND", "SpaceToDepthND")),
        purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]),
    ),
    (
        set(("BlockLSTM",)),
        purpose_from_list(
            [
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.Weights,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
            ]
        ),
    ),
    (set(("SplitSliceRead",)), purpose_from_list([TensorPurpose.FeatureMap, TensorPurpose.FeatureMap])),
    (set(("Shape", "ConcatSliceWrite", "AudioSpectrogram")), purpose_from_list([TensorPurpose.FeatureMap])),
    (
        set(("StridedSlice",)),
        purpose_from_list(
            [TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap, TensorPurpose.FeatureMap]
        ),
    ),
    (set(("Fill", "Pack", "Range")), all_parameter),
    (
        set(("Requantize",)),
        purpose_from_list(
            [
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
                TensorPurpose.FeatureMap,
            ]
        ),
    ),
    (set(("Placeholder", "SubgraphInput", "Const", "VariableV2")), purpose_from_list([])),
    (set(("FakeQuantWithMinMaxArgs", "FakeQuantWithMinMaxVars")), input0_from_output_rest_parameter),
    (
        set(("Square", "Sqrt", "Log", "Less", "Enter", "Exit", "Identity", "StopGradient", "Merge", "Switch")),
        inputs_from_output,
    ),
    (None, all_fm),
]


for ops, input_purpose in tensor_purposes:
    if ops is None:
        continue
    for op in ops:
        assert len(op) > 1, "string literal has been decomposed"


def mark_tensor_purpose(nng, arch, verbose_tensor_purpose=False):
    def mark_tensor_helper(tens, purpose):

        if tens.purpose == TensorPurpose.Unknown or tens.purpose == purpose:
            tens.purpose = purpose
        else:
            assert 0, "Cannot resolve tensor purpose %s and %s for tensor %s" % (tens.purpose, purpose, tens)
        tens.mem_area = arch.tensor_storage_mem_area[tens.purpose]

        if len(tens.ops) == 1 and tens.ops[0].type == "Const":
            tens.mem_area = (
                arch.permanent_storage_mem_area
            )  # special case constants, as they must be in permanent storage

    def rewrite_mark_tensor_purpose(op, arch):
        # find disconnected outputs and mark as parameters
        for tens in op.outputs:
            if not tens.consumers():
                mark_tensor_helper(tens, TensorPurpose.FeatureMap)

        for ops, input_purpose in tensor_purposes:
            if ops is None or op.type in ops:
                if ops is None:
                    print(
                        "warning: don't know how to mark up purpose for",
                        op.type,
                        op.inputs,
                        "triggering all feature map fallback",
                    )
                for idx, tens in enumerate(op.inputs):
                    purpose = input_purpose(op, idx)
                    mark_tensor_helper(tens, purpose)
                if op.type == "Reshape":
                    # Reshape's input and output point to same data
                    op.outputs[0].mem_area = op.inputs[0].mem_area
                break
        return op

    for sg in nng.subgraphs:
        sg = rewrite_graph.rewrite_graph_pre_order(sg, arch, [], [rewrite_mark_tensor_purpose])
        for tens in sg.output_tensors:
            mark_tensor_helper(tens, TensorPurpose.FeatureMap)

    if verbose_tensor_purpose:
        nng.print_graph_with_tensors()

    return nng


reshape_operations = set(
    (
        "Reshape",
        "QuantizedReshape",
        "ExpandDims",
        "Squeeze",
        "BatchToSpaceND",
        "SpaceToBatchND",
        "DepthToSpaceND",
        "SpaceToDepthND",
        "Placeholder",
    )
)


def mark_tensor_format(nng, arch, verbose_tensor_format=False):
    formats_for_tensor = {}

    def init_tens(tens):
        if tens.purpose == TensorPurpose.FeatureMap:
            fmt = arch.default_feature_map_format
        elif tens.purpose == TensorPurpose.Weights:
            fmt = arch.default_weight_format
        else:
            assert 0, "unknown tensor purpose %s" % (tens.purpose,)
        return fmt

    def find_npu_usage_of_tensor(tens):
        for op in tens.consumers():
            if op.type == "DMA":
                return find_npu_usage_of_tensor(op.outputs[0])
            if "npu_block_type" in op.attrs:
                return op.attrs["npu_block_type"]
            return NpuBlockType.Default

    def visit_tens(tens, ps):
        if tens not in formats_for_tensor:
            fmt = init_tens(tens)
        else:
            fmt = formats_for_tensor[tens]

        formats_for_tensor[tens] = fmt

    for sg in nng.subgraphs:
        for ps in sg.passes:
            for tens in ps.outputs:
                visit_tens(tens, ps)
            for tens in ps.intermediates:
                visit_tens(tens, ps)
            for tens in ps.inputs:
                visit_tens(tens, ps)

    for tens, fmt in formats_for_tensor.items():
        tens.set_format(fmt, arch)
        if fmt == TensorFormat.WeightsCompressed and tens.values is not None:
            npu_block_type = find_npu_usage_of_tensor(tens)
            if len(tens.ops) == 1 and tens.ops[0].type == "DMA":
                weight_compressor.compress_weights(tens, arch, npu_block_type, Block(32, 32, 32), 32)
                # Alias compressed weights back into source tensor
                src_tens = tens.ops[0].inputs[0]
                src_tens.compressed_values = tens.compressed_values
                src_tens.storage_shape = tens.storage_shape
                src_tens.brick_size = tens.brick_size
                src_tens.weight_compression_scales = tens.weight_compression_scales
                src_tens.weight_compressed_offsets = tens.weight_compressed_offsets
                src_tens.compression_scale_for_worst_weight_stream = tens.compression_scale_for_worst_weight_stream
                src_tens.storage_compression_scale = tens.storage_compression_scale

    if verbose_tensor_format:
        nng.print_passes_with_tensors()