import logging
from typing import List, Optional, Sequence, Tuple, Union

import numpy as np
import tensorrt as trt
import torch
from torch._subclasses.fake_tensor import unset_fake_temporarily
from torch.fx.node import Target
from torch_tensorrt import ENABLED_FEATURES
from torch_tensorrt.dynamo._SourceIR import SourceIR
from torch_tensorrt.dynamo.conversion import impl
from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
from torch_tensorrt.dynamo.conversion.converter_utils import (
    cast_trt_tensor,
    create_constant,
    get_axes_for_reduce_op,
    get_positive_dim,
    get_trt_tensor,
    has_dynamic_shape,
    set_layer_name,
    to_trt_weights,
)
from torch_tensorrt.dynamo.conversion.impl.cat import cat
from torch_tensorrt.dynamo.conversion.impl.elementwise.ops import ge
from torch_tensorrt.dynamo.conversion.impl.shape import shape as get_shape
from torch_tensorrt.dynamo.utils import DYNAMIC_DIM

_LOGGER: logging.Logger = logging.getLogger(__name__)


def batch_norm(
    ctx: ConversionContext,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
    input: trt.ITensor,
    momentum: float,
    eps: float,
    return_mean_rstd: bool,
    weight: Optional[Union[trt.ITensor, torch.Tensor, np.ndarray]] = None,
    bias: Optional[Union[trt.ITensor, torch.Tensor, np.ndarray]] = None,
    running_mean: Optional[Union[trt.ITensor, torch.Tensor, np.ndarray]] = None,
    running_var: Optional[Union[trt.ITensor, torch.Tensor, np.ndarray]] = None,
    training: bool = False,
    cudnn_enabled: bool = False,
) -> Union[trt.ITensor, Tuple[trt.ITensor, torch.Tensor, torch.Tensor]]:
    if has_dynamic_shape(input.shape):
        assert input.shape[1] != -1, "Channel dim can't be dynamic for batch norm."

    # Save the original output shape for later use
    output_shape = input.shape
    feature_num = output_shape[1]
    # We perform constant folding for batch norm when the weight, bias, running_mean, and running_var are all tensors.
    # Batch norm operation can be fused into a single layer, which is more efficient than the original implementation.
    # In this way, the batch norm layer will be fused with the Convolution layer and get a performance boost.
    # TODO: lanl: to remove this once we have solved the batchnorm constant folding issue in RTX
    # https://github.com/pytorch/TensorRT/issues/3699
    if ENABLED_FEATURES.tensorrt_rtx or any(
        [
            isinstance(weight, trt.ITensor),
            isinstance(bias, trt.ITensor),
            isinstance(running_mean, trt.ITensor),
            isinstance(running_var, trt.ITensor),
        ]
    ):
        # We name the weight here according to the state_dict name
        with unset_fake_temporarily():
            weight = (
                get_trt_tensor(
                    ctx, torch.ones((feature_num,)), f"{name}_weight", dtype=input.dtype
                )
                if weight is None
                else get_trt_tensor(ctx, weight, f"{name}_weight")
            )
            bias = (
                get_trt_tensor(
                    ctx, torch.zeros((feature_num,)), f"{name}_bias", dtype=input.dtype
                )
                if bias is None
                else get_trt_tensor(ctx, bias, f"{name}_bias")
            )
            running_mean = (
                get_trt_tensor(
                    ctx,
                    torch.zeros((feature_num,)),
                    f"{name}_running_mean",
                    dtype=input.dtype,
                )
                if running_mean is None
                else get_trt_tensor(ctx, running_mean, f"{name}_running_mean")
            )
            running_var = (
                get_trt_tensor(
                    ctx,
                    torch.ones((feature_num,)),
                    f"{name}_running_var",
                    dtype=input.dtype,
                )
                if running_var is None
                else get_trt_tensor(ctx, running_var, f"{name}_running_var")
            )

        # eps_tensor for numerical stability
        eps_tensor = get_trt_tensor(ctx, eps, f"{name}_eps", dtype=input.dtype)

        # adjusted_var = running_var + eps
        adjusted_var = impl.elementwise.add(
            ctx, target, source_ir, f"{name}_adjusted_var", running_var, eps_tensor
        )

        # sqrt_adjusted_var = sqrt(adjusted_var)
        sqrt_adjusted_var = impl.unary.sqrt(
            ctx, target, source_ir, f"{name}_sqrt", adjusted_var
        )

        # scale = weight / sqrt_adjusted_var
        scale = impl.elementwise.div(
            ctx, target, source_ir, f"{name}_scale", weight, sqrt_adjusted_var
        )

        # scaled_running_mean = running_mean * scale
        scaled_running_mean = impl.elementwise.mul(
            ctx, target, source_ir, f"{name}_scaled_running_mean", running_mean, scale
        )

        # bias_adjusted = bias - scaled_running_mean
        bias_adjusted = impl.elementwise.sub(
            ctx, target, source_ir, f"{name}_bias_adjusted", bias, scaled_running_mean
        )

        # Reshape scale and bias_adjusted to match input shape for broadcasting
        expanded_shape = [1] * len(output_shape)
        expanded_shape[1] = feature_num  # Set channel dimension
        scale_reshape = impl.shuffle.reshape(
            ctx,
            target,
            source_ir,
            f"{name}_reshape_scale",
            scale,
            tuple(expanded_shape),
        )
        bias_adjusted_reshape = impl.shuffle.reshape(
            ctx,
            target,
            source_ir,
            f"{name}_reshape_bias",
            bias_adjusted,
            tuple(expanded_shape),
        )

        # Apply the scale and bias to the input
        scaled_input = impl.elementwise.mul(
            ctx, target, source_ir, f"{name}_scaled_input", input, scale_reshape
        )
        output = impl.elementwise.add(
            ctx,
            target,
            source_ir,
            f"{name}_output",
            scaled_input,
            bias_adjusted_reshape,
        )

    else:
        with unset_fake_temporarily():
            if weight is None:
                weight = torch.ones((feature_num,))

            if bias is None:
                bias = torch.zeros((feature_num,))

            if running_mean is None:
                running_mean = torch.zeros((feature_num,))

            if running_var is None:
                running_var = torch.ones((feature_num,))

            power = torch.ones_like(weight)

            adjusted_scale, adjusted_bias = batch_norm_constant_folding(
                weight, bias, running_mean, running_var, eps
            )

        adjusted_scale = to_trt_weights(
            ctx,
            adjusted_scale,
            name,
            layer_type_name="SCALE",
            weight_type_name="SCALE",
            target=target,
            source_ir=source_ir,
        )
        adjusted_bias = to_trt_weights(
            ctx,
            adjusted_bias,
            name,
            layer_type_name="SCALE",
            weight_type_name="SHIFT",
            target=target,
            source_ir=source_ir,
        )

        power = to_trt_weights(
            ctx,
            power,
            name,
            layer_type_name="SCALE",
            weight_type_name="POWER",
            target=target,
            source_ir=source_ir,
        )

        if len(input.shape) < 4:
            new_shape = (
                (input.shape[0], input.shape[1], 1, 1)
                if len(input.shape) == 2
                else (input.shape[0], input.shape[1], input.shape[2], 1)
            )
            input = impl.shuffle.reshape(
                ctx, target, source_ir, f"{name}_reshape_2d", input, new_shape
            )

        layer = ctx.net.add_scale_nd(
            input, trt.ScaleMode.CHANNEL, adjusted_bias, adjusted_scale, power, 1
        )
        set_layer_name(layer, target, name, source_ir)
        output = layer.get_output(0)

    # For BatchNorm1d, reshape output back to original shape if necessary
    if len(output_shape) < 4:
        output = impl.shuffle.reshape(
            ctx,
            target,
            source_ir,
            f"{name}_reshape_1d",
            output,
            output_shape,
        )

    if return_mean_rstd:
        # return fake mean and rstd for now
        return output, None, None

    return output


