import time
from copy import deepcopy
from multiprocessing import Queue
from typing import Callable, List, Optional, Union

from evola.epso.academic_version.swarm import EvolutiveSwarm
from evola.simulation import Simulation


class EPSO(Simulation):
    def __init__(
        self,
        generations,
        size: int,
        chromossome_length: int,
        chromossome_low: Union[List[Union[float, int]], Union[float, int]],
        chromossome_high: Union[List[Union[float, int]], Union[float, int]],
        cost_function: Callable,
        cost_function_args: tuple,
        chromossome_dtypes: Union[List[type], type] = float,
        wi: float = 0.5,
        wm: float = 0.5,
        wc: float = 0.5,
        communication_probability: float = 1.0,
        export_top=0,
        description="",
    ) -> None:
        desc = description + " EPSO (p=%.2f)" % communication_probability

        pop = EvolutiveSwarm(
            size,
            chromossome_length,
            chromossome_low,
            chromossome_high,
            chromossome_dtypes,
            cost_function,
            cost_function_args,
            wi,
            wm,
            wc,
            communication_probability,
        )

        super().__init__(generations, pop, desc, export_top)

        self.population: EvolutiveSwarm
        return

    @property
    def best_particle(self):
        return self.population.global_best

    @property
    def best_solution(self):
        return self.population.global_best.chromossome

    def run_gui(self, hold=False):
        try:
            from matplotlib import pyplot as plt
        except ImportError as exc:
            raise RuntimeError(
                "Install matplotlib if you wish to use GUI simulation: `pip install matplotlib`"
            ) from exc

        self._init_graph(plt)
        history = []
        start_time = time.time()

        pbar, _ = self._get_pbar_if_tqdm_installed(range(self.generations))

        for _ in pbar:
            # EPSO steps
            self.population.reproduce()
            self.population.move()
            self.population.select()

            # Obter o custo melhor da população
            new_data = self.population.global_best.cost
            history.append(new_data)
            self.cost_history.append(new_data)
            # Atualizar gráfico
            self._update_graph(history)

        self.population.last_gen_update()

        end_time = time.time()

        self.best_cost = self.population.global_best.cost

        print(self._final_message(start_time, end_time))

        if self.export_top > 0:
            self.export(self.export_top)

        self._finish_graph(hold, plt)

        return

    def _run_once(self, thread_num):
        pbar, has_tqdm = self._get_pbar_if_tqdm_installed(range(self.generations), thread_num + 1)
        pop = deepcopy(self.population)
        for _ in pbar:
            # EPSO
            self.cost_history.append(deepcopy(pop.global_best.cost))
            if has_tqdm:
                pbar.set_description(
                    f"({self.population.size},{self.generations}) C = {pop.global_best.cost:.0f} euros"
                )
            pop.reproduce()
            pop.move()
            pop.select()
        pop.last_gen_update()
        return pop

    def _run_multiple(self, itera, thread_num):
        pbar, _ = self._get_pbar_if_tqdm_installed(range(itera), thread_num + 1)
        for _ in pbar:
            pop = deepcopy(self.population)
            for _ in range(self.generations):
                # EPSO
                pop.reproduce()
                pop.move()
                pop.select()
            pop.last_gen_update()
            self.best_hist.append(deepcopy(pop.global_best.cost))
            self.avg_hist.append(pop.average_cost())
        return pop

    def _run_no_verbose(self, itera):
        for _ in range(itera):
            pop = deepcopy(self.population)
            for _ in range(self.generations):
                # EPSO
                if itera == 1:
                    self.cost_history.append(deepcopy(pop.global_best.cost))
                pop.reproduce()
                pop.move()
                pop.select()
            pop.last_gen_update()
            self.best_hist.append(deepcopy(pop.global_best.cost))
            self.avg_hist.append(pop.average_cost())
        return pop

    def run_cli(self, q: Optional[Queue] = None, thread_num=0, verbose=True, itera=1):

        start_time = time.time()

        pop: EvolutiveSwarm
        if verbose and itera == 1:
            pop = self._run_once(thread_num)
        elif verbose and itera > 1:
            pop = self._run_multiple(itera, thread_num)
        else:
            pop = self._run_no_verbose(itera)

        end_time = time.time()

        if itera == 1:
            self.best_cost = pop.global_best.cost
        else:
            self.best_cost = min(self.best_hist)
        if verbose:
            print(self._final_message(start_time, end_time))

        self.population = pop

        if self.export_top > 0:
            self.export(self.export_top)
        if q:
            q.put(self)  # returns self in a queue, if multiprocessing is enabled
        return
