from langchain.prompts import PromptTemplate
from .hotpotQA_fewshot import *

####### Standard Prompt & CoT Prompt #######
cot_prompt = """Solve a question answering task by having a Thought, then Finish with your answer. Thought can reason about the current situation. Finish[answer] returns the answer and finishes the task.
{context}
Question: {question}{scratchpad}"""
cot_context_prompt = """Solve a question answering task by having a Thought, then Finish with your answer. Thought can reason about the current situation. Finish[answer] returns the answer and finishes the task.
Relavant Context: {context}
Question: {question}{scratchpad}"""

def cot_fewshot_prompt(fewshot):
    cot_fewshot_prompt = """Solve a question answering task by having a Thought, then Finish with your answer. Thought can reason about the current situation. Finish[answer] returns the answer and finishes the task.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)
{context}
Question: {question}{scratchpad}"""
    return cot_fewshot_prompt
def cot_fewshot_context_prompt(fewshot):
    cot_fewshot_context_prompt = """Solve a question answering task by having a Thought, then Finish with your answer. Thought can reason about the current situation. Finish[answer] returns the answer and finishes the task. You will be given context that you should use to help you answer the question.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

Relevant Context: {context}
Question: {question}{scratchpad}"""
    return cot_fewshot_context_prompt

# No ReAct
cot_noreact_zeroshot_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be a phrase, in the form [answer], at the end of your response.
{context}
Question: {question}{scratchpad}
Answer:"""
def cot_noreact_fewshot_prompt(fewshot):
    cot_noreact_fewshot_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be a phrase, in the form [answer], at the end of your response.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)
{context}
Question: {question}{scratchpad}
Answer:"""
    return cot_noreact_fewshot_prompt
def cot_noreact_fewshot_context_prompt(fewshot):
    cot_noreact_fewshot_context_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be a phrase, in the form [answer], at the end of your response.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

Relevant Context: {context}
Question: {question}{scratchpad}
Answer:"""
    return cot_noreact_fewshot_context_prompt


### io
io_prompt = """Solve a question answering task by Finish with your answer. Finish[answer] returns the answer and finishes the task.
{context}
Question: {question}{scratchpad}"""
def io_fewshot_prompt(fewshot):
    io_fewshot_prompt = """Solve a question answering task by Finish with your answer. Finish[answer] returns the answer and finishes the task.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)
{context} 
Question: {question}{scratchpad}"""
    return io_fewshot_prompt 
def io_fewshot_context_prompt(fewshot):
    io_fewshot_context_prompt = """Solve a question answering task by Finish with your answer. Finish[answer] returns the answer and finishes the task. You will be given context that you should use to help you answer the question.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

Relevant Context: {context} 
Question: {question}{scratchpad}"""
    return io_fewshot_context_prompt

def io_noreact_fewshot_prompt(fewshot):
    io_noreact_fewshot_prompt = """Solve a question answering task. You final answer should be a single name phrase phrase, in the form [answer]. Don't output anything else.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)
{context} 
Question: {question}{scratchpad}
Asnwer:"""
    return io_noreact_fewshot_prompt 
def io_noreact_fewshot_context_prompt(fewshot):
    io_noreact_fewshot_context_prompt = """Solve a question answering task. You final answer should be a single name phrase phrase, in the form [answer]. Don't output anything else. You will be given context that you should use to help you answer the question.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

Relevant Context: {context} 
Question: {question}
Answer:"""##
    return io_noreact_fewshot_context_prompt


### Answer then CoT
# def answerreason_noreact_fewshot_prompt(fewshot):
#     answerreason_noreact_fewshot_prompt = """Solve a question answering task. You final answer should be a single name phrase phrase, in the form [answer]. Don't output anything else.
# Here are some examples:""" + \
# f"\n{fewshot}\n" + \
# """(END OF EXAMPLES)
# {context} 
# Question: {question}{scratchpad}
# Asnwer:"""
#     return answerreason_noreact_fewshot_prompt 
# def answerreason_noreact_fewshot_context_prompt(fewshot):
#     answerreason_noreact_fewshot_context_prompt = """Solve a question answering task. You final answer should be a single name phrase phrase, in the form [answer]. Don't output anything else. You will be given context that you should use to help you answer the question.
# Here are some examples:""" + \
# f"\n{fewshot}\n" + \
# """(END OF EXAMPLES)

