import re
import constant
def add_prompt(text, se_order="aos"):
    prompt = []
    for counter, se in enumerate(se_order):
        prompt.append(constant.SENTIMENT_ELEMENT[se] + " : " + f"<extra_id_{counter}>")
    prompt = " ,".join(prompt)
    result = text + "| " + prompt
    return result

def catch_answer(out, se_order):
        if out == constant.NO_TARGET:
            return []
        pattern = r""
        for se in se_order:
            if se != 's':
                pattern += f"<extra_id_\d+>\s*(?P<{constant.SENTIMENT_ELEMENT[se]}>[^;]+)\s*"
            else:
                pattern += f"<extra_id_\d+>\s*(?P<{constant.SENTIMENT_ELEMENT['s']}>positive|negative|neutral)\s*"
        result = [found_iter.groupdict() for found_iter in re.finditer(pattern,out)]
        for i in range(len(result)):
            for k, v in result[i].items():
                result[i][k] = v.strip()
        return result