import torch


def get_hard_feat_map(logits, spixel_feat):
    """
    calculate a feature map after hard assignment (Not used)

    Args:
        logits: torch.Tensor
            A Tensor of shape (b, nSpixel, h, w)
            that could be regarded as a soft assignment matrix
        spixel_feat: torch.Tensor
            A Tensor of shape (b, c, nSpixel)

    Returns:
        H: torch.Tensor
            A Tensor of shape (b, c, h*w)
            where each c-dim vector is the feature of pixel's clustering center
        label: torch.Tensor
            A Tensor of shape (b, h, w)
            that indicates #spixel of each pixel
    """
    b, nSpixel, h, w = logits.shape
    c = spixel_feat.shape[1]
    label = logits.argmax(1).view(b, -1).long()
    b_idx = torch.arange(b)[:, None].repeat(1, h * w)
    H = spixel_feat[b_idx, :, label].squeeze(-2)
    H = H.permute(0, 2, 1).contiguous()

    return H, label.view(b, h, w)