# Relevant Context: {context} 
# Question: {question}
# Answer:"""##
#     return answerreason_noreact_fewshot_context_prompt


def make_template(prompt):
    return PromptTemplate(
                        input_variables=["context", "question", "scratchpad"],
                        template = prompt,
                        )
# print(cot_fewshot_context_prompt("hello"))
# print(cot_prompt.format(context="", question="world", scratchpad=""))

COT_ZEROSHOT_PROMPT = PromptTemplate(
                        input_variables=["context", "question", "scratchpad"],
                        template = cot_prompt,
                        )
COT_ZEROSHOT_CONTEXT_PROMPT = PromptTemplate(
                        input_variables=["context", "question", "scratchpad"],
                        template = cot_prompt,
                        )
IO_ZEROSHOT_PROMPT = PromptTemplate(
                        input_variables=["context", "question", "scratchpad"],
                        template = io_prompt,
                        )



COT_PROMPT_TEMPLATES = {"COT_ZEROSHOT_PROMPT": COT_ZEROSHOT_PROMPT,
                        "COT_ZEROSHOT_CONTEXT_PROMPT": COT_ZEROSHOT_CONTEXT_PROMPT,
                        "COT_FEWSHOT_PROMPT": make_template(cot_fewshot_prompt(cot_fewshot_SIMPLE6)),
                        "COT_FEWSHOT_CONTEXT_PROMPT": make_template(cot_fewshot_prompt(cot_fewshot_context)),
                        "IO_ZEROSHOT_PROMPT": IO_ZEROSHOT_PROMPT,
                        "IO_ZEROSHOT_CONTEXT_PROMPT": IO_ZEROSHOT_PROMPT,

                    
                        "COT_NOREACT_ZEROSHOT_PROMPT": make_template(cot_noreact_zeroshot_prompt),
                        "COT_NOREACT_FEWSHOT_PROMPT": make_template(cot_noreact_fewshot_prompt(cot_noreact_fewshot_SIMPLE6)),
                        "COT_NOREACT_FEWSHOT_CONTEXT_PROMPT": make_template(cot_noreact_fewshot_context_prompt(cot_noreact_fewshot_context)),
                        }
############################################

####### CoT Reflexion Prompt #######
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'
def cot_reason_fewshot_prompt(fewshot):
    cot_fewshot_prompt = """Solve a question answering task by having a Thought, then Finish with your answer. Thought can reason about the current situation. Finish[answer] returns the answer and finishes the task.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)
{reflections}
{context}
Question: {question}{scratchpad}"""
    return cot_fewshot_prompt
def cot_noreact_reason_fewshot_prompt(fewshot):
    cot_noreact_reason_fewshot_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be a phrase, in the form [answer], at the end of your response.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)
{reflections}
{context}
Question: {question}{scratchpad}
Answer:"""
    return cot_noreact_reason_fewshot_prompt
def cot_reason_fewshot_context_prompt(fewshot):
    cot_fewshot_context_prompt = """Solve a question answering task by having a Thought, then Finish with your answer. Thought can reason about the current situation. Finish[answer] returns the answer and finishes the task. You will be given context that you should use to help you answer the question.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

{reflections}

Relevant Context: {context}
Question: {question}{scratchpad}"""
    return cot_fewshot_context_prompt
def cot_noreact_reason_fewshot_context_prompt(fewshot):
    cot_noreact_reason_fewshot_context_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be a phrase, in the form [answer], at the end of your response.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

{reflections}

Relevant Context: {context}
Question: {question}{scratchpad}
Answer:"""
    return cot_noreact_reason_fewshot_context_prompt