def batch_norm_constant_folding(
    weight: torch.Tensor,
    bias: torch.Tensor,
    running_mean: torch.Tensor,
    running_var: torch.Tensor,
    eps: float,
) -> Tuple[torch.Tensor, torch.Tensor]:
    adjusted_scale = weight / torch.sqrt(running_var + eps)
    adjusted_bias = bias - running_mean * adjusted_scale
    return adjusted_scale, adjusted_bias


def native_layer_norm(
    ctx: ConversionContext,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
    input: trt.ITensor,
    normalized_shape: List[int],
    weight: Optional[Union[trt.ITensor, torch.Tensor, np.ndarray]],
    bias: Optional[Union[trt.ITensor, torch.Tensor, np.ndarray]],
    eps: float,
) -> Tuple[trt.ITensor, torch.Tensor, torch.Tensor]:
    dims = list(range(len(input.shape) - len(normalized_shape), len(input.shape)))
    axes = get_axes_for_reduce_op(dims)

    weight = get_trt_tensor(
        ctx, weight if weight is not None else 1.0, f"{name}_weight"
    )
    bias = get_trt_tensor(ctx, bias if bias is not None else 0.0, f"{name}_bias")

    # Cast weight and bias to have same dtype as input
    weight = cast_trt_tensor(
        ctx, weight, input.dtype, f"{name}_weight_cast", target, source_ir
    )
    bias = cast_trt_tensor(
        ctx, bias, input.dtype, f"{name}_bias_cast", target, source_ir
    )

    if tuple(input.shape) != tuple(weight.shape):
        weight = impl.slice.expand(
            ctx, target, source_ir, f"{name}_expand_weight", weight, input.shape
        )

    if tuple(input.shape) != tuple(bias.shape):
        bias = impl.slice.expand(
            ctx, target, source_ir, f"{name}_expand_bias", bias, input.shape
        )

    layer = ctx.net.add_normalization(input, weight, bias, axes)
    layer.epsilon = eps
    set_layer_name(layer, target, name, source_ir)

    # return fake mean and rstd for now
    return layer.get_output(0), None, None


