from langchain import PromptTemplate, LLMChain
from langchain.chat_models import ChatOpenAI

from experts.base_expert import BaseExpert


class Reducer(BaseExpert):

    ROLE_DESCRIPTION = 'You are an expert that responsible for summarize the comment of all other experts then conclude the final answer'
    FORWARD_TASK = '''Now, you are an expert of Operations Research.
You are supposed to give the final code of an problem.
Text description of the problem: {problem_description}
Your colleagues are all experts in various related fields. They have given their own insights. I hope you will carefully refer to these comments when giving the final code:
{comment_text}

Based on all the expert comments above, please generate the final Python code that solves this optimization problem.
The code should:
1. Import necessary libraries (like gurobipy, numpy, etc.)
2. Define a function that implements the optimization model
3. Include proper variable definitions, constraints, and objective function
4. Return the optimal solution

No code is required outside the function except for the import package (No test code).
Your final code is as following:
'''

    def __init__(self, model):
        super().__init__(
            name='Reducer',
            description='Reduce all comments given by other experts',
            model=model
        )

        self.llm = ChatOpenAI(
            model_name=model,
            temperature=0,
            api_key="/",  # Dummy key for RITS
            base_url="https://inference-3scale-apicast-production.apps.rits.fmaas.res.ibm.com/deepseek-v3-h200/v1",
            default_headers={"RITS_API_KEY": ""}
        )
                
        # Recreate the chain with updated LLM
        self.forward_chain = LLMChain(
            llm=self.llm,
            prompt=PromptTemplate.from_template(self.forward_prompt_template)
        )

    def forward(self, problem, comment_pool):
        """
        Generate final answer by reducing all expert comments.
        
        Args:
            problem: Problem dictionary with description
            comment_pool: CommentPool containing all expert comments
            
        Returns:
            str: Final generated answer/code
        """
        try:
            comment_text = comment_pool.get_current_comment_text()
            
            if not comment_text or comment_text.strip() == 'There is no comment available, please ignore this section.\n':
                # Fallback: generate code directly from problem description
                fallback_prompt = f"""
As an Operations Research expert, generate Python code to solve this optimization problem:

Problem: {problem['description']}

Generate complete Python code with proper imports, variable definitions, constraints, and objective function.
"""
                answer = self.llm.predict(fallback_prompt)
            else:
                answer = self.forward_chain.predict(
                    problem_description=problem['description'], 
                    comment_text=comment_text
                )
            
            return answer
            
        except Exception as e:
            print(f"Error in reducer: {str(e)}")
            # Fallback: return basic template
            return f"""
# Error in generating solution: {str(e)}
# Problem: {problem['description'][:200]}...

import gurobipy as gp
from gurobipy import GRB

def solve_optimization_problem():
    # This is a placeholder implementation
    # The optimization model should be implemented here
    try:
        model = gp.Model("optimization_problem")
        
        # Variables, constraints, and objective should be defined here
        # based on the problem description
        
        model.optimize()
        
        if model.status == GRB.OPTIMAL:
            return model.objVal
        else:
            return None
            
    except Exception as e:
        print(f"Error in optimization: {{e}}")
        return None

# Call the function
result = solve_optimization_problem()
print(f"Optimal solution: {{result}}")
"""