from dataclasses import dataclass
from enum import Enum
from typing import Dict, List, Set

import numpy as np
import tensorrt as trt
import torch

from tensorrt_llm._common import _is_building
from tensorrt_llm._utils import (trt_dtype_to_np, trt_dtype_to_str,
                                 trt_dtype_to_torch)
from tensorrt_llm.logger import logger

from .pipeline_graph import PipelineGraph
from .utils import (get_builder_flags, get_cache_key, get_sorted_layer_ids,
                    set_trt_network, to_base_class_layer, to_subclass_layer,
                    to_trt_weights)


class ShapeType(Enum):
    MIN = 0
    OPT = 1
    MAX = 2


_trt_to_type_dict = {
    trt.int64: int,
    trt.bool: bool,
}


def get_shape_layers(trt_network):
    shape_layers = set()
    for i in range(trt_network.num_layers):
        layer = trt_network.get_layer(i)
        if (layer.num_inputs > 0 and np.all([
                layer.get_input(j).is_shape_tensor
                for j in range(layer.num_inputs)
                if layer.get_input(j) is not None
        ])) or (layer.num_outputs > 0 and np.all([
                layer.get_output(j).is_shape_tensor
                for j in range(layer.num_outputs)
        ])):
            shape_layers.add(layer.name)
    return shape_layers


def get_layers_in_shape_network(trt_network, shape_layers, sorted_layer_ids):
    layers = set()
    shape_tensors = set()
    for layer_id in reversed(sorted_layer_ids):
        layer = trt_network.get_layer(layer_id)
        in_shape_network = False
        if layer.name in shape_layers:
            in_shape_network = True
        else:
            for j in range(layer.num_outputs):
                output = layer.get_output(j)
                if output.name in shape_tensors:
                    in_shape_network = True
                    break
        if in_shape_network:
            layers.add(layer.name)
            for j in range(layer.num_inputs):
                input = layer.get_input(j)
                if input is not None:
                    shape_tensors.add(input.name)
    return layers


def get_shape_network(trt_network,
                      shapes,
                      values,
                      sorted_layer_ids,
                      profile=None,
                      shape_type: ShapeType = ShapeType.OPT):
    shape_layers = get_shape_layers(trt_network)
    layers_in_shape_network = get_layers_in_shape_network(
        trt_network, shape_layers, sorted_layer_ids)
    shape_graph = PipelineGraph.create_graph()
    shape_network = shape_graph.as_trt()
    shape_builder = shape_network.builder
    shape_profile = shape_builder.create_optimization_profile()
    for i in range(trt_network.num_inputs):
        input = trt_network.get_input(i)
        shapes[input.name] = input.shape
        new_input = shape_graph.add_input(input)
        if profile is not None:
            if -1 in input.shape:
                shape = profile.get_shape(input.name)
                shape = shape[shape_type.value]
                shapes[input.name] = shape
                new_input.raw_shape = shape
            if input.is_shape_tensor:
                shape_values = profile.get_shape_input(input.name)
                value = shape_values[shape_type.value]
                values[input.name] = value
                shape_profile.set_shape_input(input.name, value, value, value)
    output_mapping = {}
    for layer_id in sorted_layer_ids:
        layer = trt_network.get_layer(layer_id)
        if layer.name in shape_layers:
            new_layer = shape_graph.add_layer(layer)
            for i in range(layer.num_outputs):
                output = layer.get_output(i)
                if output.dtype == trt.DataType.BOOL:
                    proxy_layer = shape_network.add_cast(
                        new_layer.as_trt().get_output(i),
                        trt.DataType.INT32,
                    )
                    proxy_output = proxy_layer.get_output(0)
                    shape_graph.register_layer(proxy_layer)
                    shape_graph.add_output_shape(proxy_output)
                    output_mapping[proxy_output.name] = (output.name,
                                                         output.dtype)
                else:
                    shape_graph.add_output_shape(output)
        elif layer.name in layers_in_shape_network:
            if layer.type == trt.LayerType.CONSTANT:
                shape_graph.add_input(layer.get_output(0))
            else:
                shape_graph.add_layer(layer)
    return shape_network, shape_profile, shape_layers, output_mapping


