# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved
#
# 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
#
#     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.

import numpy as np
import tensorrt as trt

from paddle.tensorrt.converter_utils import (
    add_1D_constant_layer,
    add_constant_layer,
    set_layer_name,
    trt_concat,
    trt_div,
    trt_min,
    trt_pow,
    trt_prod,
    trt_sub,
    trt_sum,
)
from paddle.tensorrt.register import converter_registry

activation_type_map = {
    "pd_op.tanh": trt.ActivationType.TANH,
    "pd_op.relu": trt.ActivationType.RELU,
    "pd_op.sigmoid": trt.ActivationType.SIGMOID,
    "pd_op.silu": trt.ActivationType.SIGMOID,
    "pd_op.swish": trt.ActivationType.SIGMOID,
}


@converter_registry.register("pd_op.relu")
@converter_registry.register("pd_op.tanh")
@converter_registry.register("pd_op.sigmoid")
def activation_converter(network, paddle_op, inputs):
    layer = network.add_activation(
        inputs[0], activation_type_map[paddle_op.name()]
    )
    set_layer_name(layer, paddle_op)
    return layer.get_output(0)


@converter_registry.register("pd_op.logsigmoid")
def logsigmoid_converter(network, paddle_op, inputs):
    sigmoid_layer = network.add_activation(
        inputs[0], trt.ActivationType.SIGMOID
    )
    set_layer_name(sigmoid_layer, paddle_op)
    layer = network.add_unary(
        sigmoid_layer.get_output(0), trt.UnaryOperation.LOG
    )
    set_layer_name(layer, paddle_op)
    return layer.get_output(0)


@converter_registry.register("pd_op.relu6")
def relu6_converter(network, paddle_op, inputs):
    layer = network.add_activation(inputs[0], trt.ActivationType.CLIP)
    layer.alpha = 0.0
    layer.beta = 6.0
    set_layer_name(layer, paddle_op)
    return layer.get_output(0)


@converter_registry.register("pd_op.softmax")
def softmax_converter(network, paddle_op, inputs):
    from paddle.tensorrt.util import support_fp32_mix_precision

    input1 = inputs[0]
    input_shape = input1.shape
    input_dims = len(input_shape)
    axis = paddle_op.attrs().get("axis", -1)

    # support 0 or 1 dims input
    is_0_dims = input_dims == 0
    is_1_dims = input_dims == 1
    if is_0_dims or is_1_dims:
        reshaped_layer = network.add_shuffle(input1)
        reshaped_dims = (1, 1 if is_0_dims else input_shape[0])
        reshaped_layer.reshape_dims = reshaped_dims
        set_layer_name(reshaped_layer, paddle_op)
        input1 = reshaped_layer.get_output(0)
        input_shape = input1.shape
        input_dims = len(input_shape)
        axis = -1

    layer = network.add_softmax(input1)
    set_layer_name(layer, paddle_op)
    support_fp32_mix_precision(paddle_op.name(), layer)
    axes = max(0, input_dims - 3)

    # Handle padded dimensions
    padded_dims = 0
    explicit_batch = 1
    for i in range(input_dims - 1, explicit_batch, -1):
        if input_shape[i] == 1:
            padded_dims += 1
        else:
            break

    if axis < 0:
        axes = input_dims + axis
    else:
        axes = axis

    layer.axes = 1 << axes

    # Support 0 or 1 dims input
    if is_0_dims or is_1_dims:
        reshaped_layer = network.add_shuffle(layer.get_output(0))
        reshaped_layer.reshape_dims = inputs[0].shape
        layer = reshaped_layer
        set_layer_name(layer, paddle_op)

    return layer.get_output(0)


