from .base_agent import BaseAgent
from prompt.template import PROBLEM_MODELING_PROMPT, PROBLEM_MODELING_CRITIQUE_PROMPT, PROBLEM_MODELING_IMPROVEMENT_PROMPT
# from prompt.constants import modeling_methods


class ProblemModeling(BaseAgent):
    def __init__(self, llm):
        super().__init__(llm)
    
    def modeling_actor(self, modeling_problem: str, problem_analysis: str, modeling_methods: str, user_prompt: str=''):
        prompt = PROBLEM_MODELING_PROMPT.format(modeling_methods=modeling_methods, modeling_problem=modeling_problem, problem_analysis=problem_analysis, user_prompt=user_prompt).strip()
        return self.llm.generate(prompt)

    def modeling_critic(self, modeling_problem: str, problem_analysis: str, modeling_solution: str):
        prompt = PROBLEM_MODELING_CRITIQUE_PROMPT.format(modeling_problem=modeling_problem, problem_analysis=problem_analysis, modeling_solution=modeling_solution).strip()
        return self.llm.generate(prompt)

    def modeling_improvement(self, modeling_problem: str, problem_analysis: str, modeling_solution: str, modeling_solution_critique: str, user_prompt: str=''):
        prompt = PROBLEM_MODELING_IMPROVEMENT_PROMPT.format(modeling_problem=modeling_problem, problem_analysis=problem_analysis, modeling_solution=modeling_solution, modeling_solution_critique=modeling_solution_critique, user_prompt=user_prompt).strip()
        return self.llm.generate(prompt)

    def modeling(self, modeling_problem: str, problem_analysis: str, modeling_methods: str, round: int = 3, user_prompt: str = ''):
        modeling_solution = self.modeling_actor(modeling_problem, problem_analysis, modeling_methods, user_prompt)
        for i in range(round):
            print(f'Problem Modeling Round {i+1}')
            modeling_solution_critique = self.modeling_critic(modeling_problem, problem_analysis, modeling_solution)
            modeling_solution_improvement = self.modeling_improvement(modeling_problem, problem_analysis, modeling_solution, modeling_solution_critique, user_prompt)
            modeling_solution = modeling_solution_improvement
        return modeling_solution