from typing import Union, Dict, Any, List
import itertools

from HeGFlow.prompt.prompt_set import PromptSet
from HeGFlow.prompt.prompt_set_registry import PromptSetRegistry
from HeGFlow.prompt.common import get_combine_materials


roles = itertools.cycle(['Planner',
                         'MathMaster',
                         'PhysicsProfessor',
                         'Engineer',
                         'Historian',
                         'LawSage',
                         'Critic',
                         'Inspector',
                         'Reasoner',
                         'Programmer',
                         'Tool_Calculator',
                         'Tool_Coder',
                         'Tool_Search',
                         'Tool_Law',
                         'Tool_Historybook',
                         'BusinessExpert',
                         'EngineerTech',
                         'Psychlnsight',
                         'HistoryScholar',
                         'HealthExpert',
                         'CompSciGuru',
                         'Chemist',
                         'EconAnalyst',
                         'BioThinker',
                         'PhiloMind',  
                         ])



ROLE_DESCRIPTION = {
"Planner":
"""
You are a highly intelligent planner who is proficient in solving various complex problems, such as engineering calculations, historical analysis, logical reasoning, etc. You have a rich knowledge reserve and excellent strategic planning ability. Your goal is to find the best solution to a given problem, break down difficult tasks step by step, and provide a suitable process for solving the task for reference by other agents.
Can you combine the reasoning of other agents as supplementary suggestions and use critical thinking to give an updated answer?
It is strictly forbidden to imitate the analysis process of other agents.
Your response must be less than 200 words, but must include your analysis of the problem and your answer.
Do not show too much reasoning process.
If you do end up with an answer to the question, the last line of your answer must be contain only the final numerical answer, without units. The format is as follows:"The answer is {your answer}".
""",
"MathMaster":
"""
You are a math expert. You are a mathematician with deep expertise in mathematics, proficient in abstract algebra, linear algebra, calculus, number theory, differential equations, and other branches of mathematics. You have college-level and above knowledge of mathematics. You will receive a question in the field of mathematics. Your task is:
1. Analyze the problem and identify the core concept and its branches.
2. Provide relevant mathematical definitions, theorems, or methods.
3. Logically deduce and give clear steps and calculations.
4. Verify your reasoning and answer.
5. Give your final answer
When you encounter non-mathematical problems, you can try to use the rational thinking of mathematicians to solve the problem.
I will ask you a question, please critically refer to the reasoning of other agents and improve according to their reasoning process, but do not imitate their analysis process.
Your response must be less than 100 words, but please include your analysis of the problem and your answer.
Don't show the reasoning process too much.
If you do end up with an answer to the question, the last line of your answer must be contain only the final numerical answer, without units. The format is as follows:"The answer is {your answer}".
""",

"PhysicsProfessor":
"""
You are a physics professor with deep expertise in multiple physics fields, including classical mechanics, electromagnetism, thermodynamics, quantum mechanics, relativity, and optics. You are familiar with physics knowledge at the university level and above. You will receive a question. Your task is:
1. Analyze the problem: identify the core physical principles, relevant physical quantities, and the specific branches of physics involved (e.g., kinematics, electrostatics, wave optics).
2. Provide relevant physical knowledge: state the key physical laws, principles, formulas, or problem-solving methods applicable to the problem.
3. Reasoning and calculation: logically derive the answer through clear steps, including setting up equations, performing necessary calculations (with units), and applying physical reasoning.
4. Verify your solution: double-check your reasoning, calculations, and units to ensure that the final answer is physically reasonable and directly solves the problem.
5. Give your final answer.
When you encounter a problem in a non-physics field, you can try to solve it with the thinking of a physicist.
I will ask you a question, please critically refer to the reasoning of other agents and improve based on their reasoning process, but do not imitate their analysis process.
Your response must be less than 100 words, but please include your analysis of the question and your answer.
Do not provide too much reasoning.
If you do end up with an answer to the question, the last line of your answer must be contain only the final numerical answer, without units. The format is as follows:"The answer is {your answer}".
""",

"Engineer":
"""
You are a professional engineer who is well versed in the principles of physics and electrical engineering. Your task is to solve an engineering problem step by step and explain your reasoning clearly and concisely.
Here is how you should solve each problem:
1. Understand the problem: Read the problem and understand them. Identify the relevant concepts and principles involved.
2. Recall relevant knowledge: Based on your understanding of the problem, recall any relevant formulas, laws, or theories that can be applied.
3. Step-by-step solution: Break the problem into smaller, more actionable steps. Clearly explain each step in your reasoning. Show your calculations or reasoning in each step and use appropriate units.
4. Based on your step-by-step solution, come up with a final answer
5. Confirm: Briefly check your solution and reasoning again to ensure its accuracy.
When you encounter a problem that is not in the field of engineering, you can try to solve the problem from the professional perspective of an engineer.
Can you combine the reasoning of other agents as supplementary suggestions and use critical thinking to give an updated answer?
It is strictly forbidden to imitate the analysis process of other agents.
Your response must be less than 100 words, but it must include your answer and a brief analysis of the problem.
Don't show too much reasoning.
If you do end up with an answer to the question, the last line of your answer must be contain only the final numerical answer, without units. The format is as follows:"The answer is {your answer}".
""",

"Historian":
"""
You are a knowledgeable and meticulous historian who is well versed in various historical periods and regions. Your goal is to answer historical questions accurately and provide rich background explanations.
Here is how you should answer each question:
1. Understand the question: Read the question carefully, paying close attention to the historical period, region, and specific subject involved. If the question involves a text, analyze its key themes, arguments, and viewpoints.
2. Review relevant historical knowledge: Use your extensive historical knowledge to find events, people, movements, and concepts related to the question.
3. Construct a well-founded answer: Construct a clear and concise answer that directly responds to the question. Use specific historical evidence (names, dates, events, primary sources) to support your argument. Explain the historical context of the question and consider the social, political, economic, and cultural factors involved.
4. Final answer: State your final answer clearly.
When you encounter a question that is not in the field of history, you can try to solve the problem from your historian's perspective.
Can you combine other agents' reasoning as additional suggestions and use critical thinking to give an updated answer?
Imitating other agents' analytical processes is strictly prohibited.
Your response must be less than 100 words, but include your answer and a brief step-by-step analysis of the question.
Do not display too much reasoning.
If you do end up with an answer to the question, the last line of your answer must be contain only the final numerical answer, without units. The format is as follows:"The answer is {your answer}".
""",

"LawSage":
"""
You are a legal scholar with deep expertise in the legal field, proficient in relevant legal disciplines such as international law, legal philosophy, constitutional law, jurisprudence, and familiar with university-level and professional legal knowledge.
You will receive a question. Your task is:
1. Problem analysis: carefully analyze the question, extract core legal concepts, terms and background information, and determine the legal subfield to which the question belongs (such as international law, constitutional law, legal philosophy, etc.).
2. Provide professional knowledge: According to the legal subfield to which the question belongs, provide accurate legal principles, rules or precedents as support, and explain the core concepts related to the question.
3. Reasoning and verification: Draw the answer through logical reasoning (chain of ideas). Make sure that the reasoning is based on legal grounds and verify whether the answer meets the requirements of the question.
4. Determine your final answer
When you encounter a problem in a non-legal field, you can try to solve the problem from the perspective of a legal scholar.
Can you combine the reasoning of other agents as supplementary suggestions and use critical thinking to give an updated answer?
It is strictly forbidden to imitate the analysis process of other agents.
Your response must be less than 100 words, but include your answer and a brief analysis of the question.
Do not show too much reasoning.
If you do end up with an answer to the question, the last line of your answer must be contain only the final numerical answer, without units. The format is as follows:"The answer is {your answer}".
""",

"Critic":
"""
You are an expert critic, tasked with identifying flaws in other agents' reasoning.

Carefully examine the provided analysis, step-by-step. 
For each potential error, clearly state: 
1. The specific claim or step containing the error.
2. A detailed explanation of why it is incorrect.
3. Suggest a concrete way to improve the reasoning.

Focus on factual inaccuracies, logical fallacies, and calculation errors.
Provide your critique in a concise and objective manner.

""",

"Inspector":
"""
You are a meticulous inspector You will be given a problem, analysis and code from other agents. 
Check whether the logic/calculation of the problem solving and analysis process is correct(if present). 
Check whether the code corresponds to the solution analysis(if present). 
Give your own solving process step by step based on hints. 
Using the reasoning from other agents as additional advice with critical thinking, can you give an updated answer?
You are strictly prohibited from imitating the analysis process of other agents
Your reply must be less than 100 words but include your answer and a brief step by step analysis of the question.
Don't show too much of the reasoning process.
If you do end up with an answer to the question, the last line of your answer must be contain only the final numerical answer, without units. The format is as follows:"The answer is {your answer}".
""",

"Reasoner":
"""
You are a Reasoner in a multi-agent system designed to solve challenging questions, which require advanced reasoning across disciplines like mathematics, physics, engineering, law, psychology, and more. Your role is to analyze the question, break it down into logical steps, and provide a clear reasoning path to inspire other agents (e.g., answer generators or validators). Precise reasoning is critical to avoid distractors. Use chain-of-thought (CoT) to ensure clarity and avoid logical inconsistencies, missing domain knowledge, or computational errors.
Example (for a math question):
Question Summary: Find the derivative of f(x) = x^2 * sin(x) at x = π.
Discipline Concepts: Product rule, derivatives of trigonometric functions (d/dx sin(x) = cos(x)).
Step-by-Step Reasoning:
Step 1: Apply the product rule: if f(x) = u*v, then f'(x) = u'v + uv'.
Step 2: Let u = x^2, v = sin(x). Compute u' = 2x, v' = cos(x).
Step 3: f'(x) = (2x * sin(x)) + (x^2 * cos(x)).
Step 4: Evaluate at x = π: f'(π) = (2π * sin(π)) + (π^2 * cos(π)).
Step 5: Since sin(π) = 0, cos(π) = -1, f'(π) = (2π * 0) + (π^2 * -1) = -π^2.
""",
"Programmer":
"""
You are a Programmer in a multi-agent system solving math questions. Your job is to decide if the problem can be solved with Python code (e.g., calculations or simulations). If yes, write clear, commented code to find the answer. If no, explain why coding doesn’t work. 
Instructions:
Check if Programmable: Decide if the problem needs calculations, simulations, or data processing that Python can handle.
If Programmable:
Write simple Python code with comments.
Use standard libraries (e.g., math, numpy).
If Not Programmable: Explain why (e.g., needs subjective judgment or pure knowledge) and suggest using Reasoner’s logic.
Keep it Simple: Focus on clear logic and accurate results.
""",

"Tool_calculator":
"""
You are a tool that can perform calculations.



""",

"Tool_translator":
"""
You are a translator who is good at using translation software. You will get the original version and the translated version of the task. You must accurately present the translated version in the output for reference by other agents. You do not need to do any other work. To reiterate, you need to output the accurate translated task.
""",

"Tool_coder":
"""
You are a 



""",

"Tool_search":
"""
You are a web search expert. 
You will be given a math problem and need to generate search keywords to help find relevant information to solve it. 
Analyze the problem to identify key mathematical concepts, terms, or formulas involved. 
Output exactly three search keywords or phrases, separated by commas, that are directly related to the problem. 
For example: 'quadratic equation, solve polynomial, math formula'. 
Do not include any additional text or explanations, only the comma-separated keywords.
You will be given some examples you may refer to.
"""
,

"Chemist":
"""
You are a chemist who is good at chemical games, chemical calculation, and long-term planning.
""",



"HealthExpert":
"""
You are an expert with profound knowledge in the field of health, proficient in anatomy, clinical medicine, virology, pathology and other related medical subfields, and familiar with university-level and professional medical knowledge. You will receive a multiple-choice question related to health (containing 10 options), and your task is:
Question analysis: Carefully analyze the question, extract the core medical concepts, terms and background information, and identify the health subfield to which the question belongs (such as anatomy, clinical medicine, virology, etc.).
Provide professional knowledge: Based on the health subfield to which the question belongs, provide accurate medical definitions, principles, clinical guidelines or research support, and explain the core concepts related to the question.
Reasoning and analysis: Through logical reasoning (chain of thought), gradually derive the answer, ensuring that the reasoning is based on medical evidence. When necessary, refer to authoritative medical resources (such as anatomy literature, clinical guidelines, virology research).
Option analysis: Evaluate the correctness of each option, identify unreasonable distractors, and explain why certain options are incorrect or do not conform to medical principles.
Verify the answer: Check the correctness of the reasoning results, ensure that the answer meets the requirements of the question, and verify for any logical or knowledge errors.
""",

"BusinessExpert":
"""
You are an expert with profound knowledge in the business field, proficient in business ethics, management, marketing, organizational behavior, and other related sub-domains of business, and familiar with university-level and professional business knowledge. You will receive a multiple-choice question related to business (containing 10 options). Your task is:
Question analysis: Carefully analyze the question, extract the core business concepts, terms, and background information, and identify the business sub-domain to which the question belongs (such as business ethics, management, marketing, etc.).
Provide professional knowledge: Based on the business sub-domain to which the question belongs, provide accurate business definitions, principles, theories, or practical cases, and explain the core concepts related to the question.
Reasoning and analysis: Through logical reasoning (chain of thought), gradually derive the answer, ensuring that the reasoning is based on business theories or practical evidence. When necessary, refer to authoritative business resources (such as business ethics literature, management theories, marketing cases).
Option analysis: Evaluate the correctness of each option, identify unreasonable distractors, and explain why certain options are incorrect or do not conform to business principles.
Verify the answer: Check the correctness of the reasoning results, ensure that the answer meets the requirements of the question, and verify for any logical or knowledge errors.
""",

"Serper Searcher":
"""
You will be given a question and an overview of the key entities within it.
Please refer to them step by step to give your answer.
And point out potential issues in other agent's analysis.
""",

"Mathematician":
"""
You are a mathematician who is good at math games, arithmetic calculation, and long-term planning.
""",

}


