# 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=invalid-name, unused-argument, import-outside-toplevel, no-value-for-parameter
"""A set of passes to legalize some of operations for the NPU"""
from typing import List, Type, Callable
import math

import numpy as np  # type: ignore
from ethosu.vela import scaling, fp_math

import tvm  # type: ignore
from tvm import relay
from tvm.relay.dataflow_pattern import DFPatternCallback  # type: ignore
from tvm.relay.dataflow_pattern import wildcard
from tvm.relay.dataflow_pattern import is_op
from tvm.relay.dataflow_pattern import rewrite
from tvm.relay.dataflow_pattern import CallPattern
from tvm.relay.backend.contrib.ethosu import op as ethosu_ops  # type: ignore
from tvm.relay.backend.contrib.ethosu import vela_api
from tvm.relay.backend.contrib.ethosu import util
from tvm.relay.op.contrib import ethosu as ethosu_patterns  # type: ignore


class SplitRewriter(DFPatternCallback):
    """This rewriting converts split operations into a sequence of
    strided_slice operations, because codegen is going to be based
    on strided_slices that will define the slice of the tensor that
    will be fed to the consumer.
    """

    def __init__(self):
        super().__init__(require_type=True)
        self.split_in = wildcard()
        self.pattern = is_op("split")(self.split_in)

    @staticmethod
    def get_section_begin_coords(split: tvm.relay.Expr) -> List[int]:
        """Currently, the split operator takes an array of indices or an integer
        indicating the number of splits. However, its an array of indices could
        represent both cases, therefore this function just make it an array of
        indices where each index represent the co-ordinate of beginning of each
        section -- defines as section begins.

        Parameters
        ----------
        split : tvm.relay.Expr
            The Relay Call expression for a split operator

        Returns
        -------
        section_begins : List[int]
            A list containing integers corresponding to section
            begins
        """
        indices_or_sections = split.attrs.indices_or_sections
        input_shape = split.args[0].checked_type.shape
        split_axis = split.attrs.axis

        if isinstance(indices_or_sections, tvm.ir.container.Array):
            # 0 is the beginning of the first section.
            return [0] + list(indices_or_sections)
        split_axis_len = input_shape[split_axis].value
        section_length = split_axis_len // indices_or_sections.value
        return list(range(0, split_axis_len, section_length))

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        split_input = post.args[0]
        split_begins = list()
        split_ends = list()
        section_begins_in_split_axis = self.get_section_begin_coords(post)
        for split_cord in section_begins_in_split_axis:
            # first begin is [0, 0, ... , 0]
            begin_shape = [0 for i in range(len(split_input.checked_type.shape))]
            begin_shape[post.attrs.axis] = split_cord
            split_begins.append(begin_shape)

            end_shape = list(split_input.checked_type.shape)
            # Only the split axis coordinate changes
            end_shape[post.attrs.axis] = split_cord
            split_ends.append(end_shape)

        # Coordinates needs to be shifted left because beginning
        # of the next section is the end of the previous
        split_ends = split_ends[1:]
        # Last section end is the shape of the tensor itself.
        split_ends.append(list(split_input.checked_type.shape))

        strided_slices = list()
        for sb, se in zip(split_begins, split_ends):
            strided_slices.append(relay.strided_slice(split_input, sb, se))

        return relay.Tuple(strided_slices)


class PartitionedSplitRewriter(DFPatternCallback):
    """This pass brings the split out of the partitioned function"""

    def __init__(self):
        super().__init__(require_type=True, rewrite_once=True)
        self.pattern = (
            wildcard().has_attr({"Composite": ethosu_patterns.SplitParams.composite_name})
        )(wildcard())

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        split_input = post.args[0]
        split_params = ethosu_patterns.SplitParams(post.op.body)
        indices_or_sections = split_params.indices_or_sections
        axis = split_params.axis
        return relay.op.split(split_input, indices_or_sections, axis=axis).astuple()


def get_lut_from_func(
    ifm_scale: float,
    ifm_zp: int,
    ofm_scale: float,
    ofm_zp: int,
    func: Callable[[float], float],
) -> List[int]:
    """Calculates the values of the lookup table based on the calculation function"""

    lut_values = list()
    # Only int8 is currently supported
    dtype = np.int8
    qmin, qmax = np.iinfo(dtype).min, np.iinfo(dtype).max
    for x in range(qmin, qmax + 1):
        x_real = ifm_scale * (x - ifm_zp)
        out_real = func(x_real)
        lut_result = int(util.round_away_zero(ofm_zp + out_real / ofm_scale))
        lut_result = min(qmax, max(qmin, lut_result))
        lut_values.append(lut_result)

    return lut_values


class LutActivationRewriter(DFPatternCallback):
    """A class to create an identity operator with the LUT"""

    def __init__(
        self,
        params_class: Type,
        activation_type: str,
        calc_func: Callable[[float], float],
    ):
        super().__init__(require_type=True, rewrite_once=True)
        self.params_class = params_class
        self.pattern = (wildcard().has_attr({"Composite": params_class.composite_name}))(wildcard())
        self.activation_type = activation_type
        self.calc_func = calc_func

    def callback(self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map):
        params = self.params_class(post.op.body)
        params.ifm.tensor = post.args[0]

        input_scale = float(params.ifm.q_params.scale_f32)
        input_zp = int(params.ifm.q_params.zero_point)
        output_scale = float(params.ofm.q_params.scale_f32)
        output_zp = int(params.ofm.q_params.zero_point)

        lut_values = get_lut_from_func(
            input_scale,
            input_zp,
            output_scale,
            output_zp,
            self.calc_func,
        )
        lut = relay.const(lut_values, dtype=params.ifm.dtype)

        # We baked the requantization into the LUT, so we don't requantize the identity operator
        identity = ethosu_ops.ethosu_identity(
            ifm=params.ifm.tensor,
            lut=lut,
            ifm_scale=input_scale,
            ifm_zero_point=input_zp,
            ofm_scale=input_scale,
            ofm_zero_point=input_zp,
            activation=self.activation_type,
        )

        return identity


class TanhRewriter(LutActivationRewriter):
    """This pass adds tanh as a LUT to the identity operator"""

    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.TanhParams, activation_type="TANH", calc_func=math.tanh
        )


