from string import Template

TYPE_VALUE_BASED_PROMPT = "value-based"
TYPE_INVERSE_PROMPT = "inverse-prompt"
TYPE_PROMPT_ENSEMBLE = "prompt-ensemble"

PROMPT_WEIGHT = 0.25

PROMPT_TEMPLATES = {
    "value-based": {
        "training": "belief states: value = $value, slot = $slot",
        "generate": "belief states: value = $value, slot ="
    },
    "inverse-prompt": "belief states: slot = $slot, value = $value",
    "prompt-ensemble": {
        "training": {
            "p1": "belief states: value = $value, slot = $slot",
            "p2": "belief states: $value = $slot",
            "p3": "$value is of slot type $slot",
            "p4": "$value is the value of $slot"
        },
        "generate": {
            "p1": "belief states: value = $value, slot =",
            "p2": "belief states: $value =",
            "p3": "$value is of slot type",
            "p4": "$value is the value of"
        }
    }
}


def get_prompt_for_training(typ, slot_value):
    if typ == TYPE_INVERSE_PROMPT:
        template = Template(PROMPT_TEMPLATES[typ])
        return template.substitute(slot=slot_value[0], value=slot_value[1])
    else:
        template = PROMPT_TEMPLATES[typ]['training']
        if isinstance(template, str):
            return Template(template).substitute(slot=slot_value[0], value=slot_value[1])
        elif isinstance(template, dict):
            template_list = template.values()
            prompt_list = []
            for template_str in template_list:
                template = Template(template_str)
                prompt_list.append(template.substitute(slot=slot_value[0], value=slot_value[1]))
            return prompt_list


def get_value_based_prompt(value):
    template = Template(PROMPT_TEMPLATES["value-based"]["generate"])
    return template.substitute(value=value)