import numpy as np
import torch
from line_profiler import profile
class TrtRunner:
    """
    TensorRT runner using the new I/O tensor API (engine.num_io_tensors).

    Features:
    - Multiple inputs / multiple outputs
    - Inputs can be torch.cuda tensors (zero-copy bind) or numpy (H2D copy)
    - Outputs are torch.cuda tensors (TRT writes directly into them)
    - Optional: execute on current torch CUDA stream (recommended for torch-only pipeline)
    """
    def __init__(self, engine_path, input_names=None, output_names=None, logger_severity=None, use_torch_stream=False):
        import tensorrt as trt

        self.trt = trt
        self.use_torch_stream = bool(use_torch_stream)

        if logger_severity is None:
            logger_severity = trt.Logger.ERROR
        self.logger = trt.Logger(logger_severity)

        with open(engine_path, "rb") as f, trt.Runtime(self.logger) as rt:
            self.engine = rt.deserialize_cuda_engine(f.read())
        if self.engine is None:
            raise RuntimeError(f"Failed to deserialize engine: {engine_path}")

        if not hasattr(self.engine, "num_io_tensors"):
            raise RuntimeError("This runner requires TensorRT 'new' I/O tensor API (engine.num_io_tensors).")

        self.context = self.engine.create_execution_context()
        if self.context is None:
            raise RuntimeError("Failed to create execution context.")

        # Discover IO tensor names
        io_names = [self.engine.get_tensor_name(i) for i in range(self.engine.num_io_tensors)]
        all_in = [n for n in io_names if self.engine.get_tensor_mode(n) == trt.TensorIOMode.INPUT]
        all_out = [n for n in io_names if self.engine.get_tensor_mode(n) == trt.TensorIOMode.OUTPUT]

        if not all_in:
            raise RuntimeError("No INPUT tensors found in engine.")
        if not all_out:
            raise RuntimeError("No OUTPUT tensors found in engine.")

        # User selection / ordering
        self.input_names = list(input_names) if input_names is not None else all_in
        self.output_names = list(output_names) if output_names is not None else all_out

        for n in self.input_names:
            if n not in all_in:
                raise ValueError(f"Requested input '{n}' not in engine inputs: {all_in}")
        for n in self.output_names:
            if n not in all_out:
                raise ValueError(f"Requested output '{n}' not in engine outputs: {all_out}")

        if not self.use_torch_stream:
            import pycuda.driver as cuda
            import pycuda.autoinit  # only here!
            self.cuda = cuda
            self.stream = self.cuda.Stream()
        else:
            self.cuda = torch.cuda
            self.stream = None


        # Caches
        self._last_in_shapes = {}            # name -> shape tuple
        self._d_in = {}                      # name -> device ptr (only used when input is numpy)
        self._out_torch = {}                 # name -> torch.cuda tensor
        self._out_torch_shape = {}           # name -> shape tuple
        self._out_torch_dtype = {}           # name -> torch dtype

    def _np_dtype_to_torch(self, np_dtype):
        dt = np.dtype(np_dtype)
        m = {
            np.dtype(np.float32): torch.float32,
            np.dtype(np.float16): torch.float16,
            np.dtype(np.int32): torch.int32,
            np.dtype(np.int8): torch.int8,
            np.dtype(np.uint8): torch.uint8,
            np.dtype(np.bool_): torch.bool,
        }
        if dt not in m:
            raise TypeError(f"Unsupported dtype for torch output: {dt}")
        return m[dt]

    def _ensure_inputs_and_bind(self, inputs):
        """
        inputs: dict(name->tensor/ndarray) OR list/tuple aligned with self.input_names
        Binds input addresses, sets shapes for dynamic inputs.
        """
        trt = self.trt
        cuda = self.cuda

        if isinstance(inputs, (list, tuple)):
            if len(inputs) != len(self.input_names):
                raise ValueError(f"Got {len(inputs)} inputs, expected {len(self.input_names)}: {self.input_names}")
            in_dict = dict(zip(self.input_names, inputs))
        elif isinstance(inputs, dict):
            in_dict = inputs
        else:
            raise TypeError("inputs must be a dict{name: tensor} or list/tuple aligned with input_names")

        # For each input: set shape (dynamic) + bind pointer
        for name in self.input_names:
            if name not in in_dict:
                raise KeyError(f"Missing required input '{name}'. Required: {self.input_names}")
            x = in_dict[name]

            if isinstance(x, torch.Tensor):
                if not x.is_cuda:
                    raise ValueError(f"Input '{name}' is torch.Tensor but not CUDA.")
                if not x.is_contiguous():
                    # This makes a copy; if you need strict zero-copy, ensure caller passes contiguous.
                    x = x.contiguous()

                # dtype check (optional but recommended)
                exp_np = trt.nptype(self.engine.get_tensor_dtype(name))
                exp_torch = self._np_dtype_to_torch(exp_np)
                if x.dtype != exp_torch:
                    raise TypeError(f"Input '{name}' dtype {x.dtype} != engine expects {exp_torch}")

                shape = tuple(x.shape)

                # Set dynamic shape
                if hasattr(self.context, "set_input_shape"):
                    if self._last_in_shapes.get(name) != shape:
                        self.context.set_input_shape(name, shape)
                        self._last_in_shapes[name] = shape
                else:
                    raise RuntimeError("context.set_input_shape not available; unexpected TRT version.")

                # ZERO-COPY bind: TRT reads directly from torch tensor memory
                self.context.set_tensor_address(name, int(x.data_ptr()))

                # Store back possibly-contiguous tensor for caller consistency
                in_dict[name] = x

            else:
                print('NUMPY')
                # numpy path: allocate device buffer + H2D copy, then bind device ptr
                x_np = np.ascontiguousarray(x)
                shape = tuple(x_np.shape)

                if hasattr(self.context, "set_input_shape"):
                    if self._last_in_shapes.get(name) != shape:
                        self.context.set_input_shape(name, shape)
                        self._last_in_shapes[name] = shape
                else:
                    raise RuntimeError("context.set_input_shape not available; unexpected TRT version.")

                in_dtype = trt.nptype(self.engine.get_tensor_dtype(name))
                if x_np.dtype != in_dtype:
                    x_np = x_np.astype(in_dtype, copy=False)

                nbytes = int(np.prod(shape)) * np.dtype(in_dtype).itemsize
                if (name not in self._d_in) or (self._d_in[name] is None):
                    self._d_in[name] = cuda.mem_alloc(nbytes)

                # H2D copy on pycuda stream
                cuda.memcpy_htod_async(self._d_in[name], x_np, self.stream)
                self.context.set_tensor_address(name, int(self._d_in[name]))
                in_dict[name] = x_np

        return in_dict

    def _ensure_outputs_and_bind(self):
        """Alloc torch CUDA outputs according to resolved shapes and bind their pointers."""
        trt = self.trt

        for out_name in self.output_names:
            out_shape = tuple(self.context.get_tensor_shape(out_name))
            if any(s < 0 for s in out_shape):
                raise RuntimeError(f"Dynamic output shape not resolved for '{out_name}': {out_shape}")

            out_np_dtype = trt.nptype(self.engine.get_tensor_dtype(out_name))
            out_torch_dtype = self._np_dtype_to_torch(out_np_dtype)

            need = (
                (out_name not in self._out_torch) or
                (self._out_torch_shape.get(out_name) != out_shape) or
                (self._out_torch_dtype.get(out_name) != out_torch_dtype)
            )
            if need:
                self._out_torch[out_name] = torch.empty(out_shape, device="cuda", dtype=out_torch_dtype)
                self._out_torch_shape[out_name] = out_shape
                self._out_torch_dtype[out_name] = out_torch_dtype

            self.context.set_tensor_address(out_name, int(self._out_torch[out_name].data_ptr()))

    @profile
    def run(self, inputs, return_dict=False, synchronize=False):
        """
        inputs:
          - dict {name: torch.cuda.Tensor OR numpy array}
          - or list/tuple aligned with self.input_names
        return:
          - list of torch.cuda.Tensor outputs (default) in self.output_names order
          - or dict if return_dict=True

        synchronize:
          - If True, synchronize stream before returning. If you run TRT on torch stream
            and immediately consume outputs in torch ops, you usually do NOT need this.
        """
        # Bind inputs + set shapes
        in_dict = self._ensure_inputs_and_bind(inputs)

        # Bind outputs (depends on resolved output shapes after setting inputs)
        self._ensure_outputs_and_bind()

        # Choose stream:
        # - torch stream handle (best if your pipeline is torch-only)
        # - otherwise pycuda stream handle (needed for numpy H2D copies above)
        if self.use_torch_stream:
            # Ensure any pending torch writes to input tensors are visible on this stream
            # If inputs were produced on current stream, this is naturally ordered.
            stream_handle = torch.cuda.current_stream().cuda_stream
            ok = self.context.execute_async_v3(stream_handle)
            if not ok:
                raise RuntimeError("TensorRT execution failed (execute_async_v3 returned False).")

            if synchronize:
                torch.cuda.current_stream().synchronize()
        else:
            ok = self.context.execute_async_v3(self.stream.handle)
            if not ok:
                raise RuntimeError("TensorRT execution failed (execute_async_v3 returned False).")
            if synchronize:
                self.stream.synchronize()

        if return_dict:
            return {n: self._out_torch[n] for n in self.output_names}
        return [self._out_torch[n] for n in self.output_names]

    def get_io_names(self):
        return list(self.input_names), list(self.output_names)

    def warmup(self, shapes_by_name, dtypes_by_name=None, iters=10):
        """
        shapes_by_name: dict input_name -> shape tuple
        dtypes_by_name: dict input_name -> torch dtype (default float32)
        """
        dtypes_by_name = dtypes_by_name or {}
        inputs = {}
        for n, sh in shapes_by_name.items():
            dt = dtypes_by_name.get(n, torch.float32)
            inputs[n] = torch.zeros(sh, device="cuda", dtype=dt)
        for _ in range(iters):
            _ = self.run(inputs, synchronize=False)