def sigmoid_calc_func(x: float) -> float:
    """Function to calculate the values for sigmoid"""
    # These limits are inherited from TFLite
    upper_limit = 8.0
    lower_limit = -8.0

    if x <= lower_limit:
        y = 0.0
    elif x >= upper_limit:
        y = 1.0
    else:
        y = 1 / (1 + math.exp(-x))
    return y


class SigmoidRewriter(LutActivationRewriter):
    """This pass adds sigmoid as a LUT for identity op"""

    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.SigmoidParams,
            activation_type="SIGMOID",
            calc_func=sigmoid_calc_func,
        )


def leaky_relu_calc_func(x: float, alpha: float) -> float:
    """Function to calculate the values for leaky relu."""
    return x if x >= 0 else x * alpha


class LeakyReLURewriter(DFPatternCallback):
    """This pass adds leaky relu as a LUT for identity op."""

    def __init__(self):
        super().__init__(require_type=True, rewrite_once=True)
        self.params_class = ethosu_patterns.LeakyReLUParams
        self.pattern = wildcard().has_attr({"Composite": self.params_class.composite_name})(
            wildcard()
        )

    def callback(self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map):
        params = self.params_class(post.op.body)
        params.ifm.tensor = post.args[0]

        input_scale = np.double(float(params.ifm.q_params.scale_f32))
        input_zp = int(params.ifm.q_params.zero_point)
        output_scale = np.double(float(params.ofm.q_params.scale_f32))
        output_zp = int(params.ofm.q_params.zero_point)

        alpha = params.alpha

        # The calculation of the LUT values is similar to that in Vela
        # convert_lrelu_to_lut(op, arch)
        # (https://review.mlplatform.org/plugins/gitiles/ml/ethos-u/ethos-u-vela/+/refs/tags/3.2.0/ethosu/vela/tflite_graph_optimiser.py#864)  # pylint: disable=line-too-long
        alpha_scalar = 1
        alpha_scale, alpha_shift = scaling.elementwise_mul_scale(input_scale, alpha, output_scale)
        identity_scale, identity_shift = scaling.elementwise_mul_scale(input_scale, 1, output_scale)

        dtype = params.ifm.dtype
        qmin, qmax = np.iinfo(dtype).min, np.iinfo(dtype).max

        def calculate_lut_value(i):
            zp_shift = (
                fp_math.multiply_by_quantized_multiplier(
                    alpha_scalar * (i - input_zp), alpha_scale, alpha_shift
                )
                if i < input_zp
                else fp_math.multiply_by_quantized_multiplier(
                    i - input_zp, identity_scale, identity_shift
                )
            )

            return min(qmax, max(qmin, output_zp + zp_shift))

        values = list(map(calculate_lut_value, range(qmin, qmax + 1)))
        lut = relay.const(values, dtype=dtype)

        # We baked the requantization into the LUT, so we don't requantize the identity operator
        identity = ethosu_ops.ethosu_identity(
            ifm=params.ifm.tensor,
            lut=lut,
            ifm_scale=input_scale,
            ifm_zero_point=input_zp,
            ofm_scale=input_scale,
            ofm_zero_point=input_zp,
            activation="LUT",
        )

        return identity


class HardSwishRewriter(DFPatternCallback):
    """Convert ethosu.hard_swish composite function to add operation with LUT."""

    def __init__(self):
        super().__init__(require_type=True, rewrite_once=True)
        self.params_class = ethosu_patterns.HardSwishParams
        self.pattern = wildcard().has_attr({"Composite": self.params_class.composite_name})(
            wildcard()
        )

    def callback(self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map):
        params = self.params_class(post.op.body)
        params.ifm.tensor = post.args[0]

        # The calculation of the LUT values is similar to that in Vela
        # convert_hardswish_to_lut(op, arch, nng)
        # (https://review.mlplatform.org/plugins/gitiles/ml/ethos-u/ethos-u-vela/+/refs/tags/3.2.0/ethosu/vela/tflite_graph_optimiser.py#719)  # pylint: disable=line-too-long
        input_scale = np.double(params.ifm.q_params.scale_f32)
        input_zp = int(params.ifm.q_params.zero_point)
        hires_input_scale = (1 / 128) * input_scale

        output_scale = np.double(params.ofm.q_params.scale_f32)
        output_zp = int(params.ofm.q_params.zero_point)
        output_scale, output_shift = scaling.quantise_scale(hires_input_scale / output_scale)
        output_scale_16 = fp_math.downscale_multiplier_int32_to_int16(output_scale)
        output_shift = 31 - output_shift
        output_shift = -output_shift if output_shift < 0 else 0

        dtype = params.ifm.dtype
        qmin, qmax = np.iinfo(dtype).min, np.iinfo(dtype).max

        def calculate_relu_multiplier(inp, input_scale):
            rmultiplier = np.double(3 / 32768)
            rscale, rshift = scaling.quantise_scale(input_scale / rmultiplier)
            rscale_16 = fp_math.downscale_multiplier_int32_to_int16(rscale)

            rvalue = np.int16(inp)
            if rshift < 31:
                rvalue = fp_math.shift_left16(rvalue, 30 - rshift)
                rvalue = fp_math.saturating_rounding_mul16(rvalue, rscale_16)
                rvalue = fp_math.shift_left16(rvalue, 1)
            elif rshift > 31:
                rvalue = fp_math.saturating_rounding_mul16(rvalue, rscale_16)
                rvalue = fp_math.rounding_divide_by_pot(rvalue, rshift - 31)
            else:
                rvalue = fp_math.saturating_rounding_mul16(rvalue, rscale_16)

            rvalue = (rvalue + (1 << 15)) >> 1
            return rvalue

        def calculate_lut_values(i):
            hires_input_value = (i - input_zp) * 128
            preshift_input_value = fp_math.saturating_rounding_mul16(
                hires_input_value, output_scale_16
            )
            relu_value = calculate_relu_multiplier(hires_input_value, hires_input_scale)
            lut_result = fp_math.saturating_mul16(relu_value, preshift_input_value)
            lut_result = fp_math.rounding_divide_by_pot(lut_result, output_shift) + output_zp
            return min(qmax, max(qmin, lut_result))

        values = list(map(calculate_lut_values, range(-128, 128)))
        lut = relay.const(values, dtype=dtype)

        # We baked the requantization into the LUT, so we don't requantize the identity operator
        identity = ethosu_ops.ethosu_identity(
            ifm=params.ifm.tensor,
            lut=lut,
            ifm_scale=input_scale,
            ifm_zero_point=input_zp,
            ofm_scale=input_scale,
            ofm_zero_point=input_zp,
            activation="LUT",
        )

        return identity


