import torch


def farthest_point_sample(coordinates, num_output_points, norm_order=2., deterministic=False):
    """Subsamples coordinates by randomly selecting a first coordinate and then iteratively
    choosing the num_output_points - 1 remaining coordinates as the furthest from previously
    chosen coordinates.

    Args:
        coordinates (torch.Tensor): Shape (batch_size, num_points, coordinate_dim).
        num_output_points (int): Number of points to select among coordinates.
        norm_order (float): Order of the norm to use to measure distance. Defaults to L2 norm.
        deterministic (bool): If True, uses first point in batch as initial point, otherwise
            randomly samples initial point.

    Returns:
        subsampled_idx: Subsampled coordinate indices of shape (batch_size, num_output_points).
    """
    device = coordinates.device
    batch_size, num_points, coordinate_dim = coordinates.shape
    # Initialize indices of points to be selected
    subsampled_idx = torch.zeros(batch_size, num_output_points, dtype=torch.long).to(device)
    # Initialize distances to very large number
    distances = torch.ones(batch_size, num_points).to(device) * 1e10
    # Select initial "farthest" point
    if deterministic:
        # Select first point in batch as initial point
        farthest = torch.zeros(batch_size, dtype=torch.long).to(device)
    else:
        # Randomly select initial point for each batch batch element
        farthest = torch.randint(0, num_points, (batch_size,), dtype=torch.long).to(device)
    batch_indices = torch.arange(batch_size, dtype=torch.long).to(device)
    for i in range(num_output_points):
        # Set selected index to be farthest from previously selected indices
        subsampled_idx[:, i] = farthest
        # Select current farthest centroid
        centroid = coordinates[batch_indices, farthest, :].view(batch_size, 1, coordinate_dim)
        # Calculate distance from current centroid to all other points
        current_dist = torch.linalg.norm(coordinates - centroid, dim=-1, ord=norm_order)
        # Update distances with newly calculated distances
        mask = current_dist < distances
        distances[mask] = current_dist[mask]
        # Select farthest point with updated distances (note torch.max returns both values
        # and indices, so we select indices)
        farthest = torch.max(distances, dim=-1)[1]
    return subsampled_idx