import torch
from modules.trt_base import TrtRunner
from line_profiler import profile


class OmnidataTrtRunner:
    """TensorRT-accelerated Omnidata depth and normal estimator.

    Runs two separate engines (depth + normal) on the same input image.
    Falls back to the PyTorch OmnidataModel if the .engine files are missing.
    """

    def __init__(self):
        self.depth_runner  = TrtRunner("pretrained_models/omnidata_depth_512_simplified_fp16.engine",  use_torch_stream=True)
        self.normal_runner = TrtRunner("pretrained_models/omnidata_normal_512_simplified_fp16.engine", use_torch_stream=True)

    @classmethod
    def try_create(cls):
        try:
            return cls()
        except Exception as e:
            print(f"[TRT] OmnidataTrtRunner not available ({e}), falling back to PyTorch.")
            return None

    def __call__(self, image):
        with torch.cuda.amp.autocast(enabled=False):
            image_f32 = image.to(torch.float32)
            depth  = self.depth_runner.run({"input": image_f32})[0]
            normal = self.normal_runner.run({"input": image_f32})[0]
            return depth, normal


class FeatureEncoderTrtRunner:
    """TensorRT-accelerated DroidNet feature encoder (fnet).

    Falls back to the PyTorch fnet if the .engine file is missing.
    """

    def __init__(self):
        self.runner = TrtRunner("pretrained_models/droidnet_fnet_fp16.engine", use_torch_stream=True)

    @classmethod
    def try_create(cls):
        try:
            return cls()
        except Exception as e:
            print(f"[TRT] FeatureEncoderTrtRunner not available ({e}), falling back to PyTorch.")
            return None

    def __call__(self, image):
        with torch.cuda.amp.autocast(enabled=False):
            return self.runner.run({"input": image.to(torch.float32)})[0]


class UpdateModuleTRTRunner:
    """TensorRT-accelerated DroidNet update operator (partial).

    The graph-aggregation step (self.agg) cannot be exported to TRT, so the
    update operator is split: the GRU trunk runs in TRT, then self.agg runs
    in PyTorch on the TRT output.

    Falls back to the full PyTorch update operator if the .engine file is missing.
    """

    def __init__(self, net, pgba=False):
        engine = ("pretrained_models/update_module_partial_pgba_fp16.engine" if pgba
                  else "pretrained_models/update_module_partial_fp16.engine")
        self.runner = TrtRunner(engine, use_torch_stream=True)
        self.agg = net.update.agg

    @classmethod
    def try_create(cls, net, pgba=False):
        try:
            return cls(net, pgba=pgba)
        except Exception as e:
            print(f"[TRT] UpdateModuleTRTRunner not available ({e}), falling back to PyTorch.")
            return None

    @profile
    def __call__(self, net, inp, corr, flow=None, ii=None, jj=None):
        # Called in two ways:
        #   1. with flow, ii, jj  (implemented below)
        #   2. without flow, ii, jj  (TODO: not yet implemented; PyTorch fallback is fast enough)
        if flow is not None:
            net  = net.to(torch.float32)
            inp  = inp.to(torch.float32)
            corr = corr.to(torch.float32)
            flow = flow.to(torch.float32)
            with torch.cuda.amp.autocast(enabled=False):
                net, delta, weight = self.runner.run({"net": net, "inp": inp, "corr": corr, "flow": flow})
            eta, upmask = self.agg(net, ii.to(net.device))
            net     = net.to(torch.float16)
            delta   = delta.to(torch.float16)
            weight  = weight.to(torch.float16)
            upmask  = upmask.to(torch.float16)
            return net, delta, weight, eta, upmask