def get_per_layer_graph(
    layer,
    shapes,
    values,
    updated_attrs=None,
    is_shape_io: bool = None,
):
    graph = PipelineGraph.create_graph()
    network = graph.as_trt()
    is_shape_layer = layer.num_inputs != 0
    for i in range(layer.num_inputs):
        input = layer.get_input(i)
        if input is not None:
            shape = shapes[input.name]
            if (values.get(input.name) is not None
                    and not isinstance(values[input.name], torch.Tensor)):
                value = values[input.name]
                weights = np.asarray(value, dtype=trt_dtype_to_np(input.dtype))
                weights = to_trt_weights(weights)
                input_layer = network.add_constant(shape, weights)
                new_input = input_layer.get_output(0)
                new_input.name = input.name
                graph.register_layer(input_layer)
            elif graph.get_input(input.name) is None:
                new_input = graph.add_input(input)
                new_input.raw_shape = shapes[input.name]
                is_shape_layer = False
    new_layer = graph.add_layer(
        layer,
        updated_attrs=updated_attrs,
    )
    output_mapping = {}
    if layer.type == trt.LayerType.SHAPE:
        is_shape_layer = True
    if layer.num_inputs == 0:
        is_shape_layer = False
    if is_shape_io is not None:
        is_shape_layer = is_shape_io
    for i in range(layer.num_outputs):
        output = layer.get_output(i)
        value = values.get(output.name)
        if value is not None and isinstance(value, torch.Tensor):
            is_output_shape = False
        elif is_shape_layer:
            is_output_shape = True
        else:
            is_output_shape = False
        if is_output_shape:
            if output.dtype == trt.DataType.BOOL:
                proxy_layer = network.add_cast(
                    new_layer.as_trt().get_output(i),
                    trt.DataType.INT32,
                )
                proxy_output = proxy_layer.get_output(0)
                graph.register_layer(proxy_layer)
                output_mapping[proxy_output.name] = (output.name, output.dtype)
                output = proxy_output
            graph.add_output_shape(output)
        else:
            graph.add_output(output)
    return graph, output_mapping


@_is_building
def infer_shapes(network, shapes, values, profile=None):
    if network.num_outputs == 0:
        return
    builder = network.builder
    config = builder.create_builder_config()
    config.builder_optimization_level = 0
    config.flags = get_builder_flags()
    profile = profile or builder.create_optimization_profile()
    config.add_optimization_profile(profile)
    plan = builder.build_serialized_network(network, config)
    if plan is None:
        raise RuntimeError(
            'Engine building failed when inferring shapes, please check the error log.'
        )
    runtime = trt.Runtime(logger.trt_logger)
    engine = runtime.deserialize_cuda_engine(plan)
    context = engine.create_execution_context()
    for i in range(network.num_inputs):
        input = network.get_input(i)
        if input.is_shape_tensor:
            value = values[input.name]
            context.set_shape_input(engine[input.name], value)
    for i in range(network.num_outputs):
        output = network.get_output(i)
        shape = context.get_tensor_shape(output.name)
        shapes[output.name] = shape
        if output.is_shape_tensor:
            if shape == [0]:
                values[output.name] = []
            else:
                if shape == []:
                    shape = [1]
                value = torch.empty(
                    list(shape),
                    dtype=trt_dtype_to_torch(output.dtype),
                    device="cpu",
                )
                values[output.name] = value
                context.set_tensor_address(output.name, value.data_ptr())
    context.infer_shapes()
    assert context.all_binding_shapes_specified
    for i in range(network.num_outputs):
        output = network.get_output(i)
        if isinstance(values.get(output.name), torch.Tensor):
            values[output.name] = values[output.name].tolist()


@dataclass
class ShapeInfo:
    shapes: Dict[str, trt.Dims]
    values: Dict[str, List[int]]
    shape_layers: Set[str]
    max_shapes: Dict[str, trt.Dims] = None