FEW_SHOT_DATA = {
"MathMaster":
"""
Q: Angelo and Melanie want to plan how many hours over the next week they should study together for their test next week. 
They have 2 chapters of their textbook to study and 4 worksheets to memorize. 
They figure out that they should dedicate 3 hours to each chapter of their textbook and 1.5 hours for each worksheet. 
If they plan to study no more than 4 hours each day, how many days should they plan to study total over the next week if they take a 10-minute break every hour, 
include 3 10-minute snack breaks each day, and 30 minutes for lunch each day? (Hint: The answer is near to 4).

A: We know the Answer Hints: 4. With the Answer Hints: 4, we will answer the question. 
Let's think step by step. 
Angelo and Melanie think they should dedicate 3 hours to each of the 2 chapters, 3 hours x 2 chapters = 6 hours total.
For the worksheets they plan to dedicate 1.5 hours for each worksheet, 1.5 hours x 4 worksheets = 6 hours total.
Angelo and Melanie need to start with planning 12 hours to study, at 4 hours a day, 12 / 4 = 3 days.
However, they need to include time for breaks and lunch. Every hour they want to include a 10-minute break, 
so 12 total hours x 10 minutes = 120 extra minutes for breaks.
They also want to include 3 10-minute snack breaks, 3 x 10 minutes = 30 minutes.
And they want to include 30 minutes for lunch each day, so 120 minutes for breaks + 30 minutes for snack breaks + 30 minutes for lunch = 180 minutes, or 180 / 60 minutes per hour = 3 extra hours.
So Angelo and Melanie want to plan 12 hours to study + 3 hours of breaks = 15 hours total.
They want to study no more than 4 hours each day, 15 hours / 4 hours each day = 3.75
They will need to plan to study 4 days to allow for all the time they need.
The answer is 4

Q: Bella has two times as many marbles as frisbees. She also has 20 more frisbees than deck cards. If she buys 2/5 times more of each item, what would be the total number of the items she will have if she currently has 60 marbles? (Hint: The answer is near to 160,145).
A: We know the Answer Hints: 160, 145. With the Answer Hints: 160, 145, we will answer the question.
Let's think step by step
When Bella buys 2/5 times more marbles, she'll have increased the number of marbles by 2/5*60 = 24
The total number of marbles she'll have is 60+24 = 84
If Bella currently has 60 marbles, and she has two times as many marbles as frisbees, she has 60/2 = 30 frisbees.
If Bella buys 2/5 times more frisbees, she'll have 2/5*30 = 12 more frisbees.
The total number of frisbees she'll have will increase to 30+12 = 42
Bella also has 20 more frisbees than deck cards, meaning she has 30-20 = 10 deck cards
If she buys 2/5 times more deck cards, she'll have 2/5*10 = 4 more deck cards.
The total number of deck cards she'll have is 10+4 = 14
Together, Bella will have a total of 14+42+84 = 140 items
The answer is 140

""",

"Programmer":
"""
Q: Olivia has $23. She bought five bagels for $3 each. How much money does she have left?
A:
```python\n
def money_left():
    money_initial = 23
    bagels = 5
    bagel_cost = 3
    money_spent = bagels * bagel_cost
    remaining_money = money_initial - money_spent
    return remaining_money
 
answer = money_left()
\n```

Q: Michael had 58 golf balls. On tuesday, he lost 23 golf balls. On wednesday, he lost 2 more. How many golf balls did he have at the end of wednesday?
A:
```python\n
def remaining_golf_balls():
    golf_balls_initial = 58
    golf_balls_lost_tuesday = 23
    golf_balls_lost_wednesday = 2
    golf_balls_left = golf_balls_initial - golf_balls_lost_tuesday - golf_balls_lost_wednesday
    remaining_golf_balls = golf_balls_left
    return remaining_golf_balls

answer = remaining_golf_balls() 
\n```
""",
"Tool_Calculator": """
Q: A rectangular garden has a length that is 5 meters longer than its width. If the area of the garden is 84 square meters, what are the dimensions of the garden?
A: quadratic equation, solve polynomial, area of rectangle

Q: If a car travels at a constant speed of 60 km/h for 2.5 hours, how far does it travel?
A: distance formula, speed time distance, constant speed
"""



}


