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, 10, 8, 68), 'torch.Tensor', 'torch.Tensor', ({(8, 23, 8, 67): '[[0, 0, 0], [2, 0, 0], [0, 2, 0], [2, 2, 0]]'}, {}), '([[0, 0, 0], [2, 0, 0], [0, 2, 0], [2, 2, 0]])', False, 'import torch\n'), ((11, 10, 11, 38), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', ({(11, 28, 11, 31): 'pos', (11, 33, 11, 37): 'face'}, {}), '(pos, face)', False, 'from torch_geometric.utils import geodesic_distance\n'), ((21, 10, 21, 50), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (), '', False, 'from torch_geometric.utils import geodesic_distance\n'), ((30, 10, 30, 36), 'torch.tensor', 'torch.tensor', ({(30, 23, 30, 35): '[0, 0, 0, 0]'}, {}), '([0, 0, 0, 0])', False, 'import torch\n'), ((31, 11, 31, 37), 'torch.tensor', 'torch.tensor', ({(31, 24, 31, 36): '[0, 1, 2, 3]'}, {}), '([0, 1, 2, 3])', False, 'import torch\n'), ((32, 10, 32, 58), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (), '', False, 'from torch_geometric.utils import geodesic_distance\n'), ((36, 10, 36, 52), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (), '', False, 'from torch_geometric.utils import geodesic_distance\n'), ((40, 10, 40, 49), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (), '', False, 'from torch_geometric.utils import geodesic_distance\n'), ((18, 31, 18, 53), 'torch.tensor', 'torch.tensor', ({(18, 44, 18, 52): 'expected'}, {}), '(expected)', False, 'import torch\n'), ((19, 31, 19, 75), 'torch_geometric.utils.geodesic_distance', 'geodesic_distance', (), '', False, 'from torch_geometric.utils import geodesic_distance\n'), ((28, 31, 28, 53), 'torch.tensor', 'torch.tensor', ({(28, 44, 28, 52): 'expected'}, {}), '(expected)', False, 'import torch\n'), ((33, 25, 33, 32), 'math.sqrt', 'sqrt', ({(33, 30, 33, 31): '(2)'}, {}), '(2)', False, 'from math import sqrt\n'), ((34, 31, 34, 53), 'torch.tensor', 'torch.tensor', ({(34, 44, 34, 52): 'expected'}, {}), '(expected)', False, 'import torch\n'), ((37, 25, 37, 32), 'math.sqrt', 'sqrt', ({(37, 30, 37, 31): '(2)'}, {}), '(2)', False, 'from math import sqrt\n'), ((38, 31, 38, 53), 'torch.tensor', 'torch.tensor', ({(38, 44, 38, 52): 'expected'}, {}), '(expected)', False, 'import torch\n'), ((42, 31, 42, 53), 'torch.Tensor', 'torch.Tensor', ({(42, 44, 42, 52): 'expected'}, {}), '(expected)', False, 'import torch\n'), ((9, 11, 9, 47), 'torch.tensor', 'torch.tensor', ({(9, 24, 9, 46): '[[0, 1, 3], [0, 2, 3]]'}, {}), '([[0, 1, 3], [0, 2, 3]])', False, 'import torch\n'), ((13, 18, 13, 25), 'math.sqrt', 'sqrt', ({(13, 23, 13, 24): '(2)'}, {}), '(2)', False, 'from math import sqrt\n'), ((14, 15, 14, 22), 'math.sqrt', 'sqrt', ({(14, 20, 14, 21): '(2)'}, {}), '(2)', False, 'from math import sqrt\n'), ((15, 12, 15, 19), 'math.sqrt', 'sqrt', ({(15, 17, 15, 18): '(2)'}, {}), '(2)', False, 'from math import sqrt\n'), ((16, 9, 16, 16), 'math.sqrt', 'sqrt', ({(16, 14, 16, 15): '(2)'}, {}), '(2)', False, 'from math import sqrt\n'), ((23, 22, 23, 29), 'math.sqrt', 'sqrt', ({(23, 27, 23, 28): '(2)'}, {}), '(2)', False, 'from math import sqrt\n'), ((24, 19, 24, 26), 'math.sqrt', 'sqrt', ({(24, 24, 24, 25): '(2)'}, {}), '(2)', False, 'from math import sqrt\n'), ((25, 16, 25, 23), 'math.sqrt', 'sqrt', ({(25, 21, 25, 22): '(2)'}, {}), '(2)', False, 'from math import sqrt\n'), ((26, 13, 26, 20), 'math.sqrt', 'sqrt', ({(26, 18, 26, 19): '(2)'}, {}), '(2)', False, 'from math import sqrt\n')]