class Conv2DRewriter(DFPatternCallback):
    """Convert conv2d related composite functions into ethosu_conv2d operators"""

    def __init__(self):
        super().__init__(require_type=True)
        self.pattern = (wildcard().has_attr({"Composite": "ethos-u.qnn_conv2d"}))(wildcard())

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = ethosu_patterns.QnnConv2DParams(post.op.body)
        params.ifm.tensor = post.args[0]
        channels_map = {
            "NHWC": 3,
        }
        kernel_size_map = {
            "HWIO": params.weights.shape[0:2],
            "OHWI": params.weights.shape[1:3],
            "HWOI": params.weights.shape[0:2],
        }
        activation_map = {"clip": "CLIP"}
        weight_to_ohwi_transform_map = {"HWIO": [3, 0, 1, 2]}
        weights_values = params.weights.values
        weights_values_ohwi = np.transpose(
            weights_values, weight_to_ohwi_transform_map[str(params.weights.layout)]
        )
        if params.activation:
            activation = activation_map[params.activation.op.name]
            clip_min = int(params.activation.attrs.a_min)
            clip_max = int(params.activation.attrs.a_max)
        else:
            activation = "NONE"
            clip_min = 0
            clip_max = 0
        scale_bias = vela_api.pack_biases(
            biases=params.biases.tensor.data.asnumpy(),
            ifm_scale=params.ifm.q_params.scale_f32,
            ifm_dtype=np.dtype(params.ifm.dtype),
            weight_scales=params.weights.q_params.scale_f32,
            ofm_scale=params.ofm.q_params.scale_f32,
            is_activation_tanh_or_sigmoid=activation in ["TANH", "SIGMOID"],
        )
        ethosu_conv2d = ethosu_ops.ethosu_conv2d(
            ifm=post.args[0],
            weight=relay.const(weights_values_ohwi, params.weights.values.dtype),
            scale_bias=relay.const(scale_bias, "uint8"),
            lut=relay.const([], dtype="int8"),
            ifm_scale=float(params.ifm.q_params.scale_f32),
            ifm_zero_point=int(params.ifm.q_params.zero_point),
            weight_zero_point=int(params.weights.q_params.zero_point),
            ofm_scale=float(params.ofm.q_params.scale_f32),
            ofm_zero_point=int(params.ofm.q_params.zero_point),
            kernel_shape=kernel_size_map[str(params.weights.layout)],
            ofm_channels=params.ofm.shape[channels_map[str(params.ofm.layout)]],
            strides=params.strides,
            padding=params.padding,
            dilation=params.dilation,
            activation=activation,
            clip_min=clip_min,
            clip_max=clip_max,
            upscale="NONE",
            ifm_layout=str(params.ifm.layout),
            ofm_layout=str(params.ofm.layout),
        )
        return ethosu_conv2d


class Conv2DTransposeRewriter(DFPatternCallback):
    """Convert conv2d_transpose related composite functions into
    ethosu_conv2d_transpose operators."""

    def __init__(self):
        super().__init__(require_type=True)
        self.pattern = (wildcard().has_attr({"Composite": "ethos-u.qnn_conv2d_transpose"}))(
            wildcard()
        )

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = ethosu_patterns.QnnConv2DTransposeParams(post.op.body)
        params.ifm.tensor = post.args[0]

        ofm_shape = params.ofm.shape
        legalize_padding = params.legalize_padding

        weight_to_ohwi_transform_map = {"IOHW": [1, 2, 3, 0]}
        weights_values = params.weights.values
        weights_values_ohwi = np.transpose(
            weights_values, weight_to_ohwi_transform_map[str(params.weights.layout)]
        )
        weights_values_ohwi = np.flip(weights_values_ohwi, (1, 2))
        weights = relay.const(weights_values_ohwi, dtype=params.weights.values.dtype)

        bias_values = (
            params.biases.tensor.data.asnumpy()
            if params.biases
            else np.zeros((params.ifm.shape[-1]))
        )
        scale_bias = vela_api.pack_biases(
            biases=bias_values,
            ifm_scale=params.ifm.q_params.scale_f32,
            ifm_dtype=np.dtype(params.ifm.dtype),
            weight_scales=params.weights.q_params.scale_f32,
            ofm_scale=params.ofm.q_params.scale_f32,
            is_activation_tanh_or_sigmoid=False,
        )

        reduced_op = ethosu_ops.ethosu_conv2d(
            ifm=post.args[0],
            weight=weights,
            scale_bias=relay.const(scale_bias, "uint8"),
            lut=relay.const([], dtype="int8"),
            ifm_scale=float(params.ifm.q_params.scale_f32),
            ifm_zero_point=int(params.ifm.q_params.zero_point),
            weight_zero_point=int(params.weights.q_params.zero_point),
            ofm_scale=float(params.ofm.q_params.scale_f32),
            ofm_zero_point=int(params.ofm.q_params.zero_point),
            kernel_shape=params.kernel_shape,
            ofm_channels=int(ofm_shape[-1]),
            strides=(1, 1),
            padding=legalize_padding,
            dilation=params.dilation,
            ifm_layout=str(params.ifm.layout),
            ofm_layout=str(params.ofm.layout),
            upscale="ZEROS",
        )

        # Remove additional padding by 'cropping' back to expected size
        return relay.strided_slice(reduced_op, (0, 0, 0, 0), ofm_shape)


