"""Utility functions for GraphBolt."""

import numpy as np
import torch


def _read_torch_data(path):
    return torch.load(path)


def _read_numpy_data(path, in_memory=True):
    if in_memory:
        return torch.from_numpy(np.load(path))
    return torch.as_tensor(np.load(path, mmap_mode="r+"))


def read_data(path, fmt, in_memory=True):
    """Read data from disk."""
    if fmt == "torch":
        return _read_torch_data(path)
    elif fmt == "numpy":
        return _read_numpy_data(path, in_memory=in_memory)
    else:
        raise RuntimeError(f"Unsupported format: {fmt}")


def tensor_to_tuple(data):
    """Split a torch.Tensor in column-wise to a tuple."""
    assert isinstance(data, torch.Tensor), "data must be a torch.Tensor"
    return tuple(data.t())
