from langchain.prompts import PromptTemplate

# 5-shot
standard_prompt_template = '''Use numbers and basic arithmetic operations (+ - * /) to obtain 24.
Input: 4 4 6 8
Answer: (4 + 8) * (6 - 4) = 24
Input: 2 9 10 12
Answer: 2 * 12 * (10 - 9) = 24
Input: 4 9 10 13
Answer: (13 - 9) * (10 - 4) = 24
Input: 1 4 8 8
Answer: (8 / 4 + 1) * 8 = 24
Input: 5 5 5 9
Answer: 5 + 5 + 5 + 9 = 24
Input: {input}
'''

# 5-shot
cot_prompt_template = '''Use numbers and basic arithmetic operations (+ - * /) to obtain 24. Each step, you are only allowed to choose two of the remaining numbers to obtain a new number.

###Examples###
Input: 4 4 6 8
Steps:
4 + 8 = 12 (left: 4 6 12)
6 - 4 = 2 (left: 2 12)
2 * 12 = 24 (left: 24)
Answer: (6 - 4) * (4 + 8) = 24
Input: 2 9 10 12
Steps:
12 * 2 = 24 (left: 9 10 24)
10 - 9 = 1 (left: 1 24)
24 * 1 = 24 (left: 24)
Answer: (12 * 2) * (10 - 9) = 24
Input: 4 9 10 13
Steps:
13 - 10 = 3 (left: 3 4 9)
9 - 3 = 6 (left: 4 6)
4 * 6 = 24 (left: 24)
Answer: 4 * (9 - (13 - 10)) = 24
Input: 1 4 8 8
Steps:
8 / 4 = 2 (left: 1 2 8)
1 + 2 = 3 (left: 3 8)
3 * 8 = 24 (left: 24)
Answer: (1 + 8 / 4) * 8 = 24
Input: 5 5 5 9
Steps:
5 + 5 = 10 (left: 5 9 10)
10 + 5 = 15 (left: 9 15)
15 + 9 = 24 (left: 24)
Answer: ((5 + 5) + 5) + 9 = 24
End Examples###

Input: {input}
'''

propose_prompt_template = '''Input: 2 8 8 14
Possible next steps:
2 + 8 = 10 (left: 8 10 14)
8 / 2 = 4 (left: 4 8 14)
14 + 2 = 16 (left: 8 8 16)
2 * 8 = 16 (left: 8 14 16)
8 - 2 = 6 (left: 6 8 14)
14 - 8 = 6 (left: 2 6 8)
14 /  2 = 7 (left: 7 8 8)
14 - 2 = 12 (left: 8 8 12)
Input: {input}
Possible next steps:
'''

# 3-shot
propose_prompt_template_123 = '''Propose possible next steps given input.

###Examples###
Input: 2 8 8 14
Possible next steps:
2 + 8 = 10 (left: 8 10 14)
8 / 2 = 4 (left: 4 8 14)
14 + 2 = 16 (left: 8 8 16)
2 * 8 = 16 (left: 8 14 16)
8 - 2 = 6 (left: 6 8 14)
14 - 8 = 6 (left: 2 6 8)
14 /  2 = 7 (left: 7 8 8)
14 - 2 = 12 (left: 8 8 12)
Input: 6 7 12
Possible next steps:
6 + 7 = 13 (left: 13 12)
12 - 6 = 6 (left: 6 7)
12 + 7 = 19 (left: 19 6)
6 * 7 = 42 (left: 42 12)
7 - 6 = 1 (left: 1 12)
12 - 7 = 5 (left: 5 6)
12 /  6 = 2 (left: 2 7)
12 + 6 = 18 (left: 7 18)
Input: 12 2
Possible next steps:
12 + 2 = 14 (left: 14)
12 - 2 = 10 (left: 10)
12 * 2 = 24 (left: 24)
12 / 2 = 6 (left: 6)
###End Examples###

Input: {input}
Possible next steps:
'''