def native_group_norm(
    ctx: ConversionContext,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
    input: trt.ITensor,
    weight: Optional[Union[trt.ITensor, torch.Tensor, np.ndarray]],
    bias: Optional[Union[trt.ITensor, torch.Tensor, np.ndarray]],
    N: int,
    C: int,
    HxW: int,
    group: int,
    eps: float,
) -> Tuple[trt.ITensor, torch.Tensor, torch.Tensor]:
    rank = len(input.shape)

    assert rank >= 3, f"Expected at least 3 dimensions for input tensor but got {rank}"

    assert (
        C == input.shape[1]
    ), f"num_channels ({C}) must be equal to number of channels in input ({input.shape[1]})"

    shape = [1, group] + [1] * (rank - 2)

    with unset_fake_temporarily():
        weight_torch = torch.ones(shape)
        bias_torch = torch.zeros(shape)

    weight_one = get_trt_tensor(ctx, weight_torch, f"{name}_weight_one", input.dtype)
    bias_zero = get_trt_tensor(ctx, bias_torch, f"{name}_bias_zero", input.dtype)

    axes = get_axes_for_reduce_op(list(range(1 if group == 1 else 2, rank)))

    # INormalizationLayer scales the normalized output per-group, but PyTorch scales the normalized output per-channel,
    # hence causing diverse result. Let TensorRT does no-op for scaling here, and do scaling ourselves later
    layer = ctx.net.add_normalization(input, weight_one, bias_zero, axes)
    layer.epsilon = eps
    layer.num_groups = group
    set_layer_name(layer, target, name, source_ir)
    output = layer.get_output(0)

    shape[1] = C

    if weight is not None:
        weight = get_trt_tensor(ctx, weight, f"{name}_weight")
        weight = cast_trt_tensor(
            ctx, weight, input.dtype, f"{name}_cast_weight", target, source_ir
        )
        weight = impl.shuffle.reshape(
            ctx, target, source_ir, f"{name}_reshape_weight", weight, shape
        )
        output = impl.elementwise.mul(
            ctx, target, source_ir, f"{name}_mul_weight", output, weight
        )

    if bias is not None:
        bias = get_trt_tensor(ctx, bias, f"{name}_bias")
        bias = cast_trt_tensor(
            ctx, bias, input.dtype, f"{name}_cast_bias", target, source_ir
        )
        bias = impl.shuffle.reshape(
            ctx, target, source_ir, f"{name}_reshape_bias", bias, shape
        )
        output = impl.elementwise.add(
            ctx, target, source_ir, f"{name}_add_bias", output, bias
        )

    # return fake mean and rstd for now
    return output, None, None


def softmax(
    ctx: ConversionContext,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
    input: trt.ITensor,
    dim: int,
    half_to_float: bool,
) -> Union[trt.ITensor, Sequence[trt.ITensor]]:
    dim = get_positive_dim(dim, len(input.shape))

    if half_to_float:
        input = cast_trt_tensor(ctx, input, torch.float, name, target, source_ir)

    layer = ctx.net.add_softmax(input)
    layer.axes = 1 << dim
    set_layer_name(layer, target, name, source_ir)
    return layer.get_output(0)


