import torch
import vox2seq


if __name__ == "__main__":
    RES = 256
    coords = torch.meshgrid(torch.arange(RES), torch.arange(RES), torch.arange(RES))
    coords = torch.stack(coords, dim=-1).reshape(-1, 3).int().cuda()
    code_z_cuda = vox2seq.encode(coords, mode='z_order')
    code_z_pytorch = vox2seq.pytorch.encode(coords, mode='z_order')
    code_h_cuda = vox2seq.encode(coords, mode='hilbert')
    code_h_pytorch = vox2seq.pytorch.encode(coords, mode='hilbert')
    assert torch.equal(code_z_cuda, code_z_pytorch)
    assert torch.equal(code_h_cuda, code_h_pytorch)

    code = torch.arange(RES**3).int().cuda()
    coords_z_cuda = vox2seq.decode(code, mode='z_order')
    coords_z_pytorch = vox2seq.pytorch.decode(code, mode='z_order')
    coords_h_cuda = vox2seq.decode(code, mode='hilbert')
    coords_h_pytorch = vox2seq.pytorch.decode(code, mode='hilbert')
    assert torch.equal(coords_z_cuda, coords_z_pytorch)
    assert torch.equal(coords_h_cuda, coords_h_pytorch)

    print("All tests passed.")