import torch


def query_feature(feature_map, query_uv):
    """
    query features from feature map
    :param feature_map: B,C,res1,res2
    :param query_uv: B,2,K,1 in [0,1]
    :return B,C,K,1
    """
    assert float(query_uv.max()) < 1 + 1e-9
    assert query_uv.shape[1] == 2
    res1 = feature_map.shape[2]
    res2 = feature_map.shape[3]
    query_index = query_uv.clone()
    query_index[:, 0, ...] = torch.clamp((query_uv[:, 0, ...] * res1).long(), 0, res1 - 1)
    query_index[:, 1, ...] = torch.clamp((query_uv[:, 1, ...] * res2).long(), 0, res2 - 1)
    if query_index.ndimension() > 3:
        index = query_index.squeeze(3)  # B,2,K
    else:
        index = query_index  # B*2*K
    # cvt to 1D index
    index = index[:, 0, :] * feature_map.shape[3] + index[:, 1, :]  # B,K
    index = index.unsqueeze(2).repeat(1, 1, feature_map.shape[1])
    flatten_feature_map = feature_map.reshape(feature_map.shape[0],
                                              feature_map.shape[1], -1).permute(0, 2, 1)  # B,N,C
    query = torch.gather(flatten_feature_map, 1, index.long()).contiguous()  # B,K,C
    query = query.permute(0, 2, 1).unsqueeze(3)
    return query