class DepthwiseConv2DRewriter(DFPatternCallback):
    """Convert ethosu.qnn_depthwise_conv2d composite functions to ethosu_depthwise_conv2d
    operators"""

    def __init__(self):
        super().__init__(require_type=True)
        self.pattern = (
            wildcard().has_attr(
                {"Composite": ethosu_patterns.QnnDepthwiseConv2DParams.composite_name}
            )
        )(wildcard())

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = ethosu_patterns.QnnDepthwiseConv2DParams(post.op.body)
        params.ifm.tensor = post.args[0]
        channels_map = {
            "NHWC": 3,
        }
        kernel_shape_map = {
            "HWOI": params.weights.shape[0:2],
        }

        weights_values = params.weights.values
        weights_values_ohwi = np.moveaxis(weights_values, [0, 1, 2, 3], [1, 2, 0, 3])

        activation = "NONE"
        # Activations requiring LUT is currently not supported, so setting it to an empty list
        lut = relay.const([], "int8")
        clip_min = 0
        clip_max = 0
        if params.activation:
            activation = ethosu_patterns.QnnDepthwiseConv2DParams.activation_map[
                params.activation.op.name
            ]
            if activation == "CLIP":
                clip_min = int(params.activation.attrs.a_min)
                clip_max = int(params.activation.attrs.a_max)
        scale_bias = vela_api.pack_biases(
            biases=params.biases.tensor.data.asnumpy(),
            ifm_scale=params.ifm.q_params.scale_f32,
            ifm_dtype=np.dtype(params.ifm.dtype),
            weight_scales=params.weights.q_params.scale_f32,
            ofm_scale=params.ofm.q_params.scale_f32,
            is_activation_tanh_or_sigmoid=activation in ["TANH", "SIGMOID"],
        )

        ethosu_depthwise_conv2d = ethosu_ops.ethosu_depthwise_conv2d(
            post.args[0],  # IFM
            relay.const(weights_values_ohwi, params.weights.values.dtype),
            relay.const(scale_bias, "uint8"),
            lut,
            float(params.ifm.q_params.scale_f32),
            int(params.ifm.q_params.zero_point),
            int(params.weights.q_params.zero_point),
            float(params.ofm.q_params.scale_f32),
            int(params.ofm.q_params.zero_point),
            kernel_shape_map[str(params.weights.layout)],
            params.ofm.shape[channels_map[str(params.ofm.layout)]],
            strides=params.strides,
            padding=params.padding,
            dilation=params.dilation,
            activation=activation,
            clip_min=clip_min,
            clip_max=clip_max,
            upscale="NONE",
            ifm_layout=str(params.ifm.layout),
            ofm_layout=str(params.ofm.layout),
            ofm_dtype=str(params.ofm.dtype),
        )
        return ethosu_depthwise_conv2d


class PoolingRewriter(DFPatternCallback):
    """Convert ethosu.avgpool2d and ethosu.maxpool2d composite functions to
    ethosu_pooling operators"""

    def __init__(
        self,
        params_class: Type,
        pattern: CallPattern,
    ):
        super().__init__(require_type=True)
        self.params_class = params_class
        self.pattern = pattern

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = self.params_class(post.op.body)
        params.ifm.tensor = post.args[0]
        channels_map = {
            "NHWC": 3,
        }

        activation_map = {"clip": "CLIP"}
        if params.activation:
            activation = activation_map[params.activation.op.name]
            clip_min = int(params.activation.attrs.a_min)
            clip_max = int(params.activation.attrs.a_max)
        else:
            activation = "NONE"
            clip_min = 0
            clip_max = 0

        # Activations requiring LUT is currently not supported, so setting it to an empty list
        lut = relay.const([], dtype="int8")

        return ethosu_ops.ethosu_pooling(
            ifm=post.args[0],
            lut=lut,
            pooling_type=params.pooling_type,
            ifm_scale=params.ifm.q_params.scale_f32,
            ifm_zero_point=params.ifm.q_params.zero_point,
            ofm_scale=params.ofm.q_params.scale_f32,
            ofm_zero_point=params.ofm.q_params.zero_point,
            pool_shape=params.pool_shape,
            ofm_channels=params.ofm.shape[channels_map[str(params.ofm.layout)]],
            strides=params.strides,
            padding=params.padding,
            activation=activation,
            clip_min=clip_min,
            clip_max=clip_max,
            upscale="NONE",
            ifm_layout=str(params.ifm.layout),
            ofm_layout=str(params.ofm.layout),
        )


class MaxPoolingRewriter(PoolingRewriter):
    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.MaxPool2DParams,
            pattern=(
                wildcard().has_attr({"Composite": ethosu_patterns.MaxPool2DParams.composite_name})
            )(wildcard()),
        )


class AvgPoolingRewriter(PoolingRewriter):
    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.AvgPool2DParams,
            pattern=(
                wildcard().has_attr({"Composite": ethosu_patterns.AvgPool2DParams.composite_name})
            )(wildcard()),
        )


