import numpy as np
from abc import ABC, abstractmethod

from metanetwork.utils._base import Base

def get_optimizer(name):
	"""
	Lookup table of default optimizers.

	Parameters
	----------
	name : Optimizer, None, str
		Optimizer to look up. Must be one of:
		 - 'sgd' : Gradient Descent.
		 - 'momentum' : Gradient Descent with Momentum.
		 - 'rmsprop' : RMSProp.
		 - 'adam' : Adam.
		 - Regularizer : A custom implementation.
		 - None : Return None.
		Custom Optimizer must implement `update`
		functions.

	Returns
	-------
	optimizer : Optimizer or None
		The optimizer.
	"""
	if name == 'sgd' : return SGD()
	elif name == 'momentum' : return Momentum()
	elif name == 'rmsprop' : return RMSProp()
	elif name == 'adam' : return Adam()
	elif isinstance(name, (type(None), Optimizer)) : return name
	else : raise ValueError("Invalid optimizer")

class Optimizer(Base, ABC):
	"""
	Base Optimizer
	"""
	def __init__(self, *args, **kwargs):
		self.name = 'optimizer'
		self.memory = {}

	def setup(self, keys, *args, **kwargs):
		"""
		Setup the memory cache.

		Parameters
		----------
		keys : list
			List of key names for the cache.
		"""
		return

	def update(self, key, gradient, *args, **kwargs):
		"""
		Calculate the parameter update for a given
		parameter.

		Parameters
		----------
		key : str
			Name of the key as stored in the memory cache.

		gradient : float, array-like
			Gradient of the respective parameter.

		Returns
		-------
		update : float, array-like
			Update to be applied.
		"""
		if key not in self.memory.keys():
			raise ValueError("Optimizer not setup to handle these gradients")

	def _update_mem(self, key, d, beta):
		"""
		Update cache with momentum.

		Parameters
		----------
		key : str
			Name of the key as stored in the memory cache.

		d : float, array-like
			Current value.

		beta : float
			Momentum.
		"""
		if self.memory[key] is None : self.memory[key] = d
		else : self.memory[key] = beta * self.memory[key] + d

class SGD(Optimizer):
	"""
	Gradient Descent.

	Parameters
	----------
	alpha : float, default=0.001
		Learning Rate.
	"""
	def __init__(self, alpha=0.001):
		super().__init__()
		self.alpha = alpha
		self.name = 'sgd'

	def update(self, key, gradient):
		return self.alpha * gradient


class Momentum(Optimizer):
	"""
	Gradient Descent with Momentum.

	Parameters
	----------
	alpha : float, default=0.001
		Learning Rate.

	beta : float, default=0.9
		Momentum for past gradients.
	"""
	def __init__(self, alpha=0.001, beta=0.9):
		super().__init__()
		self.alpha = alpha
		self.beta = beta
		self.name = 'sgd'

	def setup(self, keys):
		for k in keys:
			self.memory['M'+k] = None

	def update(self, key, gradient):
		key = 'M' + key
		super().update(key, gradient)
		d = self.alpha + gradient
		self._update_mem(key, d, self.beta)
		return self.memory[key]


class RMSProp(Optimizer):
	"""
	RMSProp.

	Parameters
	----------
	alpha : float, default=0.001
		Learning Rate.

	beta : float, default=0.9
		Momentum for past squared gradients.
	"""
	def __init__(self, alpha=0.001, beta=0.9):
		super().__init__()
		self.alpha = alpha
		self.beta = beta
		self.name = 'rmsprop'

	def setup(self, keys):
		for k in keys:
			self.memory['E'+k] = None

	def update(self, key, gradient):
		key = 'E' + key
		super().update(key, gradient)
		d = (1 - self.beta) * self.square(gradient)
		self._update_mem(key, d, self.beta)
		return self.alpha / np.sqrt(self.memory[key] + 1e-8) * gradient


class Adam(Optimizer):
	"""
	Adam.

	Parameters
	----------
	alpha : float, default=0.001
		Learning Rate.

	beta1 : float, default=0.9
		Momentum for past gradients.

	beta2 : float, default=0.999
		Momentum for past squared gradients.
	"""
	def __init__(self, alpha=0.001, beta1=0.9, beta2=0.999):
		super().__init__()
		self.alpha = alpha
		self.beta1 = beta1
		self.beta2 = beta2
		self.name = 'adam'

	def setup(self, keys):
		for k in keys:
			self.memory['M'+k] = None
			self.memory['V'+k] = None
			self.memory['t_m'] = 0
			self.memory['t_v'] = 0

	def update(self, key, gradient):
		m_key = 'M' + key
		v_key = 'V' + key
		super().update(m_key, gradient)
		super().update(v_key, gradient)
		M = (1 - self.beta1) * gradient
		V = (1 - self.beta2) * np.square(gradient)
		self._update_mem(m_key, M, self.beta1)
		self._update_mem(v_key, V, self.beta2)
		M = self._bias_correct(m_key, 't_m', self.beta1)
		V = self._bias_correct(v_key, 't_v', self.beta2)
		return self.alpha / (np.sqrt(V) + 1e-8) * M

	def _bias_correct(self, key, t, beta):
		self.memory[t] += 1
		return self.memory[key] / (1 - np.power(beta, self.memory[t]))