propose_prompt_specific_template = '''Given a set of numbers, propose possible next steps to obtain 24 using numbers and basic arithmetic operations (+ - * /).
Make sure the next steps are using numbers in the inputs. Each number in the input should only be used once. The next steps should be correct and diverse. In the parentheses, write out the number(s) left after the operation. The number(s) left should contain the result of the operation and the number(s) unused in the input. The number(s) left should always contain 1 less number that the input. Stricly follows the format in the examples without any extra words.

###Examples###
Input: 2 8 8 14
Possible next steps:
2 + 8 = 10 (left: 8 10 14)
8 / 2 = 4 (left: 4 8 14)
14 + 2 = 16 (left: 8 8 16)
2 * 8 = 16 (left: 8 14 16)
8 - 2 = 6 (left: 6 8 14)
14 - 8 = 6 (left: 2 6 8)
14 /  2 = 7 (left: 7 8 8)
14 - 2 = 12 (left: 8 8 12)
### End Examples###

Input: {input}
Possible next steps:
'''
propose_prompt_specific_123_template = '''Given a set of numbers, propose possible next steps to obtain 24 using numbers and basic arithmetic operations (+ - * /).
Make sure the next steps are using numbers in the inputs. Each number in the input should only be used once. The next steps should be correct and diverse. In the parentheses, write out the number(s) left after the operation. The number(s) left should contain the result of the operation and the number(s) unused in the input. The number(s) left should always contain 1 less number that the input. Stricly follows the format in the examples without any extra words.

###Examples###
Input: 2 8 8 14
Possible next steps:
2 + 8 = 10 (left: 8 10 14)
8 / 2 = 4 (left: 4 8 14)
14 + 2 = 16 (left: 8 8 16)
2 * 8 = 16 (left: 8 14 16)
8 - 2 = 6 (left: 6 8 14)
14 - 8 = 6 (left: 2 6 8)
14 /  2 = 7 (left: 7 8 8)
14 - 2 = 12 (left: 8 8 12)

Input: 6 7 12
Possible next steps:
6 + 7 = 13 (left: 13 12)
12 - 6 = 6 (left: 6 7)
12 + 7 = 19 (left: 19 6)
6 * 7 = 42 (left: 42 12)
7 - 6 = 1 (left: 1 12)
12 - 7 = 5 (left: 5 6)
12 /  6 = 2 (left: 2 7)
12 + 6 = 18 (left: 7 18)

Input: 12 2
Possible next steps:
12 + 2 = 14 (left: 14)
12 - 2 = 10 (left: 10)
12 * 2 = 24 (left: 24)
12 / 2 = 6 (left: 6)
### End Examples###

Input: {input}
Possible next steps:
'''



