"""
Optimizers
"""


class adam_opt:
    """
    An implementation of adam optimizer for variable gradient functions.
    """
    def __init__(self, lr=1e-3, beta1=0.9, beta2=0.999, epsilon=1e-8):
        """
        Initialize for parameters
        """
        self.lr = lr
        self.beta1 = beta1
        self.beta2 = beta2
        self.epsilon = epsilon
        self.m = 0
        self.v = 0
        self.t = 0

    def get_dir(self, grad):
        """
        Take one step to update moments and return descend direction
        """
        self.t += 1
        self.m = self.beta1*self.m + (1 - self.beta1)*grad
        self.v = self.beta2*self.v + (1 - self.beta2)*grad**2
        m_hat = self.m/(1 - self.beta1**self.t)
        v_hat = self.v/(1 - self.beta2**self.t)

        return -self.lr*m_hat/(v_hat**.5 + self.epsilon)
