from torch.utils.data._utils.collate import default_collate
from typing import Sequence
from typing import Dict


def multi_supervision_collate_fn(batch: Sequence[Dict]) -> Dict:
    """
    Custom collate_fn for K2SpeechRecognitionDataset.

    It merges the items provided by K2SpeechRecognitionDataset into the following structure:

    .. code-block::

        {
            'features': float tensor of shape (B, T, F)
            'supervisions': [
                {
                    'sequence_idx': Tensor[int] of shape (S,)
                    'text': List[str] of len S
                    'start_frame': Tensor[int] of shape (S,)
                    'num_frames': Tensor[int] of shape (S,)
                }
            ]
        }

    Dimension symbols legend:
    * ``B`` - batch size (number of Cuts),
    * ``S`` - number of supervision segments (greater or equal to B, as each Cut may have multiple supervisions),
    * ``T`` - number of frames of the longest Cut
    * ``F`` - number of features
    """

    dataset_idx_to_batch_idx = {
        example['supervisions'][0]['sequence_idx']: batch_idx
        for batch_idx, example in enumerate(batch)
    }

    def update(d: Dict, **kwargs) -> Dict:
        for key, value in kwargs.items():
            d[key] = value
        return d

    supervisions = default_collate([
        update(sup, sequence_idx=dataset_idx_to_batch_idx[sup['sequence_idx']])
        for example in batch
        for sup in example['supervisions']
    ])
    feats = default_collate([example['features'] for example in batch])
    return {
        'features': feats,
        'supervisions': supervisions
    }