def cot_reflect_fewshot_prompt(fewshot):
    cot_reflect_fewshot_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 a question to answer. You were unsuccessful in answering the question either because you guessed the wrong answer with Finish[<answer>] 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.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)
{context}
Previous trial:
Question: {question}{scratchpad}

Reflection:"""
    return cot_reflect_fewshot_prompt
def cot_reflect_nohard_fewshot_prompt(fewshot):
    cot_reflect_fewshot_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 a question to answer. Evaluate you answer to the question. If you think it is successful, state it and repeat your ansewr. If you think you were unsuccessful in answering the question either because you guessed the wrong answer with Finish[<answer>] 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.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)
{context}
Previous trial:
Question: {question}{scratchpad}

Reflection:"""
    return cot_reflect_fewshot_prompt
def cot_noreact_reflect_fewshot_prompt(fewshot):
    cot_noreact_reflect_fewshot_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 a question to answer. You were unsuccessful in answering the question either because you guessed the wrong answer with [<answer>] 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.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)
{context}
Previous trial:
Question: {question}{scratchpad}

Reflection:"""
    return cot_noreact_reflect_fewshot_prompt
def cot_reflect_fewshot_context_prompt(fewshot):
    cot_reflect_fewshot_context_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 access to relevant context and a question to answer. You were unsuccessful in answering the question either because you guessed the wrong answer with Finish[<answer>] 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.  
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

Previous trial:
Relevant Context: {context}
Question: {question}{scratchpad}

Reflection:"""
    return cot_reflect_fewshot_context_prompt
def cot_noreact_reflect_fewshot_context_prompt(fewshot):
    cot_noreact_reflect_fewshot_context_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 a question to answer. You were unsuccessful in answering the question either because you guessed the wrong answer with [<answer>] 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.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

Previous trial:
Relevant Context: {context}
Question: {question}{scratchpad}

Reflection:"""
    return cot_noreact_reflect_fewshot_context_prompt




def make_reflect_template(prompt):
    return PromptTemplate(
                        input_variables=["reflections", "context", "question", "scratchpad"],
                        template = prompt,
                        )

REFLEXION_PROMPT_TEMPLATES = {"REFLEXION_FEWSHOT_PROMPT": [make_reflect_template(cot_reason_fewshot_prompt(cot_fewshot_SIMPLE6)), make_template(cot_reflect_fewshot_prompt(reflect_fewshot))],
                              "REFLEXION_NOHARD_FEWSHOT_PROMPT": [make_reflect_template(cot_reason_fewshot_prompt(cot_fewshot_SIMPLE6)), make_template(cot_reflect_nohard_fewshot_prompt(reflect_fewshot))],
                              "REFLEXION_FEWSHOT_CONTEXT_PROMPT": [make_reflect_template(cot_reason_fewshot_context_prompt(cot_fewshot_context)), make_template(cot_reflect_fewshot_context_prompt(reflect_context_fewshot))],
                              
                              "REFLEXION_NOREACT_FEWSHOT_PROMPT": [make_reflect_template(cot_noreact_reason_fewshot_prompt(cot_noreact_fewshot_SIMPLE6)), make_template(cot_noreact_reflect_fewshot_prompt(reflect_noreact_fewshot))],
                              "REFLEXION_NOREACT_FEWSHOT_CONTEXT_PROMPT": [make_reflect_template(cot_noreact_reason_fewshot_context_prompt(cot_noreact_fewshot_context)), make_template(cot_noreact_reflect_fewshot_context_prompt(reflect_noreact_context_fewshot))],
                              }
# print(REFLEXION_PROMPT_TEMPLATES["REFLEXION_FEWSHOT_CONTEXT_PROMPT"][0].format(reflections="REFLECTION", context="CONTEXT", question="QUESTION", scratchpad="SCRATCH"))
# print(REFLEXION_PROMPT_TEMPLATES["REFLEXION_FEWSHOT_CONTEXT_PROMPT"][1].format(context="CONTEXT", question="QUESTION", scratchpad="SCRATCH"))

############################################



####### MAD Prompt #######
def mad_start_noreact_fewshot_prompt(fewshot):
    cot_noreact_fewshot_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be a phrase, in the form [answer], at the end of your response.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)
{context}
Question: {question}{scratchpad}
Answer:"""
    return cot_noreact_fewshot_prompt
