import re


def unsegment_labels(labels, mt, segmentation_suffix=u'@@', heuristic='any_bad'):
    """
    Reduce label sequence according to MT segmentation, apply heuristic to decide how labels should be reduced

    Params:
      labels: sequence of labels
      mt: sequence of tokens with segmentation applied
      segmentation_suffix: suffix which indicates when a token has been segmented
      heuristic: key which tells us how to reduce segmented labels

    Returns:
      reduced labels corresponding to the original MT hypothesis without segmentation
    """

    unsegmented_labels = []
    unsegmented_mt = []
    for label_seq, token_words in zip(labels, mt):
        mapped_labels = []
        mapped_mt = []

        current_labels = []
        current_word = []
        for label, word in zip(label_seq, token_words):
            if word.endswith(segmentation_suffix):
                current_labels.append(label)
                current_word.append(word)
            else:
                if len(current_labels) > 0:
                    current_labels.append(label)
                    current_word.append(word)
                    # word is finished, now use heuristic to decide which label wins
                    if heuristic == 'any_bad':
                        if any(l == u'BAD' for l in current_labels):
                            mapped_labels.append(u'BAD')
                        else:
                            mapped_labels.append(u'OK')
                    else:
                        raise ValueError('Unknown heuristic: {}'.format(heuristic))

                    mapped_mt.append(u''.join([re.sub(segmentation_suffix, '', u) for u in current_word]))
                    current_labels = []
                    current_word = []
                else:
                    mapped_labels.append(label)
                    mapped_mt.append(word)

        unsegmented_labels.append(mapped_labels)
        unsegmented_mt.append(mapped_mt)

    return unsegmented_labels, unsegmented_mt