@converter_registry.register("pd_op.gelu")
def gelu_converter(network, paddle_op, inputs):
    input_val = inputs[0]
    approximate = paddle_op.attrs()["approximate"]

    const_shape = [1] * len(input_val.shape)

    if approximate:
        constant_layer_pow = add_constant_layer(
            network,
            [3.0],
            const_shape,
            np.float32,
            name=[paddle_op.name(), "constant_layer_pow"],
        )
        constant_layer_multiply = add_constant_layer(
            network,
            [0.044715],
            const_shape,
            np.float32,
            name=[paddle_op.name(), "constant_layer_multiply"],
        )
        constant_layer_sqrt = add_constant_layer(
            network,
            [0.79788456080286535587989211986876],
            const_shape,
            np.float32,
            name=[paddle_op.name(), "constant_layer_sqrt"],
        )
        constant_layer_one = add_constant_layer(
            network,
            [1.0],
            const_shape,
            np.float32,
            name=[paddle_op.name(), "constant_layer_one"],
        )
        constant_layer_half = add_constant_layer(
            network,
            [0.5],
            const_shape,
            np.float32,
            name=[paddle_op.name(), "constant_layer_half"],
        )

        layer_pow = trt_pow(
            network,
            input_val,
            constant_layer_pow,
            name=[paddle_op.name(), "layer_pow"],
        )
        layer_mul = trt_prod(
            network,
            layer_pow,
            constant_layer_multiply,
            name=[paddle_op.name(), "layer_mul"],
        )
        layer_add = trt_sum(
            network, layer_mul, input_val, name=[paddle_op.name(), "layer_add"]
        )
        layer_sqrt = trt_prod(
            network,
            layer_add,
            constant_layer_sqrt,
            name=[paddle_op.name(), "layer_sqrt"],
        )

        layer_tanh = network.add_activation(layer_sqrt, trt.ActivationType.TANH)
        set_layer_name(layer_tanh, paddle_op)
        layer_one = trt_sum(
            network,
            layer_tanh.get_output(0),
            constant_layer_one,
            name=[paddle_op.name(), "layer_one"],
        )
        layer_cdf = trt_prod(
            network,
            layer_one,
            constant_layer_half,
            name=[paddle_op.name(), "layer_cdf"],
        )
        y = trt_prod(
            network, layer_cdf, input_val, name=[paddle_op.name(), "y"]
        )

        return y
    else:
        constant_layer_one = add_constant_layer(
            network,
            [1.0],
            const_shape,
            np.float32,
            name=[paddle_op.name(), "constant_layer_one"],
        )
        constant_layer_half = add_constant_layer(
            network,
            [0.5],
            const_shape,
            np.float32,
            name=[paddle_op.name(), "constant_layer_half"],
        )
        constant_layer_rsqrt2 = add_constant_layer(
            network,
            [0.70710678118],
            const_shape,
            np.float32,
            name=[paddle_op.name(), "constant_layer_rsqrt2"],
        )

        layer_mul = trt_prod(
            network,
            input_val,
            constant_layer_rsqrt2,
            name=[paddle_op.name(), "layer_mul"],
        )
        layer_erf = network.add_unary(layer_mul, trt.UnaryOperation.ERF)
        set_layer_name(layer_erf, paddle_op)
        layer_add = trt_sum(
            network,
            layer_erf.get_output(0),
            constant_layer_one,
            name=[paddle_op.name(), "layer_add"],
        )
        layer_cdf = trt_prod(
            network,
            layer_add,
            constant_layer_half,
            name=[paddle_op.name(), "layer_cdf"],
        )
        y = trt_prod(
            network, layer_cdf, input_val, name=[paddle_op.name(), "y"]
        )

        return y


@converter_registry.register("pd_op.hardsigmoid")
def hardsigmoid_converter(network, paddle_op, inputs):
    x = inputs[0]
    slope = paddle_op.attrs()["slope"]
    offset = paddle_op.attrs()["offset"]
    hardsigmoid_layer = network.add_activation(
        x, trt.ActivationType.HARD_SIGMOID
    )
    hardsigmoid_layer.alpha = slope
    hardsigmoid_layer.beta = offset
    set_layer_name(hardsigmoid_layer, paddle_op)
    return hardsigmoid_layer.get_output(0)


@converter_registry.register("pd_op.hardswish")
def hardswish_converter(network, paddle_op, inputs):
    x = inputs[0]
    scale = 6.0
    offset = 3.0
    hardsigmoid_layer = network.add_activation(
        x, trt.ActivationType.HARD_SIGMOID
    )
    hardsigmoid_layer.alpha = 1.0 / scale
    hardsigmoid_layer.beta = offset / scale
    set_layer_name(hardsigmoid_layer, paddle_op)
    hardswish_layer = network.add_elementwise(
        x, hardsigmoid_layer.get_output(0), trt.ElementWiseOperation.PROD
    )
    set_layer_name(hardswish_layer, paddle_op)
    return hardswish_layer.get_output(0)


@converter_registry.register("pd_op.elu")
@converter_registry.register("pd_op.elu_")
def elu_converter(network, paddle_op, inputs):
    x = inputs[0]
    alpha = paddle_op.attrs()["alpha"]
    elu_layer = network.add_activation(x, trt.ActivationType.ELU)
    elu_layer.alpha = alpha
    set_layer_name(elu_layer, paddle_op)
    return elu_layer.get_output(0)


