from ..utils import ProgressBar
from .learner import Learner
import numpy as np


class EvolutionaryStrategies(Learner):

	def __init__(self, npop=100, lr=0.001, std=0.1):
		self.npop = npop
		self.lr = lr
		self.std = std

	def __call__(self, f, W, epochs=100):
		pbar = ProgressBar(epochs)
		rewards = np.zeros(self.npop)
		const = self.lr / (self.npop * self.std)
		for e in range(epochs):
			curr_reward = f()
			mutations = np.random.randn(*(self.npop,) + W.shape)
			jittered = self.std * mutations
			rewards *= 0
			for m, delta in enumerate(jittered):
				W += delta
				rewards[m] = f()
				W -= delta
			rewards -= curr_reward
			std_rewards = rewards.std()
			rewards -= rewards.mean()
			try:
				rewards /= std_rewards
			except Exception:
				pass
			W += const * np.dot(mutations.T, rewards)
			pbar.update()


ES = EvolutionaryStrategies
