import torch


def dcg(
    outputs: torch.Tensor,
    targets: torch.Tensor,
    k=10,
    gain_function="pow_rank",
) -> torch.Tensor:
    """
    Computes DCG@topk for the specified values of `k`.
    Graded relevance as a measure of  usefulness,
    or gain, from examining a set of items.
    Gain may be reduced at lower ranks.
    Reference:
    https://en.wikipedia.org/wiki/Discounted_cumulative_gain

    Args:
        outputs (torch.Tensor): model outputs, logits
            with shape [batch_size; slate_length]
        targets (torch.Tensor): ground truth, labels
            with shape [batch_size; slate_length]
        gain_function:
            String indicates the gain function for the ground truth labels.
            Two options available:
            - `pow_rank`: torch.pow(2, x) - 1
            - `rank`: x
            On the default, `pow_rank` is used
            to emphasize on retrievng the relevant documents.
        k (int):
            Parameter fro evaluation on top-k items

    Returns:
        torch.Tensor for dcg at k

    Raises:
        ValueError: gain function can be either `pow_rank` or `rank`
    """
    k = min(outputs.size(1), k)
    order = torch.argsort(outputs, descending=True, dim=-1)
    true_sorted_by_preds = torch.gather(targets, dim=-1, index=order)

    if gain_function == "pow_rank":
        gain_function = lambda x: torch.pow(2, x) - 1
        gains = gain_function(true_sorted_by_preds)
        discounts = torch.tensor(1) / torch.log2(
            torch.arange(true_sorted_by_preds.shape[1], dtype=torch.float)
            + 2.0
        )
        discounted_gains = (gains * discounts)[:, :k]

    elif gain_function == "rank":
        discounts = torch.tensor(1) / torch.log2(
            torch.arange(true_sorted_by_preds.shape[1], dtype=torch.float)
            + 1.0
        )
        discounts[0] = 1
        discounted_gains = (true_sorted_by_preds * discounts)[:, :k]

    else:
        raise ValueError("gain function can be either pow_rank or rank")

    sum_dcg = torch.sum(discounted_gains, dim=1)
    return sum_dcg