# @manual=//deeplearning/trt/python:py_tensorrt
import tensorrt as trt
import torch

from ..converter_registry import tensorrt_converter

from .converter_utils import get_dyn_range, mark_as_int8_layer, to_numpy


def common_linear(network, mod, input_val, layer_name, is_quantized):
    """
    TensorRT fully connected layer implicitly flatten last three dimensions at
    the start and implicitly reshape the result to (K, 1, 1) at the end.

    e.g. If input is (N, C, H, W), first it gets flatten to (N, C*H*W). Then after
    going through fully connected operation, it becomes (N, K). Before sending it
    out, it gets reshaped into (N, K, 1, 1) and this is the final output.

    TODO: We can optimize this to get rid of unneccesary transformation.
    """
    # reshape the input to (*, X, 1, 1)
    layer = network.add_shuffle(input_val)
    layer.reshape_dims = tuple(input_val.shape) + (1, 1)
    layer.name = f"{layer_name}_pre_shuffle"

    if is_quantized:
        mark_as_int8_layer(layer, input_val.dynamic_range)

    kernel = to_numpy(mod.weight if not is_quantized else mod.weight())
    bias = to_numpy(mod.bias if not is_quantized else mod.bias())

    # add fully connected
    layer = network.add_fully_connected(
        input=layer.get_output(0),
        num_outputs=mod.out_features,
        kernel=kernel,
        bias=bias,
    )
    layer.name = f"{layer_name}_linear"

    if is_quantized:
        dyn_range = get_dyn_range(mod.scale, mod.zero_point, torch.quint8)
        mark_as_int8_layer(layer, dyn_range)

    # reshape the output from (*, K, 1, 1) to (*, K)
    layer = network.add_shuffle(layer.get_output(0))
    layer.reshape_dims = tuple(input_val.shape[:-1]) + (mod.out_features,)
    layer.name = f"{layer_name}_post_shuffle"

    if is_quantized:
        mark_as_int8_layer(layer, dyn_range)

    return layer.get_output(0)


@tensorrt_converter(torch.nn.modules.linear.Linear)
def linear(network, submod, args, kwargs, layer_name):
    # args/kwargs should have already been normalized to kwargs
    assert len(args) == 0
    input_val = kwargs["input"]

    if not isinstance(input_val, trt.tensorrt.ITensor):
        raise RuntimeError(
            f"Linear received input {input_val} that is not part "
            "of the TensorRT region!"
        )

    return common_linear(network, submod, input_val, layer_name, is_quantized=False)


@tensorrt_converter(torch.nn.quantized.modules.linear.Linear)
def quantized_linear(network, submod, args, kwargs, layer_name):
    input_val = args[0]

    if not isinstance(input_val, trt.tensorrt.ITensor):
        raise RuntimeError(
            f"Quantized Linear received input {input_val} that is not part "
            "of the TensorRT region!"
        )

    return common_linear(network, submod, input_val, layer_name, is_quantized=True)