import torch
from torch.cuda import nvtx
from collections import OrderedDict
import numpy as np
from polygraphy.backend.common import bytes_from_path
from polygraphy import util
from polygraphy.backend.trt import ModifyNetworkOutputs, Profile
from polygraphy.backend.trt import (
    engine_from_bytes,
    engine_from_network,
    network_from_onnx_path,
    save_engine,
)
from polygraphy.logger import G_LOGGER
import tensorrt as trt
from logging import error, warning
from tqdm import tqdm
import copy
from collections import OrderedDict
from typing import List, Dict
from cuda import cudart
from polygraphy import cuda

TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
G_LOGGER.module_severity = G_LOGGER.ERROR

# Map of numpy dtype -> torch dtype
numpy_to_torch_dtype_dict = {
    np.uint8: torch.uint8,
    np.int8: torch.int8,
    np.int16: torch.int16,
    np.int32: torch.int32,
    np.int64: torch.int64,
    np.float16: torch.float16,
    np.float32: torch.float32,
    np.float64: torch.float64,
    np.complex64: torch.complex64,
    np.complex128: torch.complex128,
}
if np.version.full_version >= "1.24.0":
    numpy_to_torch_dtype_dict[np.bool_] = torch.bool
else:
    numpy_to_torch_dtype_dict[np.bool] = torch.bool

# Map of torch dtype -> numpy dtype
torch_to_numpy_dtype_dict = {
    value: key for (key, value) in numpy_to_torch_dtype_dict.items()
}

# https://github.com/Jeff-LiangF/streamv2v/blob/18c1a3bd56ff348d54a3300605936980bb13b03c/src/streamv2v/acceleration/tensorrt/utilities.py
def CUASSERT(cuda_ret):
    err = cuda_ret[0]
    if err != cudart.cudaError_t.cudaSuccess:
        raise RuntimeError(
            f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t"
        )
    if len(cuda_ret) > 1:
        return cuda_ret[1]
    return None

class TQDMProgressMonitor(trt.IProgressMonitor):
    def __init__(self):
        trt.IProgressMonitor.__init__(self)
        self._active_phases = {}
        self._step_result = True
        self.max_indent = 5

    def phase_start(self, phase_name, parent_phase, num_steps):
        leave = False
        try:
            if parent_phase is not None:
                nbIndents = (
                    self._active_phases.get(parent_phase, {}).get(
                        "nbIndents", self.max_indent
                    )
                    + 1
                )
                if nbIndents >= self.max_indent:
                    return
            else:
                nbIndents = 0
                leave = True
            self._active_phases[phase_name] = {
                "tq": tqdm(
                    total=num_steps, desc=phase_name, leave=leave, position=nbIndents
                ),
                "nbIndents": nbIndents,
                "parent_phase": parent_phase,
            }
        except KeyboardInterrupt:
            # The phase_start callback cannot directly cancel the build, so request the cancellation from within step_complete.
            _step_result = False

    def phase_finish(self, phase_name):
        try:
            if phase_name in self._active_phases.keys():
                self._active_phases[phase_name]["tq"].update(
                    self._active_phases[phase_name]["tq"].total
                    - self._active_phases[phase_name]["tq"].n
                )

                parent_phase = self._active_phases[phase_name].get("parent_phase", None)
                while parent_phase is not None:
                    self._active_phases[parent_phase]["tq"].refresh()
                    parent_phase = self._active_phases[parent_phase].get(
                        "parent_phase", None
                    )
                if (
                    self._active_phases[phase_name]["parent_phase"]
                    in self._active_phases.keys()
                ):
                    self._active_phases[
                        self._active_phases[phase_name]["parent_phase"]
                    ]["tq"].refresh()
                del self._active_phases[phase_name]
            pass
        except KeyboardInterrupt:
            _step_result = False

    def step_complete(self, phase_name, step):
        try:
            if phase_name in self._active_phases.keys():
                self._active_phases[phase_name]["tq"].update(
                    step - self._active_phases[phase_name]["tq"].n
                )
            return self._step_result
        except KeyboardInterrupt:
            # There is no need to propagate this exception to TensorRT. We can simply cancel the build.
            return False

