import torch
from typing import Tuple


def get_tot_objf_and_finite_mask(tot_scores: torch.Tensor, reduction: str) -> Tuple[torch.Tensor, torch.Tensor]:
    """Figures out the total score(log-prob) over all successful supervision segments
    (i.e. those for which the total score wasn't -infinity).
        Args:
            tot_scores: a Torch tensor of shape (num_segments,) containing total scores
                       from forward-backward
            reduction: a reduction type ('mean', 'sum' or 'none')
        Returns:
             Returns a tuple of 2 scalar tensors: (tot_score, finite_mask)
        where finite_mask is a tensor containing successful segment mask.
    
    Based on get_tot_objf_and_num_frames
    from https://github.com/k2-fsa/snowfall/blob/master/snowfall/objectives/common.py
    """
    finite_mask = ~torch.isnan(tot_scores) & torch.ne(tot_scores, -float("inf"))
    if reduction == "mean":
        tot_scores = tot_scores[finite_mask].mean()
    elif reduction == "sum":
        tot_scores = tot_scores[finite_mask].sum()
    return tot_scores, finite_mask