from langchain.prompts import PromptTemplate


_cot_original_prompt = """You are a mathematician, you are supposed to answer the given question. You need to output the answer in your final sentence like "Therefore, the answer is ...". The answer can only be one of the following forms:
1. a numerical value like 0.1, no symbol and no unit at all.
2. a list of number like [2, 3, 4].
3. True/False.
4. an option like (a), (b), (c), (d)
"""

_cot_prompt_theoremqa = """Can you solve the following math problem? Explain your reasoning. """ + \
"""The final answer can only be one of the following forms:
1. a numerical value like 0.1, no symbol and no unit at all.
2. a list of number like [2, 3, 4].
3. True/False.
4. an option like (a), (b), (c), (d)
Your final answer should be in the form \\boxed{{answer}}, at the end of your response.
Question: {question}
Answer:"""
_cot_prompt_theoremqa_LLAMA = """Can you solve the following math problem? Explain your reasoning. """ + \
"""The final answer can only be one of the following forms:
1. a numerical value like 0.1, no symbol and no unit at all.
2. a list of number like [2, 3, 4].
3. True/False.
4. an option like (a), (b), (c), (d)
at the end of your response.
Question: {question}
Answer: Let's think step by step."""

_cot_simple_prompt_theoremqa = """Can you solve the following math problem? Explain your reasoning. """ + \
"""Your final answer should be a single numerical value, in the form \\boxed{{answer}}, at the end of your response.
Question: {question}
Answer:"""

_mad_start_prompt_theoremqa = """Can you solve the following math problem? Explain your reasoning. """ + \
"""The final answer can only be one of the following forms:
1. a numerical value like 0.1, no symbol and no unit at all.
2. a list of number like [2, 3, 4].
3. True/False.
4. an option like (a), (b), (c), (d)
Your final answer should be in the form \\boxed{{answer}}, at the end of your response.
Question: {question}
Answer:"""

_mad_start_prompt_theoremqa_LLAMA = """Can you solve the following math problem? Explain your reasoning. """ + \
"""The final answer can only be one of the following forms:
1. a numerical value like 0.1, no symbol and no unit at all.
2. a list of number like [2, 3, 4].
3. True/False.
4. an option like (a), (b), (c), (d)
at the end of your response.
Question: {question}
Answer: Let's think step by step."""

_reflexion_reflection_prompt_theoremqa = """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 a question to answer. You were unsuccessful in answering the question either because there is a mistake in your reasoning or there is a phrasing discrepancy with your provided answer and the answer key. In a few sentences, Diagnose a possible reason for failure or phrasing discrepancy and devise a new, concise, high level plan that aims to mitigate the same failure. Use complete sentences.\n""" + \
"""Previous trial:\nQuestion: {question}\nAnswer: {prev_ans}\n\n""" + \
"""Reflection:"""
_cot_reflection_combo_zeroshot_prompt = """Solve the following math problem. Explain your reasonings. """ + \
"""The final answer can only be one of the following forms:
1. a numerical value like 0.1, no symbol and no unit at all.
2. a list of number like [2, 3, 4].
3. True/False.
4. an option like (a), (b), (c), (d)
Your final answer should be in the form \\boxed{{answer}}, at the end of your response.""" + \
"""{reflections}\n""" + \
"""Question: {question}{scratchpad}
Answer:"""
_cot_reflection_combo_zeroshot_prompt_LLAMA = """Solve the following math problem. Explain your reasonings. """ + \
"""The final answer can only be one of the following forms:
1. a numerical value like 0.1, no symbol and no unit at all.
2. a list of number like [2, 3, 4].
3. True/False.
4. an option like (a), (b), (c), (d)
at the end of your response.""" + \
"""{reflections}\n""" + \
"""Question: {question}{scratchpad}
Answer: Let's think step by step."""

_cot_eval_effect2_prompt = """Question: {question}\n\n""" + \
"""Answer: {answer}\n\n""" + \
"""Does the answer correctly answer the question?"""
_cot_eval_effect2_sure_prompt = """Question: {question}\n\n""" + \
"""Answer: {answer}\n\n""" + \
"""Given the question and the answer, give a judgement (sure/likely/impossible) if the answer is correct."""

COT_THEOREMQA = PromptTemplate(
                        input_variables=["question"],
                        template = _cot_prompt_theoremqa,
                        )
COT_THEOREMQA_LLAMA = PromptTemplate(
                        input_variables=["question"],
                        template = _cot_prompt_theoremqa_LLAMA,
                        )
COT_SIMPLE_THEOREMQA = PromptTemplate(
                        input_variables=["question"],
                        template = _cot_simple_prompt_theoremqa,
                        )
MAD_START_THEOREMQA = PromptTemplate(
                        input_variables=["question"],
                        template = _mad_start_prompt_theoremqa)
MAD_START_THEOREMQA_LLAMA = PromptTemplate(
                        input_variables=["question"],
                        template = _mad_start_prompt_theoremqa_LLAMA)
REFLEXION_COT_REFLECT_THEOREMQA = PromptTemplate(
                        input_variables=["question", "prev_ans"],
                        template = _reflexion_reflection_prompt_theoremqa)
REFLEXION_COT_COMBO_PROPOSE = PromptTemplate(
                        input_variables=["reflections", "question", "scratchpad"],
                        template = _cot_reflection_combo_zeroshot_prompt)
REFLEXION_COT_COMBO_PROPOSE_LLAMA = PromptTemplate(
                        input_variables=["reflections", "question", "scratchpad"],
                        template = _cot_reflection_combo_zeroshot_prompt_LLAMA)
RECUR_PROMPT_THEOREMQA = PromptTemplate(
                        input_variables=["question","answer"],
                        template = _cot_eval_effect2_prompt)
RECUR_SURE_PROMPT_THEOREMQA = PromptTemplate(
                        input_variables=["question","answer"],
                        template = _cot_eval_effect2_sure_prompt)

THEOREMQA_PROMPT_TEMPLATES = {"THEOREMQA_COT_PROMPT": COT_THEOREMQA,
                              "THEOREMQA_COT_PROMPT_LLAMA": COT_THEOREMQA_LLAMA,
                              "THEOREMQA_COT_SIMPLE_PROMPT": COT_SIMPLE_THEOREMQA,
                              "THEOREMQA_ORIGINAL_COT_PROMPT": _cot_original_prompt,
                              "THEOREMQA_MAD_PROMPT": MAD_START_THEOREMQA,
                              "THEOREMQA_MAD_PROMPT_LLAMA": MAD_START_THEOREMQA_LLAMA,
                              "THEOREMQA_REFLEXION_PROMPT": [COT_THEOREMQA, REFLEXION_COT_REFLECT_THEOREMQA, REFLEXION_COT_COMBO_PROPOSE],
                                "THEOREMQA_REFLEXION_PROMPT_LLAMA": [COT_THEOREMQA_LLAMA, REFLEXION_COT_REFLECT_THEOREMQA, REFLEXION_COT_COMBO_PROPOSE_LLAMA],
                              "THEOREMQA_RECURSIVE_COT_PROMPT": [COT_THEOREMQA, RECUR_PROMPT_THEOREMQA],
                              "THEOREMQA_RECURSIVE_NUMERICAL_COT_PROMPT": [COT_SIMPLE_THEOREMQA, RECUR_PROMPT_THEOREMQA],
                              "THEOREMQA_RECURSIVE_NUMERICAL_SURE_COT_PROMPT": [COT_SIMPLE_THEOREMQA, RECUR_SURE_PROMPT_THEOREMQA],
                              }