class Engine:
    def __init__(
        self,
        engine_path,
    ):
        self.engine_path = engine_path
        self.engine = None
        self.context = None
        self.buffers = OrderedDict()
        self.tensors = OrderedDict()
        self.cuda_graph_instance = None  # cuda graph

    def __del__(self):
        del self.engine
        del self.context
        del self.buffers
        del self.tensors

    def reset(self, engine_path=None):
        del self.engine
        del self.context
        del self.buffers
        del self.tensors
        self.engine_path = engine_path

        self.buffers = OrderedDict()
        self.tensors = OrderedDict()
        self.inputs = {}
        self.outputs = {}

    def build(
        self,
        onnx_path,
        fp16,
        input_profile=None,
        enable_refit=False,
        enable_preview=False,
        enable_all_tactics=False,
        timing_cache=None,
        update_output_names=None,
    ):
        print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")
        p = [Profile()]
        if input_profile:
            p = [Profile() for i in range(len(input_profile))]
            for _p, i_profile in zip(p, input_profile):
                for name, dims in i_profile.items():
                    assert len(dims) == 3
                    _p.add(name, min=dims[0], opt=dims[1], max=dims[2])

        config_kwargs = {}
        if not enable_all_tactics:
            config_kwargs["tactic_sources"] = []

        network = network_from_onnx_path(
            onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]
        )
        if update_output_names:
            print(f"Updating network outputs to {update_output_names}")
            network = ModifyNetworkOutputs(network, update_output_names)

        builder = network[0]
        config = builder.create_builder_config()
        config.progress_monitor = TQDMProgressMonitor()

        config.set_flag(trt.BuilderFlag.FP16) if fp16 else None
        config.set_flag(trt.BuilderFlag.REFIT) if enable_refit else None

        profiles = copy.deepcopy(p)
        for profile in profiles:
            # Last profile is used for set_calibration_profile.
            calib_profile = profile.fill_defaults(network[1]).to_trt(
                builder, network[1]
            )
            config.add_optimization_profile(calib_profile)

        try:
            engine = engine_from_network(
                network,
                config,
            )
        except Exception as e:
            error(f"Failed to build engine: {e}")
            return 1
        try:
            save_engine(engine, path=self.engine_path)
        except Exception as e:
            error(f"Failed to save engine: {e}")
            return 1
        return 0

    def load(self):
        self.engine = engine_from_bytes(bytes_from_path(self.engine_path))

    def activate(self, reuse_device_memory=None):
        if reuse_device_memory:
            self.context = self.engine.create_execution_context_without_device_memory()
        #    self.context.device_memory = reuse_device_memory
        else:
            self.context = self.engine.create_execution_context()

    def allocate_buffers(self, shape_dict=None, device="cuda"):
        nvtx.range_push("allocate_buffers")
        for idx in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(idx)
            binding = self.engine[idx]
            if shape_dict and binding in shape_dict:
                shape = shape_dict[binding]["shape"]
            else:
                shape = self.context.get_tensor_shape(name)

            dtype = trt.nptype(self.engine.get_tensor_dtype(name))
            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self.context.set_input_shape(name, shape)
            tensor = torch.empty(
                tuple(shape), dtype=numpy_to_torch_dtype_dict[dtype]
            ).to(device=device)
            self.tensors[binding] = tensor
        nvtx.range_pop()

    def infer(self, feed_dict, stream, use_cuda_graph=False):
        for name, buf in feed_dict.items():
            self.tensors[name].copy_(buf)

        for name, tensor in self.tensors.items():
            self.context.set_tensor_address(name, tensor.data_ptr())