class BinaryElementwiseRewriter(DFPatternCallback):
    """Convert ethosu binary elementwise composite functions to
    ethosu_binary_elementwise operators"""

    def __init__(
        self,
        params_class: Type,
        pattern: CallPattern,
    ):
        super().__init__(require_type=True)
        self.params_class = params_class
        self.pattern = pattern

    @staticmethod
    def reshape_input(
        inputs: List["TensorParams"],
    ) -> List[tvm.relay.Expr]:
        """Reshape the inputs so that the following binary elementwise
        operator receives 4-dimensional inputs.

        Parameters
        ----------
        inputs: List[TensorParams]
            The inputs to reshape.

        Returns
        -------
        reshaped_inputs: List[tvm.relay.Expr]
            The new reshaped inputs.
        """
        reshaped_inputs = []
        for i in inputs:
            in_shape = i.shape
            if len(in_shape) < 4:
                pad_size = 4 - len(in_shape)
                new_shape = ([1] * pad_size) + in_shape
                new_call = relay.reshape(i.tensor, new_shape)
                reshaped_inputs.append(new_call)
            else:
                reshaped_inputs.append(i.tensor)
        return reshaped_inputs

    @staticmethod
    def reshape_output(output: tvm.relay.Expr, ifm_input_shape: List[int]) -> tvm.relay.Expr:
        """Reshape the output back to the original dimensionality.
        Since the NPU must have the brodcastable tensor as the
        second operand, the original shape of the first ifm must
        be the output shape.

        Parameters
        ----------
        output: tvm.relay.Expr
            The output to reshape.

        ifm_input_shape: List[int]
            The shape of the non-reshaped ifm tensor.

        Returns
        -------
        reshaped_output: tvm.relay.Expr
            The reshaped output expression.
        """
        if len(ifm_input_shape) == 4:
            return output
        reshaped_output = relay.reshape(output, ifm_input_shape)
        return reshaped_output

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = self.params_class(post.op.body)
        params.ifm.tensor = post.args[1] if params.reversed_operands else post.args[0]
        params.ifm2.tensor = post.args[0] if params.reversed_operands else post.args[1]

        activation_map = {"clip": "CLIP"}
        if params.activation:
            activation = activation_map[params.activation.op.name]
            clip_min = int(params.activation.attrs.a_min)
            clip_max = int(params.activation.attrs.a_max)
        else:
            activation = "NONE"
            clip_min = 0
            clip_max = 0

        # We don't yet support activation functions that need to get legalized to LUTs.
        lut = relay.const([], dtype="int8")

        inputs = [params.ifm, params.ifm2]
        inputs = self.reshape_input(inputs)

        ethosu_binary_elementwise = ethosu_ops.ethosu_binary_elementwise(
            ifm=inputs[0],
            ifm2=inputs[1],
            lut=lut,
            operator_type=params.operator_type,
            ifm_scale=float(params.ifm.q_params.scale_f32),
            ifm_zero_point=int(params.ifm.q_params.zero_point),
            ifm2_scale=float(params.ifm2.q_params.scale_f32),
            ifm2_zero_point=int(params.ifm2.q_params.zero_point),
            ofm_scale=float(params.ofm.q_params.scale_f32),
            ofm_zero_point=int(params.ofm.q_params.zero_point),
            ifm_channels=params.ifm.shape[-1] if params.ifm.shape else 1,
            ifm2_channels=params.ifm2.shape[-1] if params.ifm2.shape else 1,
            reversed_operands=params.reversed_operands,
            ofm_dtype=params.ofm.dtype,
            activation=activation,
            clip_min=clip_min,
            clip_max=clip_max,
            ifm_layout=str(params.ifm.layout),
            ifm2_layout=str(params.ifm2.layout),
            ofm_layout=str(params.ofm.layout),
        )
        output = self.reshape_output(ethosu_binary_elementwise, params.ifm.shape)
        return output


class AddRewriter(BinaryElementwiseRewriter):
    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.AddParams,
            pattern=(wildcard().has_attr({"Composite": ethosu_patterns.AddParams.composite_name}))(
                wildcard(), wildcard()
            ),
        )


class SubRewriter(BinaryElementwiseRewriter):
    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.SubParams,
            pattern=(wildcard().has_attr({"Composite": ethosu_patterns.SubParams.composite_name}))(
                wildcard(), wildcard()
            ),
        )


class MulRewriter(BinaryElementwiseRewriter):
    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.MulParams,
            pattern=(wildcard().has_attr({"Composite": ethosu_patterns.MulParams.composite_name}))(
                wildcard(), wildcard()
            ),
        )


class MinRewriter(BinaryElementwiseRewriter):
    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.MinParams,
            pattern=(wildcard().has_attr({"Composite": ethosu_patterns.MinParams.composite_name}))(
                wildcard(), wildcard()
            ),
        )


class MaxRewriter(BinaryElementwiseRewriter):
    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.MaxParams,
            pattern=(wildcard().has_attr({"Composite": ethosu_patterns.MaxParams.composite_name}))(
                wildcard(), wildcard()
            ),
        )


class ShlRewriter(BinaryElementwiseRewriter):
    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.ShlParams,
            pattern=(wildcard().has_attr({"Composite": ethosu_patterns.ShlParams.composite_name}))(
                wildcard(), wildcard()
            ),
        )


class StridedSliceRewriter(DFPatternCallback):
    """This pass brings the strided slice out of the partitioned function"""

    def __init__(self):
        super().__init__(require_type=True, rewrite_once=True)
        self.pattern = (
            wildcard().has_attr({"Composite": ethosu_patterns.StridedSliceParams.composite_name})
        )(wildcard())

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:

        slice_input = post.args[0]

        # TODO(lhutton1) For an unknown reason compilation will fail for strides of 4
        # dimensions, so we cannot use params.strides as this will sometimes give
        # strides as [1, 1, 1, 1]. Since we only support strides of 1, hardcoding this
        # value for now.
        strides = [1]

        params = ethosu_patterns.StridedSliceParams(post.op.body)
        strided_slice = relay.op.strided_slice(
            slice_input,
            params.begin,
            params.end,
            strides=strides,
            axes=params.axes,
            slice_mode=params.slice_mode,
        )
        return strided_slice


class ReshapeRewriter(DFPatternCallback):
    """This pass brings the reshape out of the partitioned function"""

    def __init__(self):
        super().__init__(require_type=True, rewrite_once=True)
        self.pattern = (
            wildcard().has_attr({"Composite": ethosu_patterns.ReshapeParams.composite_name})
        )(wildcard())

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        reshape_input = post.args[0]
        reshape_params = ethosu_patterns.ReshapeParams(post.op.body)
        new_shape = reshape_params.new_shape
        return relay.op.reshape(reshape_input, newshape=new_shape)


class NoOpRewriter(DFPatternCallback):
    """This pass adds an idenity operator to reshape and strided slice to avoid a no op
    without a consumer"""

    def __init__(self):
        super().__init__(require_type=True, rewrite_once=True)
        self.reshape = is_op("reshape")(wildcard())
        self.strided_slice = is_op("strided_slice")(wildcard())
        self.pattern = self.reshape | self.strided_slice

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        if pre.checked_type.dtype == "int32":
            return post
        return ethosu_ops.ethosu_identity(ifm=post, lut=relay.const([], dtype="int8"))