def pdist(
    ctx: ConversionContext,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
    input: trt.ITensor,
    p: float = 2,
) -> Union[trt.ITensor, Sequence[trt.ITensor]]:
    shape = input.shape
    # Extend input from shape [N, D] to [N, 1, D]
    extend_input = impl.unsqueeze.unsqueeze(
        ctx,
        target,
        source_ir,
        f"{name}_unsqueeze",
        input,
        1,
    )

    # Expand the input from [N, 1, D] to [N, N, D]
    x = impl.slice.expand(
        ctx,
        target,
        source_ir,
        f"{name}_expand",
        extend_input,
        (shape[0], shape[0]) + shape[1:],
    )
    # Subtract the expanded input from original input. Result shape = [N, N, D]
    # This matrix has the distance of each sample to every other sample and hence the shape is [N, N, D]
    x = impl.elementwise.sub(ctx, target, source_ir, f"{name}_sub", x, input)

    if p == 0:
        # norm = torch.sum(x!=0, dim=2)
        nonzero_val = impl.elementwise.ne(ctx, target, source_ir, f"{name}_ne", x, 0)
        norm = impl.reduce.sum(
            ctx, target, source_ir, f"{name}_sum", nonzero_val, dim=2, keepdim=False
        )
        norm = cast_trt_tensor(
            ctx, norm, torch.float32, f"{name}_cast", target, source_ir
        )
    elif p == 1:
        # norm = torch.sum(torch.abs(x), dim=2)
        abs_val = impl.unary.abs(ctx, target, source_ir, f"{name}_abs", x)
        norm = impl.reduce.sum(
            ctx, target, source_ir, f"{name}_sum", abs_val, dim=2, keepdim=False
        )
    elif 0 < p < 1 or 1 < p < float("inf"):
        # norm = torch.pow(torch.sum(torch.pow(torch.abs(x), p), dim=2), 1/p)
        abs_val = impl.unary.abs(ctx, target, source_ir, f"{name}_abs", x)
        pow_val = impl.elementwise.pow(
            ctx, target, source_ir, f"{name}_pow1", abs_val, p
        )
        sum_val = impl.reduce.sum(
            ctx, target, source_ir, f"{name}_sum", pow_val, dim=2, keepdim=False
        )
        norm = impl.elementwise.pow(
            ctx, target, source_ir, f"{name}_pow2", sum_val, 1 / p
        )
    elif p == float("inf"):
        # norm = torch.max(torch.abs(x))
        abs_val = impl.unary.abs(ctx, target, source_ir, f"{name}_abs", x)
        norm = impl.reduce.max(
            ctx,
            target,
            source_ir,
            f"{name}_max",
            abs_val,
            dim=2,
            keepdim=False,
            return_indices=False,
        )
    else:
        raise RuntimeError(
            f"p should between [0, inf], currently p={p} is not supported!"
        )
    if shape[0] == DYNAMIC_DIM:
        dim = get_shape(ctx, target, source_ir, f"{name}_get_shape", input, 0)
        shuffle_layer = ctx.net.add_shuffle(dim)
        shuffle_layer.reshape_dims = trt.Dims()
        set_layer_name(shuffle_layer, target, f"{name}_shuffle", source_ir)
        dim_tensor = shuffle_layer.get_output(0)
        indices_tensor = tri_upper_indices(
            ctx, target, source_ir, f"{name}_triu_indices", dim_tensor
        )
        gather_layer = ctx.net.add_gather_v2(
            norm, indices_tensor, mode=trt.GatherMode.ND
        )
        set_layer_name(gather_layer, target, f"{name}_gather_layer", source_ir)
        gather_layer.axis = 2
        return gather_layer.get_output(0)
    else:
        indices = np.triu_indices(shape[0], k=1)
        return impl.select.index(ctx, target, source_ir, f"{name}_index", norm, indices)