def set_constant_value(layer, values):
    to_subclass_layer(layer)
    output_name = layer.get_output(0).name
    weights = layer.weights
    if isinstance(weights, trt.Weights):
        weights = weights.numpy()
    values[output_name] = list(weights)
    to_base_class_layer(layer)


def infer_per_layer_shapes(
    layer: trt.ILayer,
    shapes,
    values,
    cache=None,
    is_shape_io=False,
):
    if layer.type == trt.LayerType.CONSTANT:
        to_subclass_layer(layer)
        output_name = layer.get_output(0).name
        shape = layer.shape
        shapes[output_name] = shape
        if is_shape_io:
            set_constant_value(layer, values)
        to_base_class_layer(layer)
        return
    elif layer.type == trt.LayerType.SHAPE:
        input_name = layer.get_input(0).name
        output_name = layer.get_output(0).name
        shape = [*shapes[input_name]]
        shapes[output_name] = trt.Dims([len(shape)])
        values[output_name] = shape
        return
    if cache is not None:
        cache_key = get_cache_key(layer, shapes, values)
        if cache_key in cache:
            output_shapes, output_values = cache[cache_key]
            for i in range(layer.num_outputs):
                output = layer.get_output(i)
                shapes[output.name] = output_shapes[i]
                if output_values[i] is not None:
                    values[output.name] = output_values[i]
            return
    graph, output_mapping = get_per_layer_graph(layer, shapes, values)
    dtypes = [
        trt_dtype_to_str(layer.get_input(i).dtype)
        for i in range(layer.num_inputs)
    ]
    layer_info = (f"type={cache_key[0]}, "
                  f"attrs={dict(cache_key[1])}, "
                  f"dtypes={dtypes}, "
                  f"shapes={list(cache_key[2])}, "
                  f"values={list(cache_key[3])}")
    logger.debug(f"infer shapes for layer {layer.name} ({layer_info})")
    try:
        infer_shapes(graph.as_trt(), shapes, values)
    except RuntimeError as e:
        raise RuntimeError(
            f"infer shapes failed for layer {layer.name} ({layer_info})") from e
    for proxy_output, (output, dtype) in output_mapping.items():
        shapes[output] = shapes[proxy_output]
        del shapes[proxy_output]
        if proxy_output in values:
            values[output] = [
                *map(_trt_to_type_dict[dtype], values[proxy_output])
            ]
            del values[proxy_output]
    if cache is not None:
        logger.debug(
            f"shape inference cache miss, layer: {layer.name}, cache key: {cache_key}"
        )
        output_shapes = []
        output_values = []
        for i in range(layer.num_outputs):
            output = layer.get_output(i)
            output_shapes.append(shapes[output.name])
            output_values.append(values.get(output.name))
        cache[cache_key] = (output_shapes, output_values)


def get_shape_info(trt_network, profile, shape_type: ShapeType = ShapeType.OPT):
    shapes = {}
    values = {}
    sorted_layer_ids = get_sorted_layer_ids(trt_network)
    infer_shape_layers = False

    shape_network, shape_profile, shape_layers, output_mapping = get_shape_network(
        trt_network,
        shapes,
        values,
        sorted_layer_ids,
        profile=profile,
        shape_type=shape_type)
    try:
        infer_shapes(shape_network, shapes, values, shape_profile)
        for proxy_output, (output, dtype) in output_mapping.items():
            shapes[output] = shapes[proxy_output]
            values[output] = [
                *map(_trt_to_type_dict[dtype], values[proxy_output])
            ]
            del shapes[proxy_output]
            del values[proxy_output]
    except RuntimeError:
        infer_shape_layers = True

    cache = {}
    for layer_id in sorted_layer_ids:
        layer = trt_network.get_layer(layer_id)
        is_shape_io = layer.name in shape_layers
        if is_shape_io and not infer_shape_layers:
            continue
        set_trt_network(layer, trt_network)
        infer_per_layer_shapes(layer,
                               shapes,
                               values,
                               cache,
                               is_shape_io=is_shape_io)
    return ShapeInfo(shapes, values, shape_layers)