from math import sqrt

import torch
from torch_geometric.utils import geodesic_distance


def test_geodesic_distance():
    pos = torch.Tensor([[0, 0, 0], [2, 0, 0], [0, 2, 0], [2, 2, 0]])
    face = torch.tensor([[0, 1, 3], [0, 2, 3]]).t()

    out = geodesic_distance(pos, face)
    expected = [
        [0, 1, 1, sqrt(2)],
        [1, 0, sqrt(2), 1],
        [1, sqrt(2), 0, 1],
        [sqrt(2), 1, 1, 0],
    ]
    assert torch.allclose(out, torch.tensor(expected))
    assert torch.allclose(out, geodesic_distance(pos, face, num_workers=-1))

    out = geodesic_distance(pos, face, norm=False)
    expected = [
        [0, 2, 2, 2 * sqrt(2)],
        [2, 0, 2 * sqrt(2), 2],
        [2, 2 * sqrt(2), 0, 2],
        [2 * sqrt(2), 2, 2, 0],
    ]
    assert torch.allclose(out, torch.tensor(expected))

    src = torch.tensor([0, 0, 0, 0])
    dest = torch.tensor([0, 1, 2, 3])
    out = geodesic_distance(pos, face, src=src, dest=dest)
    expected = [0, 1, 1, sqrt(2)]
    assert torch.allclose(out, torch.tensor(expected))

    out = geodesic_distance(pos, face, src=src[0:1])
    expected = [0, 1, 1, sqrt(2)]
    assert torch.allclose(out, torch.tensor(expected))

    out = geodesic_distance(pos, face, dest=dest)
    expected = [0, 0, 0, 0]
    assert torch.allclose(out, torch.Tensor(expected))

[(8, 'torch.Tensor', 'torch.Tensor', False, 'import torch\n'), (11, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', False, 'from torch_geometric.utils import geodesic_distance\n'), (21, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', False, 'from torch_geometric.utils import geodesic_distance\n'), (30, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (31, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (32, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', False, 'from torch_geometric.utils import geodesic_distance\n'), (36, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', False, 'from torch_geometric.utils import geodesic_distance\n'), (40, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', False, 'from torch_geometric.utils import geodesic_distance\n'), (18, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (19, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', False, 'from torch_geometric.utils import geodesic_distance\n'), (28, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (33, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n'), (34, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (37, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n'), (38, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (42, 'torch.Tensor', 'torch.Tensor', False, 'import torch\n'), (9, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (13, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n'), (14, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n'), (15, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n'), (16, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n'), (23, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n'), (24, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n'), (25, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n'), (26, 'math.sqrt', 'sqrt', False, 'from math import sqrt\n')]