from langchain.prompts import PromptTemplate
from .CSQA_fewshot import *

####### Standard Prompt & CoT Prompt #######
cot_zeroshot_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be of the form 'So the answer is (choice).', at the end of your response, where choice = a/b/c/d/e
Question: {question}
Answer:"""
def cot_fewshot_prompt(fewshot):
    cot_fewshot_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be of the form 'So the answer is (choice).', at the end of your response, where choice = a/b/c/d/e
Here are some examples:""" + \
f"\n{fewshot}\n\n" + \
"""Question: {question}
Answer:"""
    return cot_fewshot_prompt
def cot_justfewshot_prompt(fewshot):
    cot_justfewshot_prompt = f"{fewshot}\n\n" + \
"""Question: {question}
Answer:"""
    return cot_justfewshot_prompt


### io
def io_fewshot_prompt(fewshot):
    io_fewshot_prompt = """Solve a question answering task. You final answer should be a single choice, in the form [choice]. Don't output anything else.
Here are some examples:""" + \
f"\n{fewshot}\n" + \
"""Question: {question}
Asnwer:"""
    return io_fewshot_prompt 


def make_template(prompt):
    return PromptTemplate(
                        input_variables=["question"],
                        template = prompt,
                        )

CSQA_COT_PROMPT_TEMPLATES = {"CSQA_COT_ZEROSHOT_PROMPT": make_template(cot_zeroshot_prompt),
                        "CSQA_COT_FEWSHOT_PROMPT": make_template(cot_fewshot_prompt(fewshot_examples)),    
                        "CSQA_COT_JUSTFEWSHOT_PROMPT": make_template(cot_justfewshot_prompt(fewshot_examples)), 
                        }
# print(CSQA_COT_PROMPT_TEMPLATES["COT_JUSTFEWSHOT_PROMPT"].format(question="What do people use to absorb extra ink from a fountain pen?\n(a) shirt pocket\n(b) calligrapher’s hand\n(c) inkwell\n(d) desk drawer\n(e) blotter"))
############################################

####### MAD Prompt #######
def mad_start_fewshot_prompt(fewshot):
    mad_start_fewshot_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be of the form 'So the answer is (choice).', at the end of your response, where choice = a/b/c/d/e
Here are some examples:""" + \
f"\n{fewshot}\n\n" + \
"""Question: {question}
Answer:"""
    return mad_start_fewshot_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,
                        )

CSQA_MAD_PROMPT_TEMPLATES = {"CSQA_MAD_FEWSHOT_START": [make_template(mad_start_fewshot_prompt(fewshot_examples)), ""],
                        "MAD_SUMMARIZE": MAD_SUMMARIZE,
}
# print(MAD_PROMPT_TEMPLATES["CSQA_MAD_FEWSHOT_START"][0].format(question="What do people use to absorb extra ink from a fountain pen?\n(a) shirt pocket\n(b) calligrapher’s hand\n(c) inkwell\n(d) desk drawer\n(e) blotter"))
############################################

####### 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_reason_fewshot_prompt = """Solve a question answering task. Explain your reasonings first, then finish with your answer. You final answer should be of the form 'So the answer is (choice).', at the end of your response, where choice = a/b/c/d/e
Here are some examples:""" + \
f"\n{fewshot}\n\n" + \
"""(END OF EXAMPLES)
{reflections}
Question: {question}
Answer:"""
    return cot_reason_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 [choice] or there is a reasoning error in your steps. 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:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

Previous trial:
Question: {question}{scratchpad}

Reflection:"""
    return cot_reflect_fewshot_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 [choice] or there is a reasoning error in your steps. 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:""" + \
f"\n{fewshot}\n" + \
"""(END OF EXAMPLES)

Previous trial:
Question: {question}{scratchpad}

Reflection:"""
    return cot_reflect_fewshot_prompt


def make_reflect_template(prompt):
    return PromptTemplate(
                        input_variables=["reflections", "question"],
                        template = prompt,
                        )
def make_reflect_prompt_template(prompt):
    return PromptTemplate(
                        input_variables=["question", "scratchpad"],
                        template = prompt,
                        )
CSQA_REFLEXION_PROMPT_TEMPLATES = {
                              "CSQA_REFLEXION_NOHARD_FEWSHOT_PROMPT": [make_reflect_template(cot_reason_fewshot_prompt(fewshot_examples)), make_reflect_prompt_template(cot_reflect_nohard_fewshot_prompt(reflect_fewshot))],             
                              "CSQA_REFLEXION_NOREACT_FEWSHOT_PROMPT": [make_reflect_template(cot_reason_fewshot_prompt(fewshot_examples)), make_reflect_prompt_template(cot_reflect_fewshot_prompt(reflect_fewshot))],
                              }
# print(CSQA_REFLEXION_PROMPT_TEMPLATES["CSQA_REFLEXION_NOREACT_FEWSHOT_PROMPT"][0].format(reflections="REFLECTION", question="QUESTION"))
# print(CSQA_REFLEXION_PROMPT_TEMPLATES["CSQA_REFLEXION_NOREACT_FEWSHOT_PROMPT"][1].format(question="QUESTION", scratchpad="SCRATCH"))

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

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


CSQA_RECURCOT_EVAL_PROMPT = PromptTemplate(
                        input_variables=["question", "answer"],
                        template = recur_cot_eval_prompt)


CSQA_RECURCOT_TEMPLATES = {"CSQA_RECURCOT_EVAL_PROMPT": [CSQA_COT_PROMPT_TEMPLATES["CSQA_COT_FEWSHOT_PROMPT"],CSQA_RECURCOT_EVAL_PROMPT],
}

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



### SELF DISCOVER ########################
reasoning_modules = [
        "1. How could I devise an experiment to help solve that problem?",
        "2. Make a list of ideas for solving this problem, and apply them one by one to the problem to see if any progress can be made.",
        #"3. How could I measure progress on this problem?",
        "4. How can I simplify the problem so that it is easier to solve?",
        "5. What are the key assumptions underlying this problem?",
        "6. What are the potential risks and drawbacks of each solution?",
        "7. What are the alternative perspectives or viewpoints on this problem?",
        "8. What are the long-term implications of this problem and its solutions?",
        "9. How can I break down this problem into smaller, more manageable parts?",
        "10. Critical Thinking: This style involves analyzing the problem from different perspectives, questioning assumptions, and evaluating the evidence or information available. It focuses on logical reasoning, evidence-based decision-making, and identifying potential biases or flaws in thinking.",
        "11. Try creative thinking, generate innovative and out-of-the-box ideas to solve the problem. Explore unconventional solutions, thinking beyond traditional boundaries, and encouraging imagination and originality.",
        #"12. Seek input and collaboration from others to solve the problem. Emphasize teamwork, open communication, and leveraging the diverse perspectives and expertise of a group to come up with effective solutions.",
        "13. Use systems thinking: Consider the problem as part of a larger system and understanding the interconnectedness of various elements. Focuses on identifying the underlying causes, feedback loops, and interdependencies that influence the problem, and developing holistic solutions that address the system as a whole.",
        "14. Use Risk Analysis: Evaluate potential risks, uncertainties, and tradeoffs associated with different solutions or approaches to a problem. Emphasize assessing the potential consequences and likelihood of success or failure, and making informed decisions based on a balanced analysis of risks and benefits.",
        #"15. Use Reflective Thinking: Step back from the problem, take the time for introspection and self-reflection. Examine personal biases, assumptions, and mental models that may influence problem-solving, and being open to learning from past experiences to improve future approaches.",
        "16. What is the core issue or problem that needs to be addressed?",
        "17. What are the underlying causes or factors contributing to the problem?",
        "18. Are there any potential solutions or strategies that have been tried before? If yes, what were the outcomes and lessons learned?",
        "19. What are the potential obstacles or challenges that might arise in solving this problem?",
        "20. Are there any relevant data or information that can provide insights into the problem? If yes, what data sources are available, and how can they be analyzed?",
        "21. Are there any stakeholders or individuals who are directly affected by the problem? What are their perspectives and needs?",
        "22. What resources (financial, human, technological, etc.) are needed to tackle the problem effectively?",
        "23. How can progress or success in solving the problem be measured or evaluated?",
        "24. What indicators or metrics can be used?",
        "25. Is the problem a technical or practical one that requires a specific expertise or skill set? Or is it more of a conceptual or theoretical problem?",
        "26. Does the problem involve a physical constraint, such as limited resources, infrastructure, or space?",
        "27. Is the problem related to human behavior, such as a social, cultural, or psychological issue?",
        "28. Does the problem involve decision-making or planning, where choices need to be made under uncertainty or with competing objectives?",
        "29. Is the problem an analytical one that requires data analysis, modeling, or optimization techniques?",
        "30. Is the problem a design challenge that requires creative solutions and innovation?",
        "31. Does the problem require addressing systemic or structural issues rather than just individual instances?",
        "32. Is the problem time-sensitive or urgent, requiring immediate attention and action?",
        "33. What kinds of solution typically are produced for this kind of problem specification?",
        "34. Given the problem specification and the current best solution, have a guess about other possible solutions."
        "35. Let’s imagine the current best solution is totally wrong, what other ways are there to think about the problem specification?"
        "36. What is the best way to modify this current best solution, given what you know about these kinds of problem specification?"
        "37. Ignoring the current best solution, create an entirely new solution to the problem."
        #"38. Let’s think step by step."
        "39. Let’s make a step by step plan and implement it with good notation and explanation."
]
csqa_task_examples = """Example 1: What do people use to absorb extra ink from a fountain pen? 
Answer Choices:
(a) shirt pocket
(b) calligrapher’s hand
(c) inkwell
(d) desk drawer
(e) blotter
Example 2: What home entertainment equipment requires cable?
Answer Choices:
(a) radio shack
(b) substation
(c) television
(d) cabinet
Example 3: The fox walked from the city into the forest, what was it looking for?
Answer Choices:
(a) pretty flowers
(b) hen house
(c) natural habitat
(d) storybook
Example 4: Sammy wanted to go to where the people were. Where might he go?
Answer Choices:
(a) populated areas
(b) race track
(c) desert
(d) apartment
(e) roadblock
Example 5: Where do you put your grapes just before checking out?
Answer Choices:
(a) mouth
(b) grocery cart
(c) supermarket
(d) fruit basket
(e) fruit market"""
hotpotQA_task_examples = """Example 1: 
Relevant Context: The Nile River is the longest river in the world, spanning approximately 6,650 kilometers (4,132 miles) in length. It flows through eleven countries in northeastern Africa, including Egypt, Sudan, and Uganda.
Question: What is the longest river in the world?
Example 2:
Relevant Context: Ludwig van Beethoven was a German composer and pianist who was a crucial figure in the transition between the Classical and Romantic eras in Western classical music. One of his most famous compositions is the Symphony No. 9, also known as the "Choral" symphony.
Question: Which composer created Symphony No. 9?
Example 3:
Relevant Context: Photosynthesis is the process by which green plants and some other organisms convert light energy into chemical energy. During this process, plants absorb sunlight, carbon dioxide, and water to produce glucose and oxygen.
Question: What do plants produce during photosynthesis?"""
SELF_DISCOVER_TASKEXAMPLES = {"CSQA": csqa_task_examples, "hotpotQA": hotpotQA_task_examples}

_select = """Select several reasoning modules that are crucial to utilize in order to solve the given task:
All reasoning module descriptions:
{reasoning_modules}

Task Examples w/o answer:
{task_examples}

Select several reasoning modules that are crucial for solving the tasks above:"""

_adapt = """Rephrase and specify each reasoning module so that it better helps solving the task:
Selected reasoning modules:
{reasoning_modules}

Task Examples w/o answer:
{task_examples}

Adapt each reasoning module to better solve the tasks above:"""

_implement = """Operationalize the reasoning modules into a step-by-step reasoning plan in JSON format:

Examples:
Adpated reasoning descriptions:
How can I break down this problem into smaller, more manageable parts?

Reasoning Plan:
reasoning_about_colored_objects
{{
    \"Type and color of each item\":
    \"Number of items of each color\":
    \"Number of items of each type\":
    \"Number of items of each color and type\":
    \"Final answer\":
}}


Adpated reasoning descriptions:
{reasoning_modules}

Task Examples w/o answer:
{task_examples}

Implement a reasoning structure for solvers to follow step-by-step and arrive at correct answers:"""

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

Solve this task. You final answer should be of the form 'So the answer is (choice).', at the end of your response, where choice = a/b/c/d/e: {question}"""

SELECT = PromptTemplate(input_variables=["reasoning_modules", "task_examples"],template = _select,)
ADAPT = PromptTemplate(input_variables=["reasoning_modules", "task_examples"],template = _adapt)
IMPLEMENT = PromptTemplate(input_variables=["reasoning_modules", "task_examples"],template = _implement)
CSQA_EXECUTE = PromptTemplate(input_variables=["reasoning_structure", "question"],template = _execute)