def tri_upper_indices(
    ctx: ConversionContext,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
    size_tensor: trt.ITensor,
) -> trt.ITensor:
    """
    Return the indices for the upper-triangle part of a square size of matrix in a N-by-2 Tensor,
    where the diagonal offset = 1. One loop is used to calculate the indices like below.
    x = 0, y = 0, y_start = 1
    out_size = size * (size - 1) // 2
    for _ in range(out_size):
        y_out.append(y_start + y)
        x_out.append(x)
        y += 1
        if (y_start + y) >= size:
            x += 1
            y_start += 1
            y = 0
    Args:
        ctx (ConversionContext): A ConversionContext containing the TensorRT network.
        target (Target): Target of calling node.
        source_ir (Optional[SourceIR]): SourceIR of calling converter.
        name (str): Name of the calling layer.
        size_tensor (trt.ITensor): number of rows in the 2-D square matrix. scalar tensor.

    Example:
        if size_tensor is 4, it will return [[0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3]]
    """
    constant_0 = create_constant(ctx, 0, f"{name}_zero", np.int32, 0)
    constant_1 = create_constant(ctx, 1, f"{name}_one", np.int32, 0)
    constant_2 = create_constant(ctx, 2, f"{name}_two", np.int32, 0)

    size_minus_one = impl.elementwise.sub(
        ctx, target, source_ir, f"{name}_size_minus_one", size_tensor, constant_1
    )

    size_mult_prev = impl.elementwise.mul(
        ctx, target, source_ir, f"{name}_size_mult_prev", size_tensor, size_minus_one
    )

    num_loop = impl.elementwise.floor_divide(
        ctx, target, source_ir, f"{name}_num_loop", size_mult_prev, constant_2
    )

    loop = ctx.net.add_loop()
    loop.add_trip_limit(num_loop, trt.TripLimit.COUNT)

    x_recurrence = loop.add_recurrence(constant_0)
    set_layer_name(x_recurrence, target, f"{name}_x_recurrence", source_ir)
    x_tensor = x_recurrence.get_output(0)

    y_recurrence = loop.add_recurrence(constant_0)
    set_layer_name(y_recurrence, target, f"{name}_y_recurrence", source_ir)
    y_tensor = y_recurrence.get_output(0)

    y_start_recurrence = loop.add_recurrence(constant_1)
    set_layer_name(y_start_recurrence, target, f"{name}_y_start_recurrence", source_ir)
    y_start_tensor = y_start_recurrence.get_output(0)

    x_inc = impl.elementwise.add(
        ctx,
        target,
        source_ir,
        f"{name}_x_inc",
        x_tensor,
        constant_1,
    )

    y_current = impl.elementwise.add(
        ctx,
        target,
        source_ir,
        f"{name}_y_current",
        y_start_tensor,
        y_tensor,
    )

    y_inc = impl.elementwise.add(
        ctx,
        target,
        source_ir,
        f"{name}_y_inc",
        y_tensor,
        constant_1,
    )

    next_y = impl.elementwise.add(
        ctx,
        target,
        source_ir,
        f"{name}_next_y",
        y_start_tensor,
        y_inc,
    )

    y_start_inc = impl.elementwise.add(
        ctx,
        target,
        source_ir,
        f"{name}_y_start_inc",
        y_start_tensor,
        constant_1,
    )
    cond = ge(ctx, target, source_ir, f"{name}_cond", next_y, size_tensor)
    x_output = impl.condition.select(
        ctx,
        target,
        source_ir,
        f"{name}_x_output",
        x_inc,
        x_tensor,
        cond,
    )
    x_recurrence.set_input(1, x_output)

    y_start_current = impl.condition.select(
        ctx,
        target,
        source_ir,
        f"{name}_y_start_current",
        y_start_inc,
        y_start_tensor,
        cond,
    )
    y_start_recurrence.set_input(1, y_start_current)

    y_val = impl.condition.select(
        ctx,
        target,
        source_ir,
        f"{name}_y_val",
        constant_0,
        y_inc,
        cond,
    )
    y_recurrence.set_input(1, y_val)

    loop_output_x = loop.add_loop_output(x_tensor, trt.LoopOutput.CONCATENATE)
    loop_output_y = loop.add_loop_output(y_current, trt.LoopOutput.CONCATENATE)
    loop_output_x.set_input(1, num_loop)
    loop_output_y.set_input(1, num_loop)

    # Cat two N tensors into 2 x N. [0, 0, 0], [1, 2, 3] -> [[0, 0, 0], [1, 2, 3]]
    x_index = impl.shuffle.reshape(
        ctx, target, source_ir, f"{name}_x_index", loop_output_x.get_output(0), (1, -1)
    )
    y_index = impl.shuffle.reshape(
        ctx, target, source_ir, f"{name}_y_index", loop_output_y.get_output(0), (1, -1)
    )

    x_y_tensor = cat(
        ctx,
        target,
        source_ir,
        f"{name}_x_y_tensor",
        [x_index, y_index],
        0,
    )

    # Reshape 2 x N output to N x 2. [[0, 0, 0], [1, 2, 3]] -> [[0, 1], [0, 2], [0, 3]]
    indices_tensor = ctx.net.add_shuffle(x_y_tensor)
    set_layer_name(indices_tensor, target, f"{name}_indices_tensor", source_ir)
    indices_tensor.first_transpose = trt.Permutation([1, 0])
    indices_tensor.reshape_dims = (-1, 2)

    return indices_tensor.get_output(0)