class UnaryElementwiseRewriter(DFPatternCallback):
    """
    Convert ethosu unary elementwise composite function to
    ethosu_unary_elementwise operators
    """

    def __init__(self, params_class: Type, pattern: CallPattern):
        super().__init__(require_type=True)
        self.params_class = params_class
        self.pattern = pattern

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = self.params_class(post.op.body)
        params.ifm.tensor = post.args[0]

        activation_map = {"clip": "CLIP"}
        if params.activation:
            activation = activation_map[params.activation.op.name]
            clip_min = int(params.activation.attrs.a_min)
            clip_max = int(params.activation.attrs.a_max)
        else:
            activation = "NONE"
            clip_min = 0
            clip_max = 0

        # We don't yet support activation functions that use LUT.
        lut = relay.const([], dtype="int8")

        unary_input_shape = params.ifm.shape
        # If the input tensor is not 4D, enter reshapes before and after the unary operator
        if len(params.ifm.shape) == 4:
            unary_input = params.ifm.tensor
        else:
            pad_size = 4 - len(unary_input_shape)
            unary_input_shape = ([1] * pad_size) + unary_input_shape
            unary_input = relay.op.reshape(params.ifm.tensor, newshape=unary_input_shape)

        ethosu_unary_elementwise = ethosu_ops.ethosu_unary_elementwise(
            ifm=unary_input,
            lut=lut,
            operator_type=params.operator_type,
            ifm_scale=float(params.ifm.q_params.scale_f32),
            ifm_zero_point=int(params.ifm.q_params.zero_point),
            ofm_scale=float(params.ofm.q_params.scale_f32),
            ofm_zero_point=int(params.ofm.q_params.zero_point),
            ofm_channels=unary_input_shape[3],
            activation=activation,
            clip_min=clip_min,
            clip_max=clip_max,
            ifm_layout=str(params.ifm.layout),
            ofm_layout=str(params.ofm.layout),
        )
        if len(params.ifm.shape) == 4:
            op = ethosu_unary_elementwise
        else:
            op = relay.op.reshape(ethosu_unary_elementwise, newshape=params.ifm.shape)
        return op


class AbsRewriter(UnaryElementwiseRewriter):
    def __init__(self):
        super().__init__(
            params_class=ethosu_patterns.AbsParams,
            pattern=(wildcard().has_attr({"Composite": ethosu_patterns.AbsParams.composite_name}))(
                wildcard()
            ),
        )


class MeanRewriter(DFPatternCallback):
    """Convert ethosu.mean composite functions to an equivalent legalization:
    - Case 1 (axis == [1, 2] and keepsdims == True):
        ethosu_depthwise_conv2d + ethosu_binary_elementwise
    - Case 2 (ifm qparams == ofm qparams): ethosu_pooling
    - Case 3 (else): ethosu_depthwise_conv2d
    """

    def __init__(self):
        super().__init__(require_type=True)
        self.pattern = (
            wildcard().has_attr({"Composite": ethosu_patterns.MeanParams.composite_name})
        )(wildcard())

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = ethosu_patterns.MeanParams(post.op.body)
        params.ifm.tensor = post.args[0]

        ifm_shape = params.ifm.shape
        ofm_shape = params.ofm.shape
        lut = relay.const([], "int8")
        axis = params.axis
        reduced_op = params.ifm.tensor

        # Enforce 4d input
        if len(ifm_shape) < 4:
            axis = [x + 1 for x in axis]
            if len(ifm_shape) == 3:
                ifm_shape = [1, params.height, params.width, ifm_shape[2]]
            else:
                ifm_shape = [1, params.height, params.width, 1]
            reduced_op = relay.reshape(reduced_op, ifm_shape)

        filter_height = ifm_shape[1] if 1 in axis else 1
        filter_width = ifm_shape[2] if 2 in axis else 1
        in_channels = out_channels = ifm_shape[-1]

        # If the height is greater than max kernel height, reshape the input
        # from [filter_height, filter_width] to [1, (filter_height*filter_width)]
        # only in the case the axis is [1, 2].
        if axis == [1, 2] and filter_height > 64:
            ifm_shape = (ifm_shape[0], 1, filter_height * filter_width, in_channels)
            filter_width = filter_height * filter_width
            filter_height = 1
            reduced_op = relay.reshape(reduced_op, ifm_shape)

        if axis == [1, 2] and params.keepdims:
            weight_scale = 1
            weight_values = np.ones([out_channels, filter_height, filter_width, 1])
            scale_bias = vela_api.pack_biases(
                biases=np.zeros(ifm_shape[-1]),
                ifm_scale=params.ifm.q_params.scale_f32,
                ifm_dtype=np.dtype(params.ifm.dtype),
                weight_scales=np.array([weight_scale], dtype=np.float),
                ofm_scale=params.ofm.q_params.scale_f32,
                is_activation_tanh_or_sigmoid=False,
            )

            reduced_op = ethosu_ops.ethosu_depthwise_conv2d(
                ifm=reduced_op,
                weight=relay.const(weight_values, params.ifm.dtype),
                scale_bias=relay.const(scale_bias, "uint8"),
                lut=lut,
                ifm_scale=float(params.ifm.q_params.scale_f32),
                ifm_zero_point=int(params.ifm.q_params.zero_point),
                weight_zero_point=0,
                ofm_scale=float(params.ofm.q_params.scale_f32),
                ofm_zero_point=int(params.ofm.q_params.zero_point),
                kernel_shape=(filter_height, filter_width),
                ofm_channels=out_channels,
                ofm_dtype="int16",
            )

            n = int(filter_height * filter_width)
            eps = 1 / (256 * (n + 1)) if n % 2 == 0 else 0

            scalar_tensor = relay.const(np.ones([1, 1, 1, 1], dtype="int16"), dtype="int16")

            reduced_op = ethosu_ops.ethosu_binary_elementwise(
                ifm=reduced_op,
                ifm2=scalar_tensor,
                lut=lut,
                operator_type="MUL",
                ifm_scale=float(params.ofm.q_params.scale_f32),
                ifm_zero_point=int(params.ofm.q_params.zero_point),
                ifm2_scale=1 / (n - eps),
                ifm2_zero_point=0,
                ofm_scale=float(params.ofm.q_params.scale_f32),
                ofm_zero_point=int(params.ofm.q_params.zero_point),
                ifm_channels=out_channels,
                ifm2_channels=out_channels,
                reversed_operands=False,
                ofm_dtype="int8",
                rounding_mode="NATURAL",
            )
        elif (
            params.ifm.q_params.scale_f32 == params.ofm.q_params.scale_f32
            and params.ifm.q_params.zero_point == params.ofm.q_params.zero_point
        ):
            reduced_op = ethosu_ops.ethosu_pooling(
                ifm=reduced_op,
                lut=lut,
                pooling_type="AVG",
                ifm_scale=float(params.ifm.q_params.scale_f32),
                ifm_zero_point=0,
                ofm_scale=float(params.ofm.q_params.scale_f32),
                ofm_zero_point=0,
                pool_shape=(filter_height, filter_width),
                ofm_channels=out_channels,
                rounding_mode="TRUNCATE",
            )
        else:
            weight_scale = 1 / (filter_height * filter_width)
            weight_values = np.ones([out_channels, filter_height, filter_width, 1])
            bias = -1 * int(params.ifm.q_params.zero_point) * filter_height * filter_width

            scale_bias = vela_api.pack_biases(
                biases=np.ones([ifm_shape[-1]]) * bias,
                ifm_scale=params.ifm.q_params.scale_f32,
                ifm_dtype=np.dtype(params.ifm.dtype),
                weight_scales=np.array([weight_scale], dtype=np.float),
                ofm_scale=params.ofm.q_params.scale_f32,
                is_activation_tanh_or_sigmoid=False,
            )
            reduced_op = ethosu_ops.ethosu_depthwise_conv2d(
                ifm=reduced_op,
                weight=relay.const(weight_values, params.ifm.dtype),
                scale_bias=relay.const(scale_bias, "uint8"),
                lut=lut,
                ifm_scale=float(params.ifm.q_params.scale_f32),
                ifm_zero_point=0,
                weight_zero_point=0,
                ofm_scale=float(params.ofm.q_params.scale_f32),
                ofm_zero_point=int(params.ofm.q_params.zero_point),
                kernel_shape=(filter_height, filter_width),
                ofm_channels=out_channels,
                rounding_mode="NATURAL",
            )

        # Reshape to original ofm shape
        if len(ofm_shape) < 4:
            reduced_op = relay.reshape(reduced_op, ofm_shape)

        return reduced_op