@converter_registry.register("pd_op.softplus")
def softplus_converter(network, paddle_op, inputs):
    x = inputs[0]
    beta = paddle_op.attrs()["beta"]
    threshold = paddle_op.attrs()["threshold"]
    layer_clip = network.add_activation(x, trt.ActivationType.CLIP)
    layer_clip.alpha = -3.40282e038
    layer_clip.beta = threshold / beta
    set_layer_name(layer_clip, paddle_op)

    softplus_layer = network.add_activation(
        layer_clip.get_output(0), trt.ActivationType.SOFTPLUS
    )
    softplus_layer.alpha = 1.0 / beta
    softplus_layer.beta = beta
    set_layer_name(softplus_layer, paddle_op)
    return softplus_layer.get_output(0)


@converter_registry.register("pd_op.swish")
@converter_registry.register("pd_op.silu")
def swish_silu_converter(network, paddle_op, inputs):
    layer_output = network.add_activation(
        inputs[0], activation_type_map[paddle_op.name()]
    )
    set_layer_name(layer_output, paddle_op)
    return trt_prod(
        network,
        inputs[0],
        layer_output.get_output(0),
        name=[paddle_op.name(), "trt_prod"],
    )


@converter_registry.register("pd_op.tanh_shrink")
def tanh_shrink_converter(network, paddle_op, inputs):
    x = inputs[0]
    tanh_layer = network.add_activation(x, trt.ActivationType.TANH)
    set_layer_name(tanh_layer, paddle_op)
    subtract_layer = network.add_elementwise(
        x, tanh_layer.get_output(0), trt.ElementWiseOperation.SUB
    )
    set_layer_name(subtract_layer, paddle_op)
    return subtract_layer.get_output(0)


@converter_registry.register("pd_op.stanh")
def stanh_converter(network, paddle_op, inputs):
    x = inputs[0]
    scale_a = paddle_op.attrs()["scale_a"]
    scale_b = paddle_op.attrs()["scale_b"]
    stanh_layer = network.add_activation(x, trt.ActivationType.SCALED_TANH)
    stanh_layer.alpha = scale_b
    stanh_layer.beta = scale_a
    set_layer_name(stanh_layer, paddle_op)
    return stanh_layer.get_output(0)


@converter_registry.register("pd_op.mish")
def mish_converter(network, paddle_op, inputs):
    x = inputs[0]
    softplus_layer = network.add_activation(x, trt.ActivationType.SOFTPLUS)
    set_layer_name(softplus_layer, paddle_op)
    softplus_output = softplus_layer.get_output(0)

    tanh_layer = network.add_activation(
        softplus_output, trt.ActivationType.TANH
    )
    set_layer_name(tanh_layer, paddle_op)
    tanh_output = tanh_layer.get_output(0)

    return trt_prod(
        network, x, tanh_output, name=[paddle_op.name(), "trt_prod"]
    )


@converter_registry.register("pd_op.celu")
def celu_converter(network, paddle_op, inputs):
    input_tensor = inputs[0]
    alpha = paddle_op.attrs()["alpha"]
    input_rank = len(input_tensor.shape)
    constant_shape = trt.Dims([1] * input_rank)
    alpha_data = add_constant_layer(
        network,
        [alpha],
        constant_shape,
        dtype="float32",
        name=[paddle_op.name(), "alpha_data"],
    )
    constant_zero_data = add_constant_layer(
        network,
        [0.0],
        constant_shape,
        dtype="float32",
        name=[paddle_op.name(), "constant_zero_data"],
    )
    constant_one_data = add_constant_layer(
        network,
        [1.0],
        constant_shape,
        dtype="float32",
        name=[paddle_op.name(), "constant_one_data"],
    )
    input_div_with_alpha = trt_div(
        network,
        input_tensor,
        alpha_data,
        name=[paddle_op.name(), "input_div_with_alpha"],
    )
    input_exp_layer = network.add_unary(
        input_div_with_alpha, trt.UnaryOperation.EXP
    )
    set_layer_name(input_exp_layer, paddle_op)
    input_sub_with_one = trt_sub(
        network,
        input_exp_layer.get_output(0),
        constant_one_data,
        name=[paddle_op.name(), "input_sub_with_one"],
    )
    input_prod_with_alpha = trt_prod(
        network,
        input_sub_with_one,
        alpha_data,
        name=[paddle_op.name(), "input_prod_with_alpha"],
    )
    min_input = trt_min(
        network,
        input_prod_with_alpha,
        constant_zero_data,
        name=[paddle_op.name(), "min_input"],
    )
    relu_layer = network.add_activation(input_tensor, trt.ActivationType.RELU)
    set_layer_name(relu_layer, paddle_op)
    output_tensor = trt_sum(
        network,
        relu_layer.get_output(0),
        min_input,
        name=[paddle_op.name(), "output_tensor"],
    )
    return output_tensor


