from typing import Optional, Sequence, Union

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 (
    get_positive_dim,
    set_layer_name,
)


def squeeze(
    ctx: ConversionContext,
    target: Target,
    source_ir: Optional[SourceIR],
    name: str,
    input: TRTTensor,
    dim: Optional[Union[int, Sequence[int]]] = None,
) -> TRTTensor:
    # Squeeze with dim=None would only work in explicit batch dim mode without any dynamic
    # dim, which is a very rare case. For now we just claim not supporting dim=None.
    assert dim is not None, "We don't support dim=None right now for squeeze."
    dims = []

    if isinstance(dim, int):
        dims.append(dim)
    else:
        for d in dim:
            dims.append(d)

    new_dims = []
    for dim in dims:
        dim = get_positive_dim(
            dim,
            len(input.shape),
        )

        assert input.shape[dim] != -1, "We don't support squeeze dynamic dim."
        new_dims.append(dim)

    dim_to_remove = []
    new_permutation = []
    for i, s in enumerate(input.shape):
        if (i in new_dims) and s == 1:
            dim_to_remove.append(i)
        else:
            new_permutation.append(i)
    # If number of reshape dimensions is less than input, 0s are resolved by aligning
    # the most significant dimensions of input
    output_shape = tuple([0] * len(new_permutation))
    new_permutation += dim_to_remove

    layer = ctx.net.add_shuffle(input)
    layer.first_transpose = new_permutation
    layer.reshape_dims = output_shape
    set_layer_name(layer, target, name, source_ir)
    return layer.get_output(0)