class ConcatRewriter(DFPatternCallback):
    """The newer versions of TFLite converters return a concatenate operator that concatenates
    tensors with same QNN params (if the QNN params of tensors were initially different,
    the converter adds a requantize node), so this rewriter replaces the QNN concatenate with
    "normal" concatenate"""

    def __init__(self):
        super().__init__(require_type=True, rewrite_once=True)
        self.pattern = (
            wildcard().has_attr({"Composite": ethosu_patterns.ConcatParams.composite_name})
        )(None)

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        # Find the tensors that are inputs to the concat and the scales and zero points
        concat_args = list()
        for arg in post.args:
            if isinstance(arg, tvm.relay.expr.Call):
                concat_args.append(arg)

        axis = post.op.body.attrs.axis
        concat = relay.op.concatenate(relay.Tuple(concat_args), axis=axis)
        return concat


class RequantizeRewriter(DFPatternCallback):
    """Convert ethos-u.requantize composite function to an identity operation."""

    def __init__(self):
        super().__init__(require_type=True)
        self.pattern = (
            wildcard().has_attr({"Composite": ethosu_patterns.RequantizeParams.composite_name})
        )(wildcard())

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = ethosu_patterns.RequantizeParams(post.op.body)
        params.ifm.tensor = post.args[0]

        lut = relay.const([], "int8")

        return ethosu_ops.ethosu_identity(
            ifm=params.ifm.tensor,
            lut=lut,
            ifm_scale=float(params.ifm.q_params.scale_f32),
            ifm_zero_point=int(params.ifm.q_params.zero_point),
            ofm_scale=float(params.ofm.q_params.scale_f32),
            ofm_zero_point=int(params.ofm.q_params.zero_point),
        )


