from agents.agent import Agent
from agents.specialist_agent import SpecialistAgent
from agents.frontier_agent import FrontierAgent
from agents.neural_network_agent import NeuralNetworkAgent
from agents.preprocessor import Preprocessor
from free_config import WEIGHTS_WITH_SPECIALIST, WEIGHTS_WITHOUT_SPECIALIST

class EnsembleAgent(Agent):
    name = "Ensemble Agent"
    color = Agent.YELLOW

    def __init__(self, collection):
        self.log("Initializing Ensemble Agent")
        self.specialist = SpecialistAgent()
        self.frontier = FrontierAgent(collection)
        self.neural_network = NeuralNetworkAgent()
        self.preprocessor = Preprocessor()
        self.use_specialist = self.specialist.available
        self.use_neural = self.neural_network.available
        parts = ["Frontier"]
        if self.use_specialist:
            parts.append(f"Specialist [{getattr(self.specialist, 'mode', 'on')}]")
        if self.use_neural:
            parts.append("Neural Network")
        self.log(f"Ensemble Agent is ready ({' + '.join(parts)})")

    def price(self, description: str) -> float:
        self.log("Running Ensemble Agent - preprocessing text")
        rewrite = self.preprocessor.preprocess(description)
        self.log(f"Pre-processed text using {self.preprocessor.model_name}")
        frontier = self.frontier.price(rewrite)

        specialist = self.specialist.price(rewrite) if self.use_specialist else 0.0
        neural = self.neural_network.price(rewrite) if self.use_neural else 0.0

        if self.use_specialist and self.use_neural:
            combined = frontier * 0.8 + specialist * 0.1 + neural * 0.1
        elif self.use_specialist:
            combined = frontier * 0.85 + specialist * 0.15
        elif self.use_neural:
            combined = frontier * 0.85 + neural * 0.15
        else:
            combined = frontier

        self.log(f"Ensemble Agent complete - returning ${combined:.2f}")
        return combined