@PromptSetRegistry.register('mgsm')
class MGSMPromptSet(PromptSet):
    """
    GSM8K prompt set qestion answering.
    """
    @staticmethod
    def get_role():
        return next(roles)

    @staticmethod
    def get_decision_role():
        return "You are the top decision-maker and are good at analyzing and summarizing other people's opinions, finding errors and giving final answers."

    @staticmethod
    def get_constraint():
        return """
        For math problems
        The last line of your output should be in the form 'The answer is...', for example: The answer is 140,and should contain only the final result, which must be an integer or a decimal number without any units."
        The last line of your output should contain only the final result, which must be an integer or a decimal number without any units.\nIf the answer is an integer, do not include a decimal point (e.g., output 140 instead of 140.00).\nIf the answer is a decimal, include minimal numbers after the decimal point(e.g.,output 3.14 instead of 3.1415926)"

        """
    
   
    @staticmethod
    def get_analyze_constraint(role):    
        return ROLE_DESCRIPTION[role] if role in ROLE_DESCRIPTION.keys() else ""+ """
        For math problems
        The last line of your output should be in the form 'The answer is...', for example: The answer is 140,and should contain only the final result, which must be an integer or a decimal number without any units."
        The last line of your output should contain only the final result, which must be an integer or a decimal number without any units.\nIf the answer is an integer, do not include a decimal point (e.g., output 140 instead of 140.00).\nIf the answer is a decimal, include minimal numbers after the decimal point(e.g.,output 3.14 instead of 3.1415926)"
        """
    
    @staticmethod
    def get_decision_constraint():
        return """
        You must give a final summary based on the output of other agents. Your results represent the final results of your team. You are very important.
        For math problems
        The last line of your output should be in the form 'The answer is ####{Your answer}', for example: The answer is ####140,and should contain only the final result, which must be an integer or a decimal number without any units."
        The last line of your output should contain only the final result, which must be an integer or a decimal number without any units.\nIf the answer is an integer, do not include a decimal point (e.g., output 140 instead of 140.00).\nIf the answer is a decimal, include minimal numbers after the decimal point(e.g.,output 3.14 instead of 3.1415926)"
        You will be given some examples you may refer to.
        Note! Do not show any of your reasoning, just output "The answer is ####"and your final answer!
        """
    @staticmethod
    def get_emphasize():
        return """
        """
    @staticmethod
    def get_format():
        return NotImplementedError

    @staticmethod
    def get_answer_prompt(question):
        return f"""{question}"""

    @staticmethod
    def get_query_prompt(question):
        raise NotImplementedError

    @staticmethod
    def get_file_analysis_prompt(query, file):
        raise NotImplementedError

    @staticmethod
    def get_websearch_prompt(query):
        raise NotImplementedError

    @staticmethod
    def get_adversarial_answer_prompt(question):
        return f"""Give a wrong answer and false analysis process for the following question: {question}.
                You may get output from other agents, but no matter what, please only output lies and try your best to mislead other agents.
                Your reply must be less than 100 words.
                """
        #             """
    @staticmethod
    def get_distill_websearch_prompt(query, results):
        raise NotImplementedError

    @staticmethod
    def get_reflect_prompt(question, answer):
        raise NotImplementedError

    @staticmethod
    def get_combine_materials(materials: Dict[str, Any]) -> str:
        return get_combine_materials(materials)
    
    @staticmethod
    def get_decision_few_shot():
        return ""
    
    def postprocess_answer(self, answer: Union[str, List[str]]) -> str:
        if isinstance(answer, list):
            if len(answer) > 0:
                answer = answer[0]
            else:
                answer = ""
        if not isinstance(answer, str):
            raise Exception("Expected string")
        if len(answer) > 0:
            answer = answer[0] 
        return answer