def mad_start_noreact_fewshot_context_prompt(fewshot):
    mad_start_noreact_fewshot_context_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be a phrase, in the form [answer], at the end of your response.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

Relevant Context: {context}
Question: {question}{scratchpad}
Answer:"""
    return mad_start_noreact_fewshot_context_prompt

_mad_summarize = """Question: {question}

Answer: {answer}

Summarize the answer to the question answering problem. Your should include the final answer in the form [answer], at the end of your summary.
"""

MAD_SUMMARIZE = PromptTemplate(
                        input_variables=["question","answer"],
                        template = _mad_summarize,
                        )

MAD_PROMPT_TEMPLATES = {"MAD_NOREACT_FEWSHOT_START": [make_template(mad_start_noreact_fewshot_prompt(cot_noreact_fewshot_SIMPLE6)), ""],
                        "MAD_NOREACT_FEWSHOT_CONTEXT_START": [make_template(mad_start_noreact_fewshot_context_prompt(cot_noreact_fewshot_context)), ""],
                        "MAD_NOREACT_FEWSHOT_REFEWSHOT_START": [make_template(mad_start_noreact_fewshot_prompt(cot_noreact_fewshot_SIMPLE6)), cot_noreact_fewshot_SIMPLE6],
                        "MAD_NOREACT_FEWSHOT_REFEWSHOT_CONTEXT_START": [make_template(mad_start_noreact_fewshot_context_prompt(cot_noreact_fewshot_context)), cot_noreact_fewshot_context],
                                
                                "MAD_SUMMARIZE": MAD_SUMMARIZE,
}

############################################


############################################

####### Recursive CoT Prompt #######
recur_cot_eval_prompt = """Question: {question}\n\n""" + \
"""Answer: {answer}\n\n""" + \
"""Does the answer correctly answer the question?"""

recur_cot_eval_context_prompt = """Relavant Context: {context}\n""" + \
"""Question: {question}\n\n""" + \
"""Answer: {answer}\n\n""" + \
"""Does the answer correctly answer the question?"""


hotpotQA_RECURCOT_EVAL_PROMPT = PromptTemplate(
                        input_variables=["question", "answer"],
                        template = recur_cot_eval_prompt)
hotpotQA_RECURCOT_EVAL_CONTEXT_PROMPT = PromptTemplate(
                        input_variables=["context", "question", "answer"],
                        template = recur_cot_eval_context_prompt)

hotpotQA_RECURCOT_TEMPLATES = {"HOTPOTQA_RECURCOT_EVAL_PROMPT": [COT_PROMPT_TEMPLATES["COT_NOREACT_FEWSHOT_PROMPT"],hotpotQA_RECURCOT_EVAL_PROMPT],
                "HOTPOTQA_RECURCOT_EVAL_CONTEXT_PROMPT": [COT_PROMPT_TEMPLATES["COT_NOREACT_FEWSHOT_CONTEXT_PROMPT"],hotpotQA_RECURCOT_EVAL_CONTEXT_PROMPT],
}

############################################




#$$$$$$$$$$$$$$$$$$$$$$$$$$$$ NEW STRATEGIES $$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$

#### SELF DISCOVER ###############################
_execute = "Task Examples with answers:\n" + f"{cot_noreact_fewshot_context}\n\n" + """Using the following reasoning structure: {reasoning_structure}

Solve this task. You final answer should be a phrase, in the form [answer], at the end of your response:
Relevant Context: {context}
Question: {question}"""

hotpotQA_EXECUTE = PromptTemplate(input_variables=["reasoning_structure", "question", "context"],template = _execute)