# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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
#
#   http://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.
# pylint: disable=unused-argument
"""Relay identity operator for Arm(R) Ethos(TM)-U NPU"""

import tvm
from tvm.relay.op import _make
from tvm.topi.generic import schedule_injective
from tvm.relay.op.op import OpStrategy
from tvm.relay.op import strategy as _strategy

from ..te import identity_compute


@tvm.ir.register_op_attr("contrib.ethosu.identity", "FTVMCompute")
def create_ethosu_identity_compute(attrs, args, out_type):
    """Create an ethosu_identity compute op."""
    ifm = args[0]
    lut = args[1]
    ifm_scale = attrs.ifm_scale
    ifm_zero_point = attrs.ifm_zero_point
    ofm_scale = attrs.ofm_scale
    ofm_zero_point = attrs.ofm_zero_point
    activation = attrs.activation
    rounding_mode = attrs.rounding_mode
    op = identity_compute(
        ifm,
        lut,
        ifm_scale,
        ifm_zero_point,
        ofm_scale,
        ofm_zero_point,
        activation,
        rounding_mode,
    )
    return [op]


@tvm.ir.register_op_attr("contrib.ethosu.identity", "FTVMStrategy")
def identity_strategy_ethosu(attrs, inputs, out_type, target):
    strategy = OpStrategy()
    strategy.add_implementation(
        create_ethosu_identity_compute,
        _strategy.wrap_topi_schedule(schedule_injective),
        name="ethosu_identity",
    )
    return strategy


def ethosu_identity(
    ifm: tvm.relay.Expr,
    lut: tvm.relay.Expr,
    ifm_scale: float = 1,
    ifm_zero_point: int = 0,
    ofm_scale: float = 1,
    ofm_zero_point: int = 0,
    activation: str = "NONE",
    rounding_mode: str = "TFL",
) -> tvm.relay.Call:
    """The Identity operator that runs on the NPU.

    This operator takes in a tensor of any shape and returns the same tensor,
    with the data optionally requantized.

    Parameters
    ----------
    ifm : tvm.relay.Expr
        The Input Feature Map tensor (IFM).
    lut : tvm.relay.Expr
        The look-up table values to use if activation = "LUT", "TANH" or "SIGMOID".
    ifm_scale : float
        The quantization scale for the Input Feature Map tensor.
    ifm_zero_point : int
        The quantization zero point for the Input Feature Map tensor.
    ofm_scale : float
        The quantization scale for the Output Feature Map tensor.
    ofm_zero_point : int
       The quantization zero point for the Output Feature Map tensor.
    activation : str, optional
        The activation function to use.
            "NONE" - no activation function.
            "TANH" - tanh activation function.
            "SIGMOID" - sigmoid activation function.
            "LUT" - use a look-up table to perform the activation function.
    rounding_mode : str, optional
        The rounding mode to apply to the Output Feature Map tensor.
            "TFL" - Tensorflow Lite rounding scheme.
            "TRUNCATE" - Truncate towards zero.
            "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.

    Returns
    -------
    out : tvm.relay.Call
        A call to the ethosu_identity op.
    """
    return _make.ethosu_identity(
        ifm, lut, ifm_scale, ifm_zero_point, ofm_scale, ofm_zero_point, activation, rounding_mode
    )