# 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

import paddle
from paddle.pir.core import datatype_to_str
from paddle.tensorrt.converter_utils import (
    add_1D_constant_layer,
    get_input_constant_value,
    resize_to_1d,
    set_layer_name,
    trt_cast,
    trt_floor_div,
    trt_max,
    trt_min,
    trt_reduce_to_scalar,
    trt_reshape,
    trt_shape,
    trt_sub,
)
from paddle.tensorrt.register import converter_registry


@converter_registry.register("pd_op.full_int_array")
def full_int_array_converter(network, paddle_op, inputs):
    value = paddle_op.attrs()["value"]
    if len(value) == 0:
        return ()
    value_weight = trt.Weights(np.array(value, dtype=np.int32))
    full_int_array_layer = network.add_constant([len(value)], value_weight)
    set_layer_name(full_int_array_layer, paddle_op)
    return full_int_array_layer.get_output(0)


@converter_registry.register("pd_op.full")
def full_converter(network, paddle_op, inputs):
    shape = paddle_op.attrs()["shape"]
    value = paddle_op.attrs().get("value", 1.0)
    dtype = paddle_op.attrs().get("dtype")
    out_dtype = np.dtype(datatype_to_str[dtype])
    if out_dtype == np.dtype("float64"):
        out_dtype = np.dtype("float32")
    if out_dtype == np.dtype("int64"):
        out_dtype = np.dtype("int32")
    full_layer = network.add_constant(
        shape, np.full(shape, value, dtype=out_dtype)
    )
    set_layer_name(full_layer, paddle_op)
    return full_layer.get_output(0)


@converter_registry.register("pd_op.assign")
@converter_registry.register("pd_op.assign_out_")
def assign_converter(network, paddle_op, inputs):
    input_tensor = inputs[0]
    identity_layer = network.add_identity(input_tensor)
    set_layer_name(identity_layer, paddle_op)
    return identity_layer.get_output(0)


@converter_registry.register("pd_op.assign_value")
@converter_registry.register("pd_op.assign_value_")
def assign_value_converter(network, paddle_op, inputs):
    attrs = paddle_op.attrs()
    shape = attrs['shape']
    dtype = attrs['dtype']
    values = attrs['values']

    paddle_to_np_dtype_map = {
        paddle.float16: np.float16,
        paddle.float32: np.float32,
        paddle.float64: np.float64,
        paddle.int32: np.int32,
        paddle.int64: np.int64,
    }

    if dtype not in paddle_to_np_dtype_map:
        raise ValueError(
            f"Unsupported dtype {dtype} for assign_value op in TRT converter."
        )

    np_dtype = paddle_to_np_dtype_map[dtype]

    arr = np.array(values, dtype=np_dtype).reshape(shape)
    if np_dtype == np.int64:
        arr = arr.astype(np.int32)
    const_layer = network.add_constant(tuple(shape), arr)
    set_layer_name(const_layer, paddle_op)
    if const_layer is None:
        raise RuntimeError("Failed to create constant layer for assign_value.")

    return const_layer.get_output(0)


@converter_registry.register("pd_op.arange")
def arange_converter(network, paddle_op, inputs):
    start, end, step = inputs
    zero_tensor = add_1D_constant_layer(
        network, 0, name=[paddle_op.name(), 'zero_tensor']
    )

    delta = trt_sub(network, start, end, name=[paddle_op.name(), 'delta'])

    f_quotient_tensor = trt_floor_div(
        network, delta, step, name=[paddle_op.name(), 'f_quotient_tensor']
    )

    dtype = paddle_op.attrs().get("dtype")

    if start.dtype == trt.DataType.FLOAT:
        quotient_tensor = trt_cast(
            network,
            f_quotient_tensor,
            trt.int32,
            name=[paddle_op.name(), 'quotient_tensor'],
        )
    else:
        quotient_tensor = f_quotient_tensor

    delta_1 = trt_sub(
        network,
        zero_tensor,
        quotient_tensor,
        name=[paddle_op.name(), 'delta_1'],
    )
    number_tensor = trt_max(
        network, delta_1, zero_tensor, name=[paddle_op.name(), 'number_tensor']
    )
    start1 = inputs[0]
    start1 = trt_reshape(network, start1, (), name=[paddle_op.name(), 'start1'])

    fill_layer = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
    fill_layer.set_input(0, number_tensor)
    fill_layer.set_input(1, start1)
    fill_layer.set_input(2, step)
    set_layer_name(fill_layer, paddle_op)

    output_tensor = fill_layer.get_output(0)

    if dtype == paddle.int64 or dtype == paddle.int32:
        output_tensor = trt_cast(
            network,
            output_tensor,
            trt.int32,
            name=[paddle_op.name(), 'output_tensor'],
        )
    return output_tensor


