import torch


def calc_accs(logits, targs, categories=None, prepender=""):
    """
    Calculates the average accuracy over the batch for each possible
    category

    Args:
        logits: torch float tensor (B, N, K)
            the model predictions. the last dimension must be the
            same number of dimensions as possible target values.
        targs: torch long tensor (B, N)
            the targets for the predictions
        categories: torch long tensor (B, N) or None
            if None, this value is ignored. Otherwise it specifies
            categories for accuracy calculations.
        prepender: str
            a string to prepend to all keys in the accs dict
    Returns:
        accs: dict
            keys: str
                <prepender>_acc: float
                    the average accuracy over all categories
                <prepender>_acc_<category>: float
                    the average accuracy over this particular
                    category. for example, if one of the categories
                    is named 1, the key will be "1" and the value
                    will be the average accuracy over that
                    particular category.
    """
    if prepender!="" and prepender[-1]!="_": prepender = prepender+"_"
    prepender = prepender + "acc"
    logits = logits.reshape(-1, logits.shape[-1])
    try:
        argmaxes = torch.argmax(logits, dim=-1).reshape(-1)
    except:
        print("logits:", logits)
        return { prepender: 0 }
    targs = targs.reshape(-1)
    acc = (argmaxes.long()==targs.long()).float()
    accs = {
        prepender: acc.mean().item()
    }
    if len(argmaxes) == 0: return accs
    pre = prepender + "lbl_"
    targ_types = {*targs.cpu().data.numpy()}
    for t in targ_types:
        idxs = targs==t
        if idxs.float().sum() == 0: continue
        accs[pre+str(t)] = acc[idxs].mean().item()
    if type(categories) == torch.Tensor: # (B, N)
        categories = categories.reshape(-1).data.long()
        pre = prepender + "ctg_"
        cats = {*categories.numpy()}
        for cat in cats:
            idxs = categories==cat
            if idxs.float().sum() <= 0: continue
            accs[pre+str(cat)] = acc[idxs].mean().item()
    return accs