@converter_registry.register("pd_op.thresholded_relu")
def thresholded_relu_converter(network, paddle_op, inputs):
    x = inputs[0]
    threshold = paddle_op.attrs()["threshold"]
    thresholded_relu_layer = network.add_activation(
        x, trt.ActivationType.THRESHOLDED_RELU
    )
    thresholded_relu_layer.alpha = threshold
    set_layer_name(thresholded_relu_layer, paddle_op)
    return thresholded_relu_layer.get_output(0)


@converter_registry.register("pd_op.leaky_relu")
@converter_registry.register("pd_op.leaky_relu_")
def leaky_relu_converter(network, paddle_op, inputs):
    x = inputs[0]
    negative_slope = paddle_op.attrs()["negative_slope"]
    leaky_relu_layer = network.add_activation(x, trt.ActivationType.LEAKY_RELU)
    leaky_relu_layer.alpha = negative_slope
    set_layer_name(leaky_relu_layer, paddle_op)
    return leaky_relu_layer.get_output(0)


@converter_registry.register("pd_op.selu")
def selu_converter(network, paddle_op, inputs):
    x = inputs[0]
    alpha = paddle_op.attrs()["alpha"]
    scale = paddle_op.attrs()["scale"]
    selu_layer = network.add_activation(x, trt.ActivationType.SELU)
    selu_layer.alpha = alpha
    selu_layer.beta = scale
    set_layer_name(selu_layer, paddle_op)
    return selu_layer.get_output(0)


@converter_registry.register("pd_op.prelu")
def prelu_converter(network, paddle_op, inputs):
    input, alpha_data = inputs
    input_dims = input.shape
    data_format = paddle_op.attrs().get("data_format", "NCHW")
    w_dims = trt.Dims(paddle_op.operands()[1].source().shape)
    trt_w_dims = w_dims

    alpha_tensor = network.add_constant(trt_w_dims, alpha_data)
    set_layer_name(alpha_tensor, paddle_op)
    alpha_tensor = alpha_tensor.get_output(0)

    alpha_dims = alpha_tensor.shape
    real_alpha_tensor = alpha_tensor

    if len(alpha_dims) != len(input_dims):
        reshape_layer = network.add_shuffle(alpha_tensor)
        set_layer_name(reshape_layer, paddle_op)
        c = alpha_dims[0]

        n_tensor = add_1D_constant_layer(
            network, [1], name=[paddle_op.name(), "n_tensor"]
        )
        c_tensor = add_1D_constant_layer(
            network, [c], name=[paddle_op.name(), "c_tensor"]
        )
        hw_tensor = None
        if len(input_dims) - 2 > 0:
            hw_tensor = add_1D_constant_layer(
                network,
                [1] * (len(input_dims) - 2),
                name=[paddle_op.name(), "hw_tensor"],
            )

        if data_format == "NCHW":
            if hw_tensor:
                shape_tensor = trt_concat(
                    network,
                    [n_tensor, c_tensor, hw_tensor],
                    name=[paddle_op.name(), "shape_tensor"],
                )
            else:
                shape_tensor = trt_concat(
                    network,
                    [n_tensor, c_tensor],
                    name=[paddle_op.name(), "shape_tensor"],
                )
        else:
            if hw_tensor:
                shape_tensor = trt_concat(
                    network,
                    [n_tensor, hw_tensor, c_tensor],
                    name=[paddle_op.name(), "shape_tensor"],
                )
            else:
                shape_tensor = trt_concat(
                    network,
                    [n_tensor, c_tensor],
                    name=[paddle_op.name(), "shape_tensor"],
                )

        reshape_layer.set_input(1, shape_tensor)
        real_alpha_tensor = reshape_layer.get_output(0)

    layer = network.add_parametric_relu(input, real_alpha_tensor)
    set_layer_name(layer, paddle_op)
    return layer.get_output(0)