import numpy as np
from hyponic.optimizers.base_optimizer import BaseOptimizer


class AntColonyOptimization(BaseOptimizer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.alpha = kwargs.get('alpha', 1)
        self.beta = kwargs.get('beta', 1)
        self.rho = kwargs.get('rho', 0.5)
        self.q = kwargs.get('q', 1)

        self.pheromone = None
        self.population = None
        self.scores = None
        self.best_score = None
        self.best_solution = None

    def initialize(self, problem_dict):
        super().initialize(problem_dict)

        self.pheromone = np.ones((self.population_size, self.dimensions))
        self.best_score = np.inf if self.minmax == 'min' else -np.inf

        self.population = np.random.uniform(low=self.lb, high=self.ub, size=(self.population_size, self.dimensions))
        self.scores = np.zeros(self.population_size)
        for i in range(self.population_size):
            self.scores[i] = self.function(self.population[i])
            if self._minmax()(self.scores[i]) < self._minmax()(self.best_score):
                self.best_score = self.scores[i]
                self.best_solution = self.population[i]

    def evolve(self, epoch):
        new_population = np.zeros((self.population_size, self.dimensions))
        new_scores = np.zeros(self.population_size)

        for i in range(self.population_size):
            ant_position = np.random.uniform(low=self.lb, high=self.ub, size=self.dimensions)

            for j in range(self.dimensions):
                pheromone_weights = self.pheromone[:, j] ** self.alpha
                heuristic_weights = 1 / (np.abs(ant_position[j] - self.population[:, j]) + 1e-6) ** self.beta
                probs = pheromone_weights * heuristic_weights
                probs /= np.sum(probs)
                ant_position[j] = np.random.choice(self.population[:, j], p=probs)

            ant_score = self.function(ant_position)
            self._update_pheromone(ant_score)

            if self._minmax()(ant_score) < self._minmax()(self.best_score):
                self.best_solution = ant_position
                self.best_score = ant_score

            new_population[i] = ant_position
            new_scores[i] = ant_score

        self.population = new_population
        self.scores = new_scores

    def get_best_score(self):
        return self.best_score

    def get_best_solution(self):
        return self.best_solution

    def _initialize_pheromone(self):
        self.pheromone = np.ones((self.population_size, self.dimensions))

    def _update_pheromone(self, ant_score):
        delta_pheromone = np.zeros((self.population_size, self.dimensions))
        delta_pheromone[:, :] = self.q / ant_score
        self.pheromone = (1 - self.rho) * self.pheromone + self.rho * delta_pheromone
