import torch


def get_ans_idx(p_s, p_e, max_len=15, no_answer=False):
    """
    Discretize soft predictions (probs) to get start and end indices
    Choose (i, j) which maximizes p1[i]*p2[j], s.t. (i <= j) & (j-i+1 <= max_len)
    Args:
        p_s: [batch_size, c_len], probs for start index
        p_e: [batch_size, c_len], probs for end index
        max_len: max length of the answer prediction
        no_answer (bool): Treat 0-idx as the no-answer prediction. Consider a prediction no-answer
                          if preds[0,0]*preds[0,1] > the prob assigned to the max-probability span
    Returns:
        s_idxs:  [batch_size], hard predictions for start index
        e_idxs: [batch_size], hard predictions for end index
        
    """
    c_len = p_s.shape[1]
    device = p_s.device
    
    if p_s.min() < 0 or p_s.max() > 1 or p_e.min() < 0 or p_e.max() > 1:
        raise ValueError('Expected p_start and p_end to have values in [0, 1]')

    # Compute pairwise probs
    p_s = p_s.unsqueeze(2)  # [batch_size, c_len, 1]
    p_e = p_e.unsqueeze(1)  # [batch_size, 1, c_len]    
    p_join = torch.bmm(p_s, p_e)  # [batch_size, c_len, c_len]

    # Restrict (i, j) s.t. (i <= j) & (j-i+1 <= max_len)
    is_legal_pair = torch.triu(torch.ones((c_len, c_len), device=device))
    is_legal_pair = is_legal_pair - torch.triu(torch.ones((c_len, c_len), device=device), diagonal=max_len)
    if no_answer:
        p_no_answer = p_join[:, 0, 0].clone()
        is_legal_pair[0, :] = 0
        is_legal_pair[:, 0] = 0
    else:
        p_no_answer = None
    p_join = p_join * is_legal_pair

    # Obtain (i, j) which maximizes p_join
    max_each_row, _ = torch.max(p_join, dim=2)  # [batch_size, c_len]
    max_each_col, _ = torch.max(p_join, dim=1)  # [batch_size, c_len]
    s_idxs = torch.argmax(max_each_row, dim=1)  # [batch_size]
    e_idxs = torch.argmax(max_each_col, dim=1)  # [batch_size]

    # Predict no-answer whenever p_no_answer > max_prob
    if no_answer:      
        max_prob, _ = torch.max(max_each_col, dim=1)
        s_idxs[p_no_answer > max_prob] = 0
        e_idxs[p_no_answer > max_prob] = 0

    return s_idxs, e_idxs