import torch


def tsp_get_min_dist(dist_matrix):
    """
    Takes a distance matrix dist_matrix of shape (n,n) or (bs,n,n)
    Returns the mean minimal distance between two distinct nodes for each of the *bs* instances 
    """
    if not isinstance(dist_matrix,torch.Tensor):
        try:
            dist_matrix = torch.tensor(dist_matrix)
        except Exception:
            raise ValueError(f"Type {type(dist_matrix)} could not be broadcasted to torch.Tensor")
    solo=False
    if len(dist_matrix.shape)==2:
        solo=True
        dist_matrix = dist_matrix.unsqueeze(0)
    bs,n,_ = dist_matrix.shape
    max_dist = torch.max(dist_matrix)
    eye_mask = torch.zeros_like(dist_matrix)
    eye_mask[:] = torch.eye(n)
    dist_modif = dist_matrix + max_dist*eye_mask
    min_dist = torch.min(dist_modif,dim=-1).values
    min_mean = torch.mean(min_dist,dim=-1)
    if solo:
        min_mean = min_mean[0]
    
    return min_mean