class Resize2dRewriter(DFPatternCallback):
    """
    Convert ethos-u.resize2d composite function to an equivalent operation that
    performs the relevant upsampling operation.

    Case 1: No upsampling (upscale factor of 1):
        Identity.
    Case 1: Nearest neighbor upsampling:
        1x1 pooling with 2x2 nearest neighbor upsampling.
    Case 2: Bilinear upsampling:
        2x2 average pool with 2x2 nearest neighbor upsampling.
    """

    def __init__(self):
        super().__init__(require_type=True)
        self.pattern = (
            wildcard().has_attr({"Composite": ethosu_patterns.Resize2dParams.composite_name})
        )(wildcard())

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = ethosu_patterns.Resize2dParams(post.op.body)
        params.ifm.tensor = post.args[0]

        lut = relay.const([], "int8")
        ifm_shape = params.ifm.shape
        in_channels = ifm_shape[-1]
        reduced_op = params.ifm.tensor
        current_size = np.array(ifm_shape[1:3])
        output_size = np.array(params.size)

        if (current_size == output_size).all():
            return ethosu_ops.ethosu_identity(
                reduced_op,
                lut,
                ifm_scale=float(params.ifm.q_params.scale_f32),
                ifm_zero_point=int(params.ifm.q_params.zero_point),
                ofm_scale=float(params.ofm.q_params.scale_f32),
                ofm_zero_point=int(params.ofm.q_params.zero_point),
            )

        padding = [0, 0, 0, 0]
        rounding_mode = "TFL"
        pool_shape = [1, 1]
        if params.method == "linear":
            pool_shape = [2, 2]
            rounding_mode = "NATURAL"
            if params.coordinate_transformation_mode == "asymmetric":
                # Use SAME padding.
                ypad = Resize2dRewriter.get_required_padding(ifm_shape[1])
                xpad = Resize2dRewriter.get_required_padding(ifm_shape[2])
                padding = [ypad // 2, xpad // 2, (ypad + 1) // 2, (xpad + 1) // 2]

        return ethosu_ops.ethosu_pooling(
            ifm=reduced_op,
            lut=lut,
            pooling_type="AVG",
            ifm_scale=float(params.ifm.q_params.scale_f32),
            ifm_zero_point=int(params.ifm.q_params.zero_point),
            ofm_scale=float(params.ofm.q_params.scale_f32),
            ofm_zero_point=int(params.ofm.q_params.zero_point),
            pool_shape=pool_shape,
            ofm_channels=in_channels,
            strides=[1, 1],
            padding=padding,
            upscale="NEAREST",
            rounding_mode=rounding_mode,
        )

    @staticmethod
    def get_required_padding(input_size: int, pool_size: int = 2) -> int:
        """Gets the amount of padding required needed to achieve
        'SAME' padding for a given axis."""
        needed_input = (input_size - 1) + pool_size
        total_padding = max(0, needed_input - input_size)
        return total_padding


class ExpandDimsRewriter(DFPatternCallback):
    """Legalize expand dims to a reshape operator."""

    def __init__(self):
        super().__init__(require_type=True, rewrite_once=True)
        self.pattern = (
            wildcard().has_attr({"Composite": ethosu_patterns.ExpandDimsParams.composite_name})
        )(None)

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = ethosu_patterns.ExpandDimsParams(post.op.body)
        return relay.op.reshape(post.args[0], newshape=params.output.shape)


class SqueezeRewriter(DFPatternCallback):
    """Legalize squeeze to a reshape operator."""

    def __init__(self):
        super().__init__(require_type=True, rewrite_once=True)
        self.pattern = (
            wildcard().has_attr({"Composite": ethosu_patterns.SqueezeParams.composite_name})
        )(None)

    def callback(
        self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map
    ) -> tvm.relay.Expr:
        params = ethosu_patterns.SqueezeParams(post.op.body)
        return relay.op.reshape(post.args[0], newshape=params.output.shape)


class FullyConnectedRewriter(DFPatternCallback):
    """Legalize Fully Connected (with bias and clip) to an NPU operator"""

    def __init__(self):
        super().__init__(require_type=True)
        self.pattern = (
            wildcard().has_attr({"Composite": ethosu_patterns.FullyConnectedParams.composite_name})
        )(wildcard())

    def callback(self, pre, post, node_map):
        params = ethosu_patterns.FullyConnectedParams(post.op.body)
        params.ifm.tensor = post.args[0]

        # IFM reshapes
        ifm = post.args[0]
        if len(params.ifm.shape) != 4 or not params.ifm.shape[1] == params.ifm.shape[2] == 1:
            ifm = relay.reshape(ifm, (1, 1, 1, params.ifm.shape[-1]))

        # Weight transformations
        weights_values = params.weights.values
        weights_values_ohwi = np.expand_dims(weights_values, axis=(1, 2))
        if params.activation:
            activation = "CLIP"
            clip_min = int(params.activation.attrs.a_min)
            clip_max = int(params.activation.attrs.a_max)
        else:
            activation = "NONE"
            clip_min = 0
            clip_max = 0
        bias_values = (
            params.biases.tensor.data.asnumpy()
            if params.biases
            else np.zeros((params.ofm.shape[-1]))
        )
        scale_bias = vela_api.pack_biases(
            biases=bias_values,
            ifm_scale=params.ifm.q_params.scale_f32,
            ifm_dtype=np.dtype(params.ifm.dtype),
            weight_scales=params.weights.q_params.scale_f32,
            ofm_scale=params.ofm.q_params.scale_f32,
            is_activation_tanh_or_sigmoid=False,
        )
        ethosu_fc = ethosu_ops.ethosu_conv2d(
            ifm=ifm,
            weight=relay.const(weights_values_ohwi, params.weights.values.dtype),
            scale_bias=relay.const(scale_bias, "uint8"),
            lut=relay.const([], dtype="int8"),
            ifm_scale=float(params.ifm.q_params.scale_f32),
            ifm_zero_point=int(params.ifm.q_params.zero_point),
            weight_zero_point=int(params.weights.q_params.zero_point),
            ofm_scale=float(params.ofm.q_params.scale_f32),
            ofm_zero_point=int(params.ofm.q_params.zero_point),
            kernel_shape=[1, 1],
            ofm_channels=params.weights.shape[0],
            strides=(1, 1),
            padding=(0, 0, 0, 0),
            dilation=(1, 1),
            activation=activation,
            clip_min=clip_min,
            clip_max=clip_max,
            upscale="NONE",
            ifm_layout="NHWC",
            ofm_layout="NHWC",
        )

        if len(params.ofm.shape) != 4 or not params.ofm.shape[1] == params.ofm.shape[2] == 1:
            ethosu_fc = relay.reshape(ethosu_fc, params.ofm.shape)
        return ethosu_fc


@util.create_npu_function_pass(opt_level=1)
class LegalizeEthosU:
    """This is the pass to call graph-rewrites to perform graph transformation
    in a way such that the operations are replaced with hardware/codegen supported
    operations.
    """

    def transform_npu_function(self, _, func: relay.Function) -> relay.Function:
        """This is the method that replaces the operations with hardware/codegen supported
        operations.
        """
        rewriters = [
            PartitionedSplitRewriter(),
            SplitRewriter(),
            Conv2DRewriter(),
            Conv2DTransposeRewriter(),
            DepthwiseConv2DRewriter(),
            FullyConnectedRewriter(),
            MaxPoolingRewriter(),
            AvgPoolingRewriter(),
            AddRewriter(),
            SubRewriter(),
            MulRewriter(),
            MinRewriter(),
            MaxRewriter(),
            ShlRewriter(),
            AbsRewriter(),
            TanhRewriter(),
            HardSwishRewriter(),
            LeakyReLURewriter(),
            MeanRewriter(),
            ConcatRewriter(),
            SigmoidRewriter(),
            RequantizeRewriter(),
            Resize2dRewriter(),
            ExpandDimsRewriter(),
            SqueezeRewriter(),
            ReshapeRewriter(),
            StridedSliceRewriter(),
            NoOpRewriter(),
        ]
        for rewriter in rewriters:
            func = rewrite(rewriter, func)

        return func

    def __call__(self, *args, **kwargs):
        # pylint is unable figure out the decorated
        # class is callable, thus adding this to
        # suppress the warning.
        pass