import torch


def AverageProportion(
    delays, src_lens, tgt_lens, ref_len=None, target_padding_mask=None):
    """
    Function to calculate Average Proportion from
    Can neural machine translation do simultaneous translation?
    (https://arxiv.org/abs/1606.02012)
    Delays are monotonic steps, range from 1 to src_len.
    Give src x tgt y, AP is calculated as:
    AP = 1 / (|x||y]) sum_i^|Y| deleys_i
    """
    if target_padding_mask is not None:
        AP = torch.sum(delays.masked_fill(target_padding_mask, 0), dim=1)
    else:
        AP = torch.sum(delays, dim=1)

    AP = AP / (src_lens * tgt_lens)
    return AP