from pathlib import Path
import sys
import torch
import torch.optim

sys.path.append(str(Path(__file__).absolute().parent.parent))

import model

torch.set_grad_enabled(False)

m = model.load(
    "/home/dengnc/dvs/data/classroom/_nets/train_hr_pano_t0.8/_hr_snerf/checkpoint_50.tar").eval().to("cuda")
print(m.cores[0])
inputs = (
    torch.rand(10, 63, device="cuda"),
    torch.rand(10, 24, device="cuda")
)
def fn(*args, **kwargs):
    return m.cores[0].infer(*args, **kwargs)
sm = torch.jit.trace(fn, inputs)
torch.nn.Module.__call__
print(sm.infer(torch.rand(5, 63, device="cuda"), torch.rand(5, 24, device="cuda")))
sm.save("test.pt")

torch.onnx.export(sm.infer,               # model being run
                  inputs,                    # model input (or a tuple for multiple inputs)
                  "core_0.onnx",      # where to save the model
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=10,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names=["x", "d"],   # the model's input names
                  output_names=["densities", "colors"],  # the model's output names
                  dynamic_axes={
                      "x": [0],
                      "d": [0],
                      "densities": [0],
                      "colors": [0]
                  })  # variable length axes