value_prompt_template = '''Evaluate if given numbers can reach 24 (sure/likely/impossible)

###Examples###
Input: 10 14
Thoughts:
10 + 14 = 24
Judge: sure
Input: 11 12
Thoughts:
11 + 12 = 23
12 - 11 = 1
11 * 12 = 132
11 / 12 = 0.91
Judge: impossible
Input: 4 4 10
Thoughts:
4 + 4 + 10 = 8 + 10 = 18
4 * 10 - 4 = 40 - 4 = 36
(10 - 4) * 4 = 6 * 4 = 24
Judge: sure
Input: 4 9 11
Thoughts:
9 + 11 + 4 = 20 + 4 = 24
Judge: sure
Input: 5 7 8
Thoughts:
5 + 7 + 8 = 12 + 8 = 20
(8 - 5) * 7 = 3 * 7 = 21
I cannot obtain 24 now, but numbers are within a reasonable range
Judge: likely
Input: 5 6 6
Thoughts:
5 + 6 + 6 = 17
(6 - 5) * 6 = 1 * 6 = 6
I cannot obtain 24 now, but numbers are within a reasonable range
Judge: likely
Input: 10 10 11
Thoughts:
10 + 10 + 11 = 31
(11 - 10) * 10 = 10
10 10 10 are all too big
Judge: impossible
Input: 1 3 3
Thoughts:
1 * 3 * 3 = 9
(1 + 3) * 3 = 12
1 3 3 are all too small
Judge: impossible
###End Examples###

Input: {input}
Thoughts:'''
value_prompt_gpt4_template = '''Evaluate if given numbers can reach 24 (sure/likely/impossible)
10 14
10 + 14 = 24
sure
11 12
11 + 12 = 23
12 - 11 = 1
11 * 12 = 132
11 / 12 = 0.91
impossible
4 4 10
4 + 4 + 10 = 8 + 10 = 18
4 * 10 - 4 = 40 - 4 = 36
(10 - 4) * 4 = 6 * 4 = 24
sure
4 9 11
9 + 11 + 4 = 20 + 4 = 24
sure
5 7 8
5 + 7 + 8 = 12 + 8 = 20
(8 - 5) * 7 = 3 * 7 = 21
I cannot obtain 24 now, but numbers are within a reasonable range
likely
5 6 6
5 + 6 + 6 = 17
(6 - 5) * 6 = 1 * 6 = 6
I cannot obtain 24 now, but numbers are within a reasonable range
likely
10 10 11
10 + 10 + 11 = 31
(11 - 10) * 10 = 10
10 10 10 are all too big
impossible
1 3 3
1 * 3 * 3 = 9
(1 + 3) * 3 = 12
1 3 3 are all too small
impossible
{input}
'''

value_last_step_prompt_template = '''Use numbers and basic arithmetic operations (+ - * /) to obtain 24. Given an input and an answer, give a judgement (sure/impossible) if the answer is correct, i.e. it uses each input exactly once and no other numbers, and reach 24.

###Examples###
Input: 4 4 6 8
Answer: (4 + 8) * (6 - 4) = 24
Judge: 
sure
Input: 2 9 10 12
Answer: 2 * 12 * (10 - 9) = 24
Judge: 
sure
Input: 4 9 10 13
Answer: (13 - 9) * (10 - 4) = 24
Judge: 
sure
Input: 4 4 6 8
Answer: (4 + 8) * (6 - 4) + 1 = 25
Judge: 
impossible
Input: 2 9 10 12
Answer: 2 * (12 - 10) = 24
Judge: 
impossible
Input: 4 9 10 13
Answer: (13 - 4) * (10 - 9) = 24
Judge: 
impossible
##End Examples###

Input: {input}
Answer: {answer}
Judge:'''


# reflexion
REFLECTION_HEADER = 'You have attempted to answer following question before and failed. The following reflection(s) give a plan to avoid failing to answer the question in the same way you did previously. Use them to improve your strategy of correctly answering the given question.\n'
REFLECTION_AFTER_LAST_TRIAL_HEADER = 'The following reflection(s) give a plan to avoid failing to answer the question in the same way you did previously. Use them to improve your strategy of correctly answering the given question.\n'
LAST_TRIAL_HEADER = 'You have attempted to answer the following question before and failed. Below is the last trial you attempted to answer the question.\n'

# possible errors:
# 1. final answer has incorrect steps to 24: "10 - 4 = 6 (left: 5 6 6)\n6 * 6 = 36 (left: 5 36)\n36 - 5 = 24 (left: 24)\nAnswer: ((10 - 6) * 4) - 5 = 24"

# 2. final answer has wrong number of numbers: "10 - 6 = 4 (left: 4 4 5)\n4 * 5 = 20 (left: 4 20)\n20 + 4 = 24 (left: 24)\nAnswer: ((10 - 6) * 5) + 4 + 4 = 24"
# 3. reasoning right, but final combo wrong "10 - 6 = 4 (left: 4 4 5)\n4 * 5 = 20 (left: 4 20)\n20 + 4 = 24 (left: 24)\nAnswer: ((10 - 6) * 5) + 4 = 24"
# 4. reasoning wrong "7 - 4 = 3 (left: 1 2 3)\n3 * 2 = 6 (left: 1 6)\n6 * 4 = 24 (left: 24)\nAnswer: (1 + (7 - 4)) * 2 = 24\n",

