from string import Template

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

PROMPT_WEIGHT = 0.25

INVERSE_PROMPTS = {
    "i1": "belief states: $slot = $value",
    "i2": "belief states: slot = $slot, value = $value",
}

PROMPT_TEMPLATES = {
    "value-based": {
        "training": "belief states: value = $value, slot = $slot",
        "generate": "belief states: value = $value, slot ="
    },
    "inverse-prompt": {
        "training": INVERSE_PROMPTS["i2"],
    },
    "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"
        }
    }
}

# TODO add the final set of answered prompts
ANSWERED_PROMPTS = """
user : hello, i'm looking for a 4 star place on the west side to stay at.\n
system : do you have a price range ?\n
user : yes, i am looking in the expensive price range. also, i need free parking.\n
belief states: west = area, 4 = stars, expensive = price, yes = parking \n
user : i am looking for a train, it should go to cambridge and should depart from norwich.\n
system : what time and day are you looking to travel?\n
user : yes, i would like travel on monday and i would need to arrive by 08:30. \n
belief states: cambridge = destination, norwich = departure, monday = day, 08:30 = arrive \n
user : i would like an expensive place to dine, centre of town.\n
system : what type of food would you like to eat?\n
user : the type of food doesn't matter, but i need a reservation for 8 people on wednesday at 12:45\n
belief states: expensive = price, centre = area, dont care = food, 8 = people, wednesday = day, 12:45 = time \n
user : i need a taxi to pick me up at curry prince at 08:15.\n
belief states: curry prince = departure, 08:15 = leave \n
user : i am looking for a place in the centre of town that is a nightclub.\n
belief states: night club = type, centre = area \n
"""


def get_prompt_for_training(typ, slot_value):
    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:
            prompt = Template(template_str).substitute(slot=slot_value[0], value=slot_value[1])
            prompt_list.append(prompt)
        return prompt_list


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


def get_ensemble_prompts(value):
    templates = PROMPT_TEMPLATES[TYPE_PROMPT_ENSEMBLE]['generate']
    template_list = templates.values()
    prompt_list = []
    for template_str in template_list:
        template = Template(template_str)
        prompt_list.append(template.substitute(value=value))
    return prompt_list


def get_answered_prompts():
    return ANSWERED_PROMPTS