import torch


def _get_idx_worst(name, ref, hyp):
    """
    Computes the batch indices for which the input metric value is worse than the current metric value.

    Parameters
    ----------
    name : str
        'mse', 'psnr', 'ssim', 'lpips', or 'fvd'. Metric to consider. For 'mse', 'fvd' and 'lpips', lower is better,
        while for 'psnr' and 'ssim', higher is better.
    ref : torch.*.Tensor
        One-dimensional tensor containing a list of current metric values.
    hyp : torch.*.Tensor
        One-dimensional tensor containing a list of new metric values to be compared agains ref. Must be of the same
        length as ref.

    Returns
    -------
    torch.*.LongTensor
        List of indices i for which the value hyp[i] is worse than the value ref[i].
    """
    if name in ('mse', 'fvd', 'lpips'):
        return torch.nonzero(hyp > ref, as_tuple=False).flatten()
    if name in ('psnr', 'ssim'):
        return torch.nonzero(hyp < ref, as_tuple=False).flatten()
    raise ValueError(f'Metric \'{name}\' not yet implemented')