if __name__ == "__main__":

    troch_type = torch.float32
    num = 24
    dummy_inputs = {
        "net":  torch.randn(1, num, 128, 43, 77, device="cuda", dtype=troch_type).contiguous(),
        "inp":  torch.randn(1, num, 128, 43, 77, device="cuda", dtype=troch_type).contiguous(),
        "corr": torch.randn(1, num, 196, 43, 77, device="cuda", dtype=troch_type).contiguous(),
        "flow": torch.randn(1, num,   4, 43, 77, device="cuda", dtype=troch_type).contiguous(),
        # "ii":   torch.arange(num, device="cuda", dtype=torch.int64).contiguous(),
        # "jj":   torch.arange(num, device="cuda", dtype=torch.int64).contiguous(),
    }


    runner = TrtRunner("update_module_partial_fp16.engine")

    # optional: print names
    print(runner.get_io_names())

    # warmup
    runner.warmup({"net": (1, num, 128, 43, 77),
                "inp": (1, num, 128, 43, 77),
                "corr": (1, num, 196, 43, 77),
                "flow": (1, num,   4, 43, 77),
                }
                , dtypes_by_name={"net": troch_type,
                                    "inp": troch_type,
                                    "corr": troch_type,
                                    "flow": troch_type,
                                    }, iters=5)
    # dummy_np = np.random.randn(1, 1, 3, 344, 616).astype(np.float32)
    # run many times
    import time
    begin_time = time.time()
    for _ in range(3717):
        y_list = runner.run(dummy_inputs)          # list outputs
        # print(y_list[0].shape)
        # # or
        # y_dict = runner.run(dummy_np, return_dict=True)
    end_time = time.time()
    print(f"Time taken: {end_time - begin_time} seconds")


    # runner = TrtRunner("droidnet_fnet_fp32.engine")

    # # optional: print names
    # print(runner.get_io_names())

    # # warmup
    # runner.warmup({"input": (1, 1, 3, 344, 616)}, dtypes_by_name={"input": torch.float32}, iters=5)
    # # dummy_np = np.random.randn(1, 1, 3, 344, 616).astype(np.float32)
    # dummy = torch.randn(1, 1, 3, 344, 616, device="cuda", dtype=torch.float32).contiguous()
    # # run many times
    # import time
    # begin_time = time.time()
    # for _ in range(2914):
    #     y_list = runner.run({"input": dummy})          # list outputs
    #     # print(y_list[0].shape)
    #     # # or
    #     # y_dict = runner.run(dummy_np, return_dict=True)
    # end_time = time.time()
    # print(f"Time taken: {end_time - begin_time} seconds")