import torch


def get_loss_cumu(loss_dict, cumu_mode):
    """Combine different losses to obtain a single scalar loss.

    Args:
        loss_dict: A dictionary or list of loss values, each of which is a torch scalar.
        cumu_mode: a 2-tuple. Choose from:
            ("generalized-mean"/"gm", {order}): generalized mean with order
            "harmonic": harmonic mean
            "geometric": geometric mean
            "mean": arithmetic mean
            "sum": summation
            "min": minimum
            "original": returns the original loss_dict.

    Returns:
        loss: the combined loss scalar computed according to cumu_mode.
    """
    if cumu_mode == "original":
        return loss_dict
    if isinstance(loss_dict, dict):
        loss_list = torch.stack([loss for loss in loss_dict.values()])
    elif isinstance(loss_dict, list):
        loss_list = torch.stack(loss_dict)
    elif isinstance(loss_dict, torch.Tensor):
        loss_list = loss_dict
    else:
        raise
    N = len(loss_list)
    if N == 1:
        return loss_list[0]
    epsilon = 1e-20  # to prevent NaN
    if cumu_mode.startswith("gm"):
        cumu_mode_str, num = cumu_mode.split("-")
        cumu_mode = (cumu_mode_str, eval(num))
    if isinstance(cumu_mode, tuple) and cumu_mode[0] in ["generalized-mean", "gm"]:
        if cumu_mode[1] == -1:
            cumu_mode = "harmonic"
        elif cumu_mode[1] == 0:
            cumu_mode = "geometric"
        elif cumu_mode[1] == 1:
            cumu_mode = "mean"
    
    if cumu_mode == "harmonic":
        loss = N / (1 / (loss_list + epsilon)).sum()
    elif cumu_mode == "geometric":
        loss = (loss_list + epsilon).prod() ** (1 / float(N))
    elif cumu_mode == "mean":
        loss = loss_list.mean()
    elif cumu_mode == "sum":
        loss = loss_list.sum()
    elif cumu_mode == "min":
        loss = loss_list.min()
    elif cumu_mode[0] in ["generalized-mean", "gm"]:
        order = cumu_mode[1]
        loss = (((loss_list + epsilon) ** order).mean()) ** (1 / float(order))
    else:
        raise
    return loss