
    :param batch_span:          This is a list of span.     The size of the list should match the batch_size
    :param max_span_length:     An optional argument that specific the max_padding value
                                Start is inclusive, and end is exclusive.
                                Remember this is just one span object indicating a start and end, not a list of span
    :return:
    ouput_seqs:         This is a span of text selected by the span from input_seq.
    output_seqs_l:      This is the length of the selected span
    """
    batch_size = input_p.size(0)
    assert batch_size == len(batch_span)

    output_seq_list = []
    output_seq_l = []
    for b_i in range(batch_size):
        cur_p: torch.Tensor = input_p[b_i]  # [T, *]
        cur_span_index = batch_span[b_i]
        cur_selected_span = cur_p[cur_span_index[0]:cur_span_index[1]]
        output_seq_list.append(cur_selected_span)
        output_seq_l.append(cur_selected_span.size(0))

    output_sents = torch_util.pack_list_sequence(output_seq_list, output_seq_l, max_span_length)
    output_l = input_p.new_tensor(output_seq_l).long()

    return output_sents, output_l


class ParagraphSpan(object):
    def __init__(self, span_list: List[Tuple[int, int]], total_length: int = None,
                 is_consecutive=True) -> None:
        """
        :param span_list:   unwrapped raw span list
        :param total_length:    total length of the paragraph for data validation.
        :param is_consecutive:  whether to check if the span if consecutive.
        """
        super().__init__()
        # Check whether it is valid span list:
        pre_end_ind = 0
        cur_length = 0
        self._span_lengths = []
        for start_ind, end_ind in span_list:
            span_length = end_ind - start_ind
            assert span_length >= 0

            if is_consecutive:
                assert start_ind == pre_end_ind

            self._span_lengths.append(span_length)
            cur_length += span_length
            pre_end_ind = end_ind

        if total_length is not None:
            assert cur_length == total_length

        self.span_list = span_list
        self._total_length = cur_length
        # self._span_lengths =

    def __iter__(self):
        return self.span_list.__iter__()

    def __repr__(self):
        return self.span_list.__repr__()

    def number_of_span(self):
        return len(self.span_list)

    @property
    def span_lengths(self):
        return self._span_lengths

    @property
    def total_length(self):
        return self._total_length


def cut_paragraph_to_sentence(input_p: torch.Tensor, p_l: torch.Tensor, batch_span: List[ParagraphSpan],
                              max_sentence_length: int = None) \
        -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[ParagraphSpan]]:
    """
    :param input_p: The input paragraph Tensor. [B (batch_size: number of paragraph), T, D]
    :param p_l:     The length of paragraph Tensor. [B] This is for checking the length
    :param batch_span:    A List of The span of the paragraph for each example in the batch.
    :param max_sentence_length:     Use this value to truncate the max length of each sentence to avoid out-of-memory

            Note: If you set max_sentence_length, then we will not be able to map the token back to the original
                Paragraph. The method will be in-revertable. Maybe fix this in the future.
    :return:
        output_sents:        The new packed sentence vector.
                            Shape: [B (new batch size: number of spans in the original batch), T, D]

        output_l:            The lengths of the outputs.
                            Shape: [B]

        batched_l:           This is the new paragraph length after sentence truncate for original paragraph format.
                            Shape: [B]

        batch_modified_span: This is the span after sentence truncate.

        if max_sentence_length is None: the batched_l will remain the same, and valid batch span will also be the same.
    """
    batch_size = input_p.size(0)
    dim = input_p.size(-1)
    assert batch_size == len(batch_span)
    assert batch_size == p_l.size(0)

    sequence_tensors_list: List[torch.Tensor] = []  # The list to save all the output sequence tensors.
    sequence_span_lengths: List[int] = []
    batch_modified_span: List[ParagraphSpan] = []
    batch_modified_l: List[int] = []

    for b_i in range(batch_size):
        cur_p: torch.Tensor = input_p[b_i]  # [T, D]
        cur_span: ParagraphSpan = batch_span[b_i]
        # assert p_l[b_i] == cur_span.total_length    # Checking that current length is consistent with span total length
        # remove this because the input paragraph might already be truncated and will not be consistent with the span

        modified_span_list: List[Tuple[int, int]] = []
        m_start, m_end = 0, 0
        m_total_l = 0
        for o_start, o_end in cur_span:
            if o_start >= cur_p.size(0):
                # If the current span is already surpass the boundary, then just ignore this span
                # This line just ignore all the span that surpass the boundary due to paragraph truncate.
                continue

            cur_sent_length = o_end - o_start
            if max_sentence_length is not None:
                cur_sent_length = min(max_sentence_length, cur_sent_length)

            # retrieve the spanned tensor
            start_point = o_start
            end_point = min(p_l[b_i], o_start + cur_sent_length)
            # input span might be greater because truncating paragraph
            retrieved_tensor = cur_p[start_point:end_point]
            cur_sent_length = retrieved_tensor.size(0)  # there might be some cut in the end of the paragraph
            sequence_tensors_list.append(retrieved_tensor)

            # calculate the new start and end of the span
            m_end = m_start + cur_sent_length
            modified_span_list.append((m_start, m_end))
            m_total_l += cur_sent_length
            sequence_span_lengths.append(cur_sent_length)
            m_start = m_end  # reset m_start

            # sequence_tensors_list.append(cur_p[o_start:o_start + cur_sent_length])

        batch_modified_span.append(ParagraphSpan(modified_span_list))
        batch_modified_l.append(m_total_l)

    output_sents = torch_util.pack_list_sequence(sequence_tensors_list, sequence_span_lengths)
    output_l = p_l.data.new_tensor(sequence_span_lengths)
    batched_l = p_l.data.new_tensor(batch_modified_l)
    assert output_sents.size(0) == sum([span.number_of_span() for span in batch_modified_span])

    return output_sents, output_l, batched_l, batch_modified_span


def merge_sentence_to_paragraph(span_inputs: torch.Tensor, span_l: torch.Tensor, batch_span: List[ParagraphSpan],
                                expected_l: torch.Tensor) \
        -> Tuple[torch.Tensor, torch.Tensor]:
    """
    This is the reverse method of the `cut_paragraph_to_sentence` method.
    :param span_inputs: