import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'vigs'))
from modules.droid_net import DroidNet
from collections import OrderedDict
import torch
droidnet = DroidNet()
weights = "pretrained_models/droid.pth"
state_dict = OrderedDict([
    (k.replace("module.", ""), v) for (k, v) in torch.load(weights).items()])
state_dict["update.weight.2.weight"] = state_dict["update.weight.2.weight"][:2]
state_dict["update.weight.2.bias"] = state_dict["update.weight.2.bias"][:2]
state_dict["update.delta.2.weight"] = state_dict["update.delta.2.weight"][:2]
state_dict["update.delta.2.bias"] = state_dict["update.delta.2.bias"][:2]
droidnet.load_state_dict(state_dict)
droidnet.to("cuda:0").eval()

# NOTE: Best practice: export to ONNX in fp32, then convert to fp16 in TRT.
# Input/output remain fp32; only internal computation runs in fp16.
# Otherwise, sometimes will have errors in TRT if ONNX directly in fp16.
# NOTE: dynamic resolution is used here. For a fixed input size, remove the dynamic H, W axes for better performance.

# Export feature encoder (fnet)
dummy = torch.randn(1, 1, 3, 344, 616, device="cuda", dtype=torch.float32)
torch.onnx.export(
    droidnet.fnet,
    dummy,
    "pretrained_models/droidnet_fnet.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input":  {3: "H", 4: "W"},      # N, T, C, H, W
        "output": {3: "H_out", 4: "W_out"}
    },
    opset_version=18,
)

# Export the GRU-based iterative update module.
# NOTE: 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.
num = 28
B, C, H, W = 1, 128, 43, 77

net  = torch.randn(B, num, 128, H, W, device="cuda", dtype=torch.float32)
inp  = torch.randn(B, num, 128, H, W, device="cuda", dtype=torch.float32)
corr = torch.randn(B, num, 196, H, W, device="cuda", dtype=torch.float32)
flow = torch.randn(B, num,   4, H, W, device="cuda", dtype=torch.float32)

torch.onnx.export(
    droidnet.update,
    (net, inp, corr, flow),
    "pretrained_models/update_module_partial.onnx",
    opset_version=17,
    input_names=["net", "inp", "corr", "flow"],
    output_names=["net_out", "delta", "weight"],
    dynamic_axes={
        "net":  {1: "num", 3: "H", 4: "W"},
        "inp":  {1: "num", 3: "H", 4: "W"},
        "corr": {1: "num", 3: "H", 4: "W"},
        "flow": {1: "num", 3: "H", 4: "W"},

        "net_out": {1: "num", 3: "H", 4: "W"},
        "delta":   {1: "num", 3: "H", 4: "W"},
        "weight":  {1: "num", 3: "H", 4: "W"},
    }
)