_cot_reflection_reflect_prompt = """You are an advanced reasoning agent that can improve based on self refection. You will be given a previous reasoning trial in which you were given numbers and basic arithmetic operations to obtain 24. You were unsuccessful in finding the final answer either because you reasoning steps don't lead to 24 or your final answer fails to combine the reasoning steps correctly to reache 24. In a few sentences, Diagnose a possible reason for failure and devise a new, concise, high level plan that aims to mitigate the same failure. Use complete sentences.
Here are some examples:
Input: 4 4 6 8
Steps:
4 + 8 = 12 (left: 4 6 12)
6 - 4 = 2 (left: 2 12)
2 * 12 = 24 (left: 24)
Answer: (6 - 4) * (4 + 8) + 4 = 28

Reflection: I failed to obtain 24 because I used 5 numbers in the final answer instead of 4. I used the number 4 twice. If I follow the steps correctly can obtain 24 by (6 - 4) * (4 + 8) = 24. In the future, when attempting this question, I should make sure to use each number only once and use 4 numbers in the final answer. I should also make sure to follow the reasoning steps correctly.

Input: 2 9 10 12
Steps:
12 * 2 = 24 (left: 9 10 24)
24 - 10 = 1 (left: 1 14)
14 * 1 = 24 (left: 24)
Answer: ((12 * 2) - 10) * 1 = 24

Reflection: My reasoning steps are wrong and I failed to obtain 24 because 14 * 1 = 14 not 24. In the future, when attempting this question, I should pick steps that lead to 24. I should also make sure to do arithmetic operations correctly.

Input: 4 9 10 13
Steps:
13 - 10 = 3 (left: 3 4 9)
9 - 3 = 6 (left: 4 6)
4 * 6 = 24 (left: 24)
Answer: 9 * (4 - (13 - 10)) = 9

Reflection: My reasoning steps are correct, but I failed to obtain 24 because in the final answer I swapped 4 and 9. It should be 4 * (9 - (13 - 10)) = 24 instead of 9 * (4 - (13 - 10)) = 9. In the future, when attempting this question, I should pay more attention to the order of the reasoning steps and apply them in the final answer correctly. I need to be careful to not swap numbers in the final answer.
(END OF EXAMPLES)
Previous trial:\nInput: {input}\nSteps:\n{prev_ans}\n\n""" + \
"""Reflection:"""
_cot_reflection_reflect_zeroshot_prompt = """You are an advanced reasoning agent that can improve based on self refection. You will be given a previous reasoning trial in which you were given numbers and basic arithmetic operations to obtain 24. You were unsuccessful in finding the final answer either because you reasoning steps don't lead to 24 or your final answer fails to combine the reasoning steps correctly to reache 24. In a few sentences, Diagnose a possible reason for failure and devise a new, concise, high level plan that aims to mitigate the same failure. Use complete sentences.\n""" + \
"""Previous trial:\nInput: {input}\nSteps:\n{prev_ans}\n\n""" + \
"""Reflection:"""
_cot_reflection_combo_zeroshot_prompt = '''Use numbers and basic arithmetic operations (+ - * /) to obtain 24. Each step, you are only allowed to choose two of the remaining numbers to obtain a new number.

###Examples###
Input: 4 4 6 8
Steps:
4 + 8 = 12 (left: 4 6 12)
6 - 4 = 2 (left: 2 12)
2 * 12 = 24 (left: 24)
Answer: (6 - 4) * (4 + 8) = 24
Input: 2 9 10 12
Steps:
12 * 2 = 24 (left: 9 10 24)
10 - 9 = 1 (left: 1 24)
24 * 1 = 24 (left: 24)
Answer: (12 * 2) * (10 - 9) = 24
Input: 4 9 10 13
Steps:
13 - 10 = 3 (left: 3 4 9)
9 - 3 = 6 (left: 4 6)
4 * 6 = 24 (left: 24)
Answer: 4 * (9 - (13 - 10)) = 24
Input: 1 4 8 8
Steps:
8 / 4 = 2 (left: 1 2 8)
1 + 2 = 3 (left: 3 8)
3 * 8 = 24 (left: 24)
Answer: (1 + 8 / 4) * 8 = 24
Input: 5 5 5 9
Steps:
5 + 5 = 10 (left: 5 9 10)
10 + 5 = 15 (left: 9 15)
15 + 9 = 24 (left: 24)
Answer: ((5 + 5) + 5) + 9 = 24
End Examples###

{reflections}
Input: {input}
'''
REFLEXION_COT_PROPOSE = PromptTemplate(
                        input_variables=["input"],
                        template = cot_prompt_template)
