import torch
from torch import Tensor


def pairwise_distances(x: Tensor) -> Tensor:
    """Compute pairwise euclidean distances.

    ---
    Args:
        x: Vector of coordinates.
            Shape of [n_points, coord_dim].

    ---
    Returns:
        Pairwise distances.
            Shape of [n_points, n_points].
    """
    d = (x[:, None, :] - x[None, :, :]) ** 2
    return torch.sqrt(d.sum(dim=-1))


def solution_cost(x: Tensor, s: Tensor) -> Tensor:
    """Compute the cost of the given solution.

    ---
    Args:
        x: Cities coordinates.
            Shape of [n_cities, (x, y)].
        s: Solution.
            Shape of [n_cities,].

    ---
    Returns:
        The solution cost.
    """
    x = x[s]
    x_ = torch.roll(x, -1, dims=0)
    return torch.norm(x - x_, p=2, dim=-1).sum()
