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):
        """
        Create an instance of Ensemble, by creating each of the models
        And loading the weights of the Ensemble
        """
        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
        if self.use_specialist:
            mode = getattr(self.specialist, "mode", "on")
            self.log(f"Ensemble Agent is ready (Frontier + Specialist [{mode}] + Neural Network)")
        else:
            self.log("Ensemble Agent is ready (Frontier + Neural Network only)")

    def price(self, description: str) -> float:
        """
        Run this ensemble model
        Ask each of the models to price the product
        Then return the weighted average price
        """
        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)
        neural_network = self.neural_network.price(rewrite)

        if self.use_specialist:
            specialist = self.specialist.price(rewrite)
            wf, ws, wn = WEIGHTS_WITH_SPECIALIST
            combined = frontier * wf + specialist * ws + neural_network * wn
        else:
            wf, wn = WEIGHTS_WITHOUT_SPECIALIST
            combined = frontier * wf + neural_network * wn

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