import onnxruntime
import torch

# ZeroGPU containers ship CUDA 13; the onnxruntime-gpu wheel links against
# libcublasLt.so.12 (CUDA 12) → CUDAExecutionProvider fails to load. Use CPU.
providers = ['CPUExecutionProvider']

def load_onnx(file_path: str):
    assert file_path.endswith(".onnx")
    sess_opt = onnxruntime.SessionOptions()
    ort_session = onnxruntime.InferenceSession(file_path, sess_opt=sess_opt, providers=providers)
    return ort_session


def load_onnx_caller(file_path: str, single_output=False):
    ort_session = load_onnx(file_path)
    def caller(*args):
        torch_input = isinstance(args[0], torch.Tensor)
        if torch_input:
            torch_input_dtype = args[0].dtype
            torch_input_device = args[0].device
            # check all are torch.Tensor and have same dtype and device
            assert all([isinstance(arg, torch.Tensor) for arg in args]), "All inputs should be torch.Tensor, if first input is torch.Tensor"
            assert all([arg.dtype == torch_input_dtype for arg in args]), "All inputs should have same dtype, if first input is torch.Tensor"
            assert all([arg.device == torch_input_device for arg in args]), "All inputs should have same device, if first input is torch.Tensor"
            args = [arg.cpu().float().numpy() for arg in args]
        
        ort_inputs = {ort_session.get_inputs()[idx].name: args[idx] for idx in range(len(args))}
        ort_outs = ort_session.run(None, ort_inputs)
        
        if torch_input:
            ort_outs = [torch.tensor(ort_out, dtype=torch_input_dtype, device=torch_input_device) for ort_out in ort_outs]
        
        if single_output:
            return ort_outs[0]
        return ort_outs
    return caller