@converter_registry.register("pd_op.full_like")
def full_like_converter(network, paddle_op, inputs):
    input_tensor = inputs[0]
    shape = input_tensor.shape
    ndims = len(shape)

    dtype = int(paddle_op.attrs().get("dtype", -1))

    dtype_map = {
        0: None,  # Undefined
        1: trt.bool,  # bool
        2: trt.int32,  # int32
        3: trt.int32,  # int64 -> int32
        4: trt.int32,  # int16 -> int32
        5: trt.float32,  # float16 -> float32
        6: trt.float32,  # float64 -> float32
        7: trt.float32,  # float32
        8: trt.int32,  # uint8 -> int32
        11: trt.float32,  # float32
    }

    target_dtype = dtype_map.get(dtype, None)
    if target_dtype is None:
        target_dtype = input_tensor.dtype

    value = get_input_constant_value(paddle_op, inputs, 1)
    if value is not None:
        if isinstance(value, (list, tuple)):
            value = value[0] if value else 0

        if target_dtype == trt.int32:
            value_tensor = add_1D_constant_layer(
                network,
                int(value),
                np.int32,
                name=[paddle_op.name(), 'value_tensor'],
            )
        else:
            value_tensor = add_1D_constant_layer(
                network,
                float(value),
                np.float32,
                name=[paddle_op.name(), 'value_tensor'],
            )
    else:
        value_tensor = inputs[1]
        if value_tensor.dtype != target_dtype:
            value_tensor = trt_cast(
                network,
                value_tensor,
                target_dtype,
                name=[paddle_op.name(), 'value_tensor'],
            )

    shape_tensor = trt_shape(
        network, input_tensor, name=[paddle_op.name(), 'shape_tensor']
    )
    one_rank_tensor = add_1D_constant_layer(
        network, [1] * ndims, name=[paddle_op.name(), 'one_rank_tensor']
    )
    input_shape_tensor = one_rank_tensor

    shuffle_layer = network.add_shuffle(value_tensor)
    shuffle_layer.set_input(1, input_shape_tensor)
    set_layer_name(shuffle_layer, paddle_op)
    start = trt.Dims([0] * ndims)
    size = trt.Dims([1] * ndims)
    stride = trt.Dims([1] * ndims)

    starts_tensor = add_1D_constant_layer(
        network, [0] * ndims, name=[paddle_op.name(), 'starts_tensor']
    )
    one_tensor = add_1D_constant_layer(
        network, 1, name=[paddle_op.name(), 'one_tensor']
    )
    sizes_tensor = trt_max(
        network,
        input_shape_tensor,
        shape_tensor,
        name=[paddle_op.name(), 'sizes_tensor'],
    )
    input_sub_tensor = trt_sub(
        network,
        input_shape_tensor,
        one_tensor,
        name=[paddle_op.name(), 'input_sub_tensor'],
    )
    strides_tensor = trt_min(
        network,
        one_tensor,
        input_sub_tensor,
        name=[paddle_op.name(), 'strides_tensor'],
    )

    layer = network.add_slice(shuffle_layer.get_output(0), start, size, stride)
    layer.set_input(1, starts_tensor)
    layer.set_input(2, sizes_tensor)
    layer.set_input(3, strides_tensor)
    set_layer_name(layer, paddle_op)

    output = layer.get_output(0)

    if output.dtype != target_dtype:
        output = trt_cast(
            network, output, target_dtype, name=[paddle_op.name(), 'output']
        )

    return output


