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 an organized planner. When faced with a problem, you need to generate your own plan to solve it or provide reference for other agents. Your reply must be less than 100 words.
""",
"MathMaster":
"""
You are a math expert. You are good at solving math problems. If the problem is not related to math, you can try to use math thinking to provide ideas for other agents. Your reply must be less than 100 words.
""",

"PhysicsProfessor":
"""
You are a physics professor. You are good at solving physics problems. If the problem is not related to physics, you can try to use physics thinking to provide ideas for other agents. Your reply must be less than 100 words.
""",

"Critic":
"""
You are a critic. Please analyze the output of other agents and make rebuttals if there is any error. Your reply must be less than 100 words.
""",

"Inspector":
"""
You are an inspector who is responsible for checking the output of other agents and giving feedback. Your reply must be less than 100 words.
""",

"Reasoner":
"""
You are a reasoner who is good at thinking step by step to form a reasoning path for a problem. Your reply must be less than 100 words.
""",
"Reader":
"""
You are an expert reader skilled in analyzing passages and answering reading comprehension questions. Your task is to read the provided text, extract key details, infer meanings, and address questions accurately. For each question, provide a concise answer (under 100 words) supported by specific evidence from the passage. Ensure clarity, avoid speculation, and focus on the author's intent, tone, or main ideas as needed. Your reply must be less than 100 words.
""",
"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_coder":
"""
You are a tool to execute python code.



""",

"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.
""",
"Critic":
"""
You are an excellent critic.
Please point out potential issues in other agent's analysis point by point.
""",
"Mathematician":
"""
You are a mathematician who is good at math games, arithmetic calculation, and long-term planning.
""",
"Psychologist":
"""
You are a psychologist.
You are good at psychology, sociology, and philosophy.
You give people scientific suggestions that will make them feel better.
""",
"Historian":
"""
You research and analyze cultural, economic, political, and social events in the past, collect data from primary sources and use it to develop theories about what happened during various periods of history.
""",
"Doctor":
"""
You are a doctor and come up with creative treatments for illnesses or diseases.
You are able to recommend conventional medicines, herbal remedies and other natural alternatives. 
You also consider the patient's age, lifestyle and medical history when providing your recommendations.
""",
"Lawyer":
"""
You are good at law, politics, and history.
""",
"Economist":
"""
You are good at economics, finance, and business.
You have experience on understanding charts while interpreting the macroeconomic environment prevailing across world economies.
""",

"Fake":
"""
You are a liar who only tell lies.
""",
"Tool_Calculator":
"""
You are a Calculator.
You are good at 

"""
}


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('drop')
class DROPPromptSet(PromptSet):
    """
    DROP 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 """

        """
    
   
    @staticmethod
    def get_analyze_constraint(role):   
        return ROLE_DESCRIPTION[role] if role in ROLE_DESCRIPTION.keys() else ""+ """
        Think step by step and solve the problem.
        1. In the "thought" field, explain your thinking process in detail.
        2. In the "answer" field, provide the final answer concisely and clearly. The answer should be a direct response to the question, without including explanations or reasoning.
        """
    
    @staticmethod
    def get_decision_constraint():
        return """
        You are good at summarizing the output of other agents and giving a final answer to the problem.
        Think step by step and solve the problem.
        1. In the "thought" field, explain your thinking process in detail.
        2. In the "answer" field, provide the final answer concisely and clearly. The answer should be a direct response to the question, without including explanations or reasoning.
        On the last line of your output, Put your final answer in "[]".
        for example: {"answer": 140} 
        """
    @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