from string import Template

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

PROMPT_TEMPLATES = {
    "value-based": {
        "training": "belief states: value = $value, slot = $slot",
        "generate": "belief states: value = $value, slot ="
    },
    "inverse-prompt": {
        "training": "belief states: slot = $slot, value = $value",
        "generate": "belief states: slot = $slot, value ="
    }
}


def get_prompt_for_training(typ, slot_value):
    template = Template(PROMPT_TEMPLATES[typ]["training"])
    return template.substitute(slot=slot_value[0], value=slot_value[1])


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


def get_inverse_prompt(slot):
    template = Template(PROMPT_TEMPLATES["inverse-prompt"]["generate"])
    return template.substitute(slot=slot)