class MultiStreamEngine(Engine):
    def __init__(self, engine_path, num_streams=2):
        super().__init__(engine_path)
        self.num_streams = num_streams
        self.streams = []
        self.contexts = []
        self.stream_tensors = []

    def set_num_streams(self, value):
        self.num_streams = value

    def activate(self):
        super().activate()
        for _ in range(self.num_streams):
            # stream = torch.cuda.Stream()
            stream = cuda.Stream()
            context = self.engine.create_execution_context()
            self.streams.append(stream)
            self.contexts.append(context)
            self.stream_tensors.append(OrderedDict())

    def get_torch_dtype(self, trt_dtype):
        """Convert TensorRT dtype to PyTorch dtype."""
        dtype_map = {
            trt.int8: torch.int8,
            trt.int32: torch.int32,
            trt.float16: torch.float16,
            trt.float32: torch.float32,
            trt.bool: torch.bool,
        }
        return dtype_map.get(trt_dtype, torch.float32)  # Default to float32 if not found

    def allocate_buffers(self, shape_dict=None, device="cuda"):
        nvtx.range_push("allocate_buffers")
        for stream_idx in range(self.num_streams):
            tensors = self.stream_tensors[stream_idx]
            context = self.contexts[stream_idx]
            
            for idx in range(self.engine.num_io_tensors):
                name = self.engine.get_tensor_name(idx)
                binding = self.engine[idx]
                if shape_dict and binding in shape_dict:
                    shape = shape_dict[binding]["shape"]
                else:
                    shape = context.get_tensor_shape(name)

                dtype = trt.nptype(self.engine.get_tensor_dtype(name))
                if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                    context.set_input_shape(name, shape)
                
                torch_dtype = numpy_to_torch_dtype_dict[dtype]
                tensor = torch.empty(tuple(shape), dtype=torch_dtype, device=device)
                tensors[binding] = tensor
        nvtx.range_pop()

    def infer(self, feed_dicts: List[Dict[str, torch.Tensor]], use_cuda_graph=True):
        results = []
        num_batches = len(feed_dicts)
        
        for i in range(0, num_batches, self.num_streams):
            batch_results = []
            
            for j in range(self.num_streams):
                if i + j >= num_batches:
                    break
                
                feed_dict = feed_dicts[i + j]
                stream = self.streams[j]
                context = self.contexts[0]
                tensors = self.stream_tensors[0]

                for name, buf in feed_dict.items():
                    tensors[name].copy_(buf)

                if (i == 0):
                    for name, tensor in tensors.items():
                        context.set_tensor_address(name, tensor.data_ptr())
                    
                if use_cuda_graph:
                    if self.cuda_graph_instance is not None:
                        CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream.ptr))
                        CUASSERT(cudart.cudaStreamSynchronize(stream.ptr))
                    else:
                        # do inference before CUDA graph capture
                        noerror = context.execute_async_v3(stream.ptr)
                        if not noerror:
                            raise ValueError("ERROR: inference failed.")
                        # capture cuda graph
                        CUASSERT(
                            cudart.cudaStreamBeginCapture(stream.ptr, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal)
                        )
                        context.execute_async_v3(stream.ptr)
                        self.graph = CUASSERT(cudart.cudaStreamEndCapture(stream.ptr))
                        self.cuda_graph_instance = CUASSERT(cudart.cudaGraphInstantiate(self.graph, 0))
                else:
                    noerror = self.context.execute_async_v3(stream.ptr)
                    if not noerror:
                        raise ValueError("ERROR: inference failed.")

                # success = context.execute_async_v3(stream.cuda_stream)
                # if not success:
                #     raise RuntimeError(f"Inference failed for batch {i + j}")

                batch_results.append({name: tensor for name, tensor in tensors.items() 
                                        if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT})

            # Synchronize all used streams
            for j in range(min(self.num_streams, num_batches - i)):
                self.streams[j].synchronize()

            results.extend(batch_results)

        return results