def cdist_forward(
    ctx: ConversionContext,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
    x1: trt.ITensor,
    x2: trt.ITensor,
    p: float,
    compute_mode: Optional[int],
) -> Union[trt.ITensor, Sequence[trt.ITensor]]:
    """
    Computes pairwise distances between sets of vectors in tensors x1 and x2 using the p-norm. The function treats the last dimension
    of x1 and x2 as feature dimensions, which must be identical for both inputs. The second-to-last dimensions can differ, reflecting
    the number of vectors in each tensor. The dimensions preceding the last are considered as batch dimensions, and pairwise distances
    are computed for each matching set in these dimensions.

    The output tensor's shape is derived by matching the batch dimensions of x1 and x2, where the mismatched batch dimensions are
    merged, and the resulting shape reflects the computed distances for each pair of vectors. It's crucial that the batch dimensions
    (except for the size of sets of vectors to compare) of x1 and x2 either match or one of them is 1 (broadcasting).

    Args:
        x1 (Tensor): input tensor of shape B x P x M.
        x2 (Tensor): input tensor of shape B x R x M.
        p (float): p value for the p-norm distance to calculate between each vector pair
        compute_mode (int): Controls the computation method based on the size of the input sets:
            - None ('use_mm_for_euclid_dist_if_necessary'): Default mode. Uses matrix multiplication to calculate
              Euclidean distance (p=2) if either the number of vectors in x1 or x2 exceeds 25 (P > 25 or R > 25).
            - 1 ('use_mm_for_euclid_dist'): Always use matrix multiplication approach to calculate
            euclidean distance (p = 2)
            - 2 ('donot_use_mm_for_euclid_dist'): Never use matrix multiplication approach to calculate
            euclidean distance (p = 2)

    Example:
    - If x1.shape = [2, 3, 10, 5] and x2.shape = [2, 3, 20, 5], both having the same batch dimensions [2, 3], the output shape will be [2, 3, 10, 20].
      This represents computing distances in two batches of three groups, each comparing 10 vectors from x1 with 20 vectors from x2.
    - For x1.shape = [10, 5] (10 vectors, each of 5 features) and x2.shape = [20, 5] (20 vectors, each of 5 features),
      since there are no batch dimensions to match, the output shape is simply [10, 20], comparing all vectors from x1 against all vectors from x2.

    Note: The `compute_mode` parameter is designed to optimize the performance of the Euclidean distance calculation,
    especially useful when working with large datasets. This parameter allows you to control how the distances are computed,
    with different modes available to leverage matrix multiplication for speed improvements.
    """
    if compute_mode is None:
        compute_mode = 0

    x1_expand_shape = list(x1.shape[:-1]) + [1, x1.shape[-1]]
    x2_expand_shape = list(x2.shape[:-2]) + [1] + list(x2.shape[-2:])

    # Reshape x1 and x2 for broadcasting
    x1_expanded = impl.shuffle.reshape(
        ctx, target, source_ir, f"{name}_x1_expand", x1, x1_expand_shape
    )
    x2_expanded = impl.shuffle.reshape(
        ctx, target, source_ir, f"{name}_x2_expand", x2, x2_expand_shape
    )

    diff = impl.elementwise.sub(
        ctx, target, source_ir, f"{name}_diff", x1_expanded, x2_expanded
    )

    if p == 0:
        diff_non_zero = impl.elementwise.ne(
            ctx, target, source_ir, f"{name}_diff_non_zero", diff, 0
        )
        diff_non_zero = cast_trt_tensor(
            ctx, diff_non_zero, torch.float32, f"{name}_cast", target, source_ir
        )
        dist = impl.reduce.sum(
            ctx,
            target,
            source_ir,
            f"{name}_sum",
            diff_non_zero,
            dim=-1,
            keepdim=False,
        )
    elif p == 1:
        abs_val = impl.unary.abs(ctx, target, source_ir, f"{name}_abs_val", diff)
        dist = impl.reduce.sum(
            ctx, target, source_ir, f"{name}_sum", abs_val, dim=-1, keepdim=False
        )
    elif p == 2:
        if (
            compute_mode == 0 and (x1.shape[-2] > 25 or x2.shape[-2] > 25)
        ) or compute_mode == 1:
            # Compute squared elements
            x1_squared = impl.elementwise.pow(
                ctx, target, source_ir, f"{name}_x1_squared", x1, 2
            )
            x2_squared = impl.elementwise.pow(
                ctx, target, source_ir, f"{name}_x2_squared", x2, 2
            )

            # Sum squares along the last dimension
            x1_sum_squared = impl.reduce.sum(
                ctx,
                target,
                source_ir,
                f"{name}_x1_sum",
                x1_squared,
                dim=-1,
                keepdim=True,
            )
            x2_sum_squared = impl.reduce.sum(
                ctx,
                target,
                source_ir,
                f"{name}_x2_sum",
                x2_squared,
                dim=-1,
                keepdim=True,
            )

            # Reshape sums for broadcasting
            rank = len(x2.shape)
            permute_shape = list(range(rank - 2)) + [rank - 1, rank - 2]
            x1_sum_expanded = x1_sum_squared
            x2_sum_expanded = impl.permutation.permute(
                ctx, target, source_ir, f"{name}_permute", x2_sum_squared, permute_shape
            )

            # Compute dot product of x1 and transposed x2
            x2_tr = impl.permutation.permute(
                ctx, target, source_ir, f"{name}_permute_mm", x2, permute_shape
            )
            dot_product = impl.matmul.matrix_multiply(
                ctx,
                target,
                source_ir,
                f"{name}_dot_product",
                x1,
                x2_tr,
                input_matrix_op=trt.MatrixOperation.NONE,
                other_matrix_op=trt.MatrixOperation.NONE,
            )

            # Combine results to get squared distances
            dist_squared = impl.elementwise.add(
                ctx,
                target,
                source_ir,
                f"{name}_dist_squared_initial",
                x1_sum_expanded,
                x2_sum_expanded,
            )
            dist_squared = impl.elementwise.sub(
                ctx,
                target,
                source_ir,
                f"{name}_dist_squared",
                dist_squared,
                impl.elementwise.mul(
                    ctx, target, source_ir, f"{name}_dot_product_scaled", dot_product, 2
                ),
            )

            # Compute the Euclidean distances
            dist = impl.unary.sqrt(ctx, target, source_ir, f"{name}_dist", dist_squared)
        else:
            diff_squared = impl.elementwise.pow(
                ctx, target, source_ir, f"{name}_diff_squared", diff, 2
            )
            dist_squared = impl.reduce.sum(
                ctx,
                target,
                source_ir,
                f"{name}_dist_sq_sum",
                diff_squared,
                dim=-1,
                keepdim=False,
            )
            dist = impl.unary.sqrt(ctx, target, source_ir, f"{name}_sqrt", dist_squared)
    elif 0 < p < 1 or 1 < p < 2 or 2 < p < float("inf"):
        abs_val = impl.unary.abs(ctx, target, source_ir, f"{name}_abs_val", diff)
        pow_val = impl.elementwise.pow(
            ctx, target, source_ir, f"{name}_pow_val_1", abs_val, p
        )
        sum_val = impl.reduce.sum(
            ctx, target, source_ir, f"{name}_sum", pow_val, dim=-1, keepdim=False
        )
        dist = impl.elementwise.pow(
            ctx, target, source_ir, f"{name}_pow_val_2", sum_val, 1 / p
        )
    elif p == float("inf"):
        abs_val = impl.unary.abs(ctx, target, source_ir, f"{name}_abs_val", diff)
        dist = impl.reduce.max(
            ctx,
            target,
            source_ir,
            f"{name}_max",
            abs_val,
            dim=-1,
            keepdim=False,
            return_indices=False,
        )
    return dist