REFLEXION_COT_REFLECT = PromptTemplate(
                        input_variables=["input", "prev_ans"],
                        template = _cot_reflection_reflect_zeroshot_prompt)
REFLEXION_COT_COMBO_PROPOSE = PromptTemplate(
                        input_variables=["reflections", "input"],
                        template = _cot_reflection_combo_zeroshot_prompt)
REFLEXION_COT_FEWSHOT_REFLECT = PromptTemplate(
                        input_variables=["input", "prev_ans"],
                        template = _cot_reflection_reflect_prompt)


REFLEXION_PROMPT_TEMPLATES = {"REFLEXION": [REFLEXION_COT_PROPOSE, REFLEXION_COT_REFLECT,REFLEXION_COT_COMBO_PROPOSE],
                    "REFLEXION_FEWSHOT": [REFLEXION_COT_PROPOSE, REFLEXION_COT_FEWSHOT_REFLECT,REFLEXION_COT_COMBO_PROPOSE],
}





standard_prompt = PromptTemplate(
                        input_variables=["input"],
                        template = standard_prompt_template,
                        )

cot_prompt = PromptTemplate(
                        input_variables=["input"],
                        template = cot_prompt_template,
                        )


# propose_prompt_template, propose_prompt_template_123
# propose_prompt_specific_template, propose_prompt_specific_123_template
propose_prompt = PromptTemplate(
                        input_variables=["input"],
                        template = propose_prompt_template,
                        )
propose_prompt_123 = PromptTemplate(
                        input_variables=["input"],
                        template = propose_prompt_template_123,
                        )
propose_prompt_123_instruct = PromptTemplate(
                        input_variables=["input"],
                        template = propose_prompt_specific_123_template,
                        )



value_prompt = PromptTemplate(
                        input_variables=["input"],
                        template = value_prompt_template,
                        )

value_gpt4_prompt = PromptTemplate(
                        input_variables=["input"],
                        template = value_prompt_gpt4_template,
                        )
value_last_step_prompt = PromptTemplate(
                        input_variables=["input","answer"],
                        template = value_last_step_prompt_template,
                        )

GAME24_TEMPLATES = {"ToT_ORIGINAL": [propose_prompt,cot_prompt, value_prompt,value_last_step_prompt],
                    "ToT_3shots": [propose_prompt_123,cot_prompt, value_prompt,value_last_step_prompt],
                    "ToT_3shots_gpt4value": [propose_prompt_123,cot_prompt, value_gpt4_prompt,value_last_step_prompt],
                    "ToT_specific3shots": [propose_prompt_123_instruct,cot_prompt, value_prompt,value_last_step_prompt],
                    "CoT_ORIGINAL": ["",cot_prompt, "",value_last_step_prompt],   
                    }