@converter_registry.register("pd_op.full_with_tensor")
def full_with_tensor_converter(network, paddle_op, inputs):
    value_input = inputs[0]

    shape_tensor = None
    dtype = paddle_op.attrs()["dtype"]

    operands = paddle_op.operands()
    num_operands = len(operands)

    if num_operands >= 2:
        shape_tensor = inputs[1]
        if isinstance(shape_tensor, list):
            shape_tensor_list = shape_tensor
        else:
            shape_tensor_list = [shape_tensor]

    shape_val = get_input_constant_value(paddle_op, inputs, 1)
    if shape_val is not None:
        shape_tensor = shape_val
    else:
        shape_tensor = inputs[1]

    tensor_rank = 0
    if isinstance(shape_tensor, trt.ITensor):
        shapes_tensor = shape_tensor
    elif isinstance(shape_tensor, (list, tuple)):
        shapes_tensor = shape_tensor
    else:
        raise TypeError(f"Unsupported shape_tensor type: {type(shape_tensor)}")

    if shape_tensor is not None and len(shape_tensor_list) == 1:
        is_dynamic_shape = True
    elif len(shape_tensor_list) >= 1:
        is_dynamic_shape = True
    else:
        is_dynamic_shape = False

    if is_dynamic_shape:
        if len(shape_tensor_list) == 1:
            shape_tensor = shape_tensor_list[0]
            if not isinstance(shape_tensor, trt.ITensor):
                raise TypeError("shape_tensor must be an ITensor")
            tensor_rank = shape_tensor.shape[0]
            shapes_tensor = shape_tensor
        else:
            shape_tensors = []
            for tensor in shape_tensor_list:
                if len(tensor.shape) == 0:
                    tensor = trt_reshape(
                        network, tensor, (1,), name=[paddle_op.name(), "tensor"]
                    )
                shape_tensors.append(tensor)

            concat_layer = network.add_concatenation(shape_tensors)
            set_layer_name(concat_layer, paddle_op)
            shapes_tensor = concat_layer.get_output(0)
            tensor_rank = len(shape_tensors)

        shapes_tensor = resize_to_1d(
            network, shapes_tensor, name=[paddle_op.name(), "shapes_tensor"]
        )
        fill_layer = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
        fill_layer.set_input(0, shapes_tensor)

    if dtype == paddle.int32 or dtype == paddle.int64:
        beta_vec = [0] * tensor_rank
        value_input = trt_reduce_to_scalar(
            network, value_input, name=[paddle_op.name(), 'value_input']
        )
        fill_layer.set_input(1, value_input)
        fill_layer.set_input(
            2, add_1D_constant_layer(network, beta_vec, np.int32)
        )
    elif dtype == paddle.float32:
        beta_vec = [0.0] * tensor_rank
        value_input = trt_reduce_to_scalar(
            network,
            value_input,
            trt.float32,
            name=[paddle_op.name(), 'value_input'],
        )
        fill_layer.set_input(1, value_input)
        fill_layer.set_input(
            2, add_1D_constant_layer(network, beta_vec, np.float32)
        )
    else:
        raise ValueError(f"Unsupported dtype for full_with_tensor: {dtype}")

    set_layer_name(fill_layer, paddle_op)
    output_tensor = fill_layer.get_output(0)
    return output_tensor


@converter_registry.register("pd_op.meshgrid")
def meshgrid_converter(network, paddle_op, vec_inputs):
    inputs = vec_inputs[0]
    n = len(inputs)
    outputs = []

    # get all input dims (all input is 1-dim)
    input_dims = [network.add_shape(inp).get_output(0) for inp in inputs]

    for k in range(n):
        # --------------------------------
        # step1:reshape k input as [1,..,Dk,..,1]
        # --------------------------------
        x = inputs[k]
        reshape_dims = []  # init dims as 1
        for i in range(n):
            one = add_1D_constant_layer(
                network,
                1,
                dtype=np.int32,
                is_scalar=False,
                name=[paddle_op.name(), f'one_{k}'],
            )
            reshape_dims.append(one)
        # replace k-th input dim as Dk
        reshape_dims[k] = input_dims[k]

        dim_concat = network.add_concatenation(reshape_dims)
        set_layer_name(dim_concat, paddle_op)
        x_reshaped = network.add_shuffle(x)
        x_reshaped.set_input(1, dim_concat.get_output(0))

        # --------------------------------
        # step2: create tensor([D1, D2, ..., 1, ..., Dn]) that filled with 1
        # --------------------------------
        ones_shape = []
        for i in range(n):
            ones_shape.append(input_dims[i])
        ones_shape[k] = add_1D_constant_layer(
            network,
            1,
            dtype=np.int32,
            is_scalar=False,
            name=[paddle_op.name(), f'ones_shape_{k}'],
        )
        dim_concat = network.add_concatenation(ones_shape)
        set_layer_name(dim_concat, paddle_op)

        # Fill constant 1
        fill_layer = network.add_fill(shape=(), op=trt.FillOperation.LINSPACE)
        fill_layer.set_input(0, dim_concat.get_output(0))
        value_input = add_1D_constant_layer(
            network,
            1,
            dtype=np.float32,
            is_scalar=True,
            name=[paddle_op.name(), 'one_for_fill'],
        )
        fill_layer.set_input(1, value_input)
        beta_vec = [0] * n
        fill_layer.set_input(
            2, add_1D_constant_layer(network, beta_vec, np.float32)
        )

        # --------------------------------
        # step3: element wise multiplication
        # --------------------------------
        grid = network.add_elementwise(
            x_reshaped.get_output(0),
            fill_layer.get_output(0),
            trt.ElementWiseOperation.PROD,
        ).get_output(0)
        outputs.append(grid)
    return outputs