from typing import Optional, Union

import numpy as np
import tensorrt as trt
import torch
from tensorrt import ITensor as TRTTensor
from torch.fx.node import Target
from torch_tensorrt.dynamo._SourceIR import SourceIR
from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
from torch_tensorrt.dynamo.conversion.converter_utils import (
    broadcastable,
    cast_trt_tensor,
    get_trt_tensor,
    prepend_ones,
    promote_trt_tensors_to_same_dtype,
    set_layer_name,
)
from torch_tensorrt.dynamo.conversion.impl.elementwise import ne


def where(
    ctx: ConversionContext,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
    input: Union[TRTTensor, np.ndarray, torch.Tensor],
    other: Union[TRTTensor, np.ndarray, torch.Tensor],
    condition: Union[TRTTensor, np.ndarray, torch.Tensor],
) -> TRTTensor:
    if not (broadcastable(input, other)):
        assert "The two torch tensors should be broadcastable"

    x_shape = list(input.shape)
    y_shape = list(other.shape)
    condition_shape = list(condition.shape)
    max_shape_len = max(len(x_shape), len(y_shape), len(condition_shape))

    if not isinstance(condition, TRTTensor):
        condition = get_trt_tensor(ctx, condition, f"{name}_condition")

    if condition.dtype != trt.bool:
        condition = cast_trt_tensor(ctx, condition, trt.float32, f"{name}_cast")
        condition = ne(ctx, target, source_ir, f"{name}_cond_zero", condition, 0)

    diff = max_shape_len - len(condition_shape)
    if diff > 0:
        condition = prepend_ones(ctx, condition, f"{name}_condition_broadcast", diff)

    if not isinstance(input, TRTTensor):
        input = get_trt_tensor(ctx, input, f"{name}_x")
    diff = max_shape_len - len(x_shape)
    if diff > 0:
        input = prepend_ones(ctx, input, f"{name}_input_broadcast", diff)

    if not isinstance(other, TRTTensor):
        other = get_trt_tensor(ctx, other, f"{name}_y")
    diff = max_shape_len - len(y_shape)
    if diff > 0:
        other = prepend_ones(ctx, other, f"{name}_other_broadcast", diff)

    # Ensure that input and other have the same TRT dtype
    input, other = promote_trt_tensors_to_same_dtype(ctx, input, other, name)

    return select(ctx, target, source_ir, name, input, other, condition)


def select(
    ctx: ConversionContext,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
    input: TRTTensor,
    other: TRTTensor,
    condition: TRTTensor,
) -> TRTTensor:
    select_layer = ctx.net.add_select(condition, input, other)
    set_layer_name(select_layer, target, name + "_select", source_ir)
    return select_layer.get_output(0)