import numpy as np

from abc import ABC, abstractmethod
from functions import Function
from typing import Iterable, Tuple


class Solver(ABC):
    def __init__(self, function: Function):
        self.function = function

    @abstractmethod
    def rank(self, samples: Iterable[Iterable[float]], elite_size: int) -> (Iterable[Tuple], Iterable[Tuple], float):
        pass

    @abstractmethod
    def sample(self, n: int) -> Iterable[float]:
        pass

    @abstractmethod
    def update(elite: Iterable[Tuple]):
        pass


class SimpleEvolutionStrategy(Solver):
    def __init__(self, function: Function, mu: Iterable[float], sigma: Iterable[float] = None):
        if sigma and len(mu) != len(sigma):
            raise Exception('Length of mu and sigma must match')
        super().__init__(function)
        self.mu = mu
        if sigma:
            self.sigma = sigma
        else:
            self.sigma = [1] * len(mu)

    def rank(self, samples: Iterable[Iterable[float]], elite_size: int) -> (Iterable[Tuple], Iterable[float]):
        fitness = self.function.eval(*samples.transpose())
        samples = samples[np.argsort(fitness)]
        fitness = np.sort(fitness)
        elite = samples[0:elite_size]
        return elite, fitness[0]

    def sample(self, n: int) -> Iterable[Iterable[float]]:
        return np.array([np.random.multivariate_normal(self.mu, np.diag(self.sigma)) for _ in range(n)])

    def update(self, elite: Iterable[Tuple]):
        self.mu = elite[0]