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', (['[[0, 0, 0], [2, 0, 0], [0, 2, 0], [2, 2, 0]]'], {}), False, 'import torch\n'), (11, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {}), False, 'from torch_geometric.utils import geodesic_distance\n'), (21, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {'norm': '(False)'}), False, 'from torch_geometric.utils import geodesic_distance\n'), (30, 'torch.tensor', 'torch.tensor', (['[0, 0, 0, 0]'], {}), False, 'import torch\n'), (31, 'torch.tensor', 'torch.tensor', (['[0, 1, 2, 3]'], {}), False, 'import torch\n'), (32, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {'src': 'src', 'dest': 'dest'}), False, 'from torch_geometric.utils import geodesic_distance\n'), (36, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {'src': 'src[0:1]'}), False, 'from torch_geometric.utils import geodesic_distance\n'), (40, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {'dest': 'dest'}), False, 'from torch_geometric.utils import geodesic_distance\n'), (18, 'torch.tensor', 'torch.tensor', (['expected'], {}), False, 'import torch\n'), (19, 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (['pos', 'face'], {'num_workers': '(-1)'}), False, 'from torch_geometric.utils import geodesic_distance\n'), (28, 'torch.tensor', 'torch.tensor', (['expected'], {}), False, 'import torch\n'), (33, 'math.sqrt', 'sqrt', (['(2)'], {}), False, 'from math import sqrt\n'), (34, 'torch.tensor', 'torch.tensor', (['expected'], {}), False, 'import torch\n'), (37, 'math.sqrt', 'sqrt', (['(2)'], {}), False, 'from math import sqrt\n'), (38, 'torch.tensor', 'torch.tensor', (['expected'], {}), False, 'import torch\n'), (42, 'torch.Tensor', 'torch.Tensor', (['expected'], {}), False, 'import torch\n'), (9, 'torch.tensor', 'torch.tensor', (['[[0, 1, 3], [0, 2, 3]]'], {}), False, 'import torch\n'), (13, 'math.sqrt', 'sqrt', (['(2)'], {}), False, 'from math import sqrt\n'), (14, 'math.sqrt', 'sqrt', (['(2)'], {}), False, 'from math import sqrt\n'), (15, 'math.sqrt', 'sqrt', (['(2)'], {}), False, 'from math import sqrt\n'), (16, 'math.sqrt', 'sqrt', (['(2)'], {}), False, 'from math import sqrt\n'), (23, 'math.sqrt', 'sqrt', (['(2)'], {}), False, 'from math import sqrt\n'), (24, 'math.sqrt', 'sqrt', (['(2)'], {}), False, 'from math import sqrt\n'), (25, 'math.sqrt', 'sqrt', (['(2)'], {}), False, 'from math import sqrt\n'), (26, 'math.sqrt', 'sqrt', (['(2)'], {}), False, 'from math import sqrt\n')]