import numpy as np


class Optimizer:
    """Base optimizer class"""

    def __init__(self, parameters, lr=0.01):
        self.parameters = parameters
        self.lr = lr

    def step(self):
        raise NotImplementedError

    def zero_grad(self):

        for layer, param_name, grad_name in self.parameters:
            grad = getattr(layer, grad_name)
            grad.fill(0)



# 1. Stochastic Gradient Descent

class SGD(Optimizer):

    def step(self):

        for layer, param_name, grad_name in self.parameters:

            param = getattr(layer, param_name)
            grad = getattr(layer, grad_name)

            param -= self.lr * grad



# 2. Momentum

class Momentum(Optimizer):

    def __init__(self, parameters, lr=0.01, beta=0.9):
        super().__init__(parameters, lr)
        self.beta = beta

        self.v = []
        for layer, param_name, _ in parameters:
            param = getattr(layer, param_name)
            self.v.append(np.zeros_like(param))

    def step(self):

        for i, (layer, param_name, grad_name) in enumerate(self.parameters):

            param = getattr(layer, param_name)
            grad = getattr(layer, grad_name)

            self.v[i] = self.beta * self.v[i] + (1 - self.beta) * grad

            param -= self.lr * self.v[i]



# 3. Nesterov Accelerated Gradient

class Nesterov(Optimizer):

    def __init__(self, parameters, lr=0.01, beta=0.9):
        super().__init__(parameters, lr)
        self.beta = beta

        self.v = []
        for layer, param_name, _ in parameters:
            param = getattr(layer, param_name)
            self.v.append(np.zeros_like(param))

    def step(self):

        for i, (layer, param_name, grad_name) in enumerate(self.parameters):

            param = getattr(layer, param_name)
            grad = getattr(layer, grad_name)

            v_prev = self.v[i]

            self.v[i] = self.beta * self.v[i] + self.lr * grad

            param -= (-self.beta * v_prev + (1 + self.beta) * self.v[i])



# 4. AdaGrad

class AdaGrad(Optimizer):

    def __init__(self, parameters, lr=0.01, eps=1e-8):
        super().__init__(parameters, lr)
        self.eps = eps

        self.G = []
        for layer, param_name, _ in parameters:
            param = getattr(layer, param_name)
            self.G.append(np.zeros_like(param))

    def step(self):

        for i, (layer, param_name, grad_name) in enumerate(self.parameters):

            param = getattr(layer, param_name)
            grad = getattr(layer, grad_name)

            self.G[i] += grad ** 2

            param -= self.lr * grad / (np.sqrt(self.G[i]) + self.eps)



# 5. RMSProp

class RMSProp(Optimizer):

    def __init__(self, parameters, lr=0.001, beta=0.9, eps=1e-8):
        super().__init__(parameters, lr)
        self.beta = beta
        self.eps = eps

        self.s = []
        for layer, param_name, _ in parameters:
            param = getattr(layer, param_name)
            self.s.append(np.zeros_like(param))

    def step(self):

        for i, (layer, param_name, grad_name) in enumerate(self.parameters):

            param = getattr(layer, param_name)
            grad = getattr(layer, grad_name)

            self.s[i] = self.beta * self.s[i] + (1 - self.beta) * grad ** 2

            param -= self.lr * grad / (np.sqrt(self.s[i]) + self.eps)



# 6. Adam

class Adam(Optimizer):

    def __init__(self, parameters, lr=0.001, beta1=0.9, beta2=0.999, eps=1e-8):

        super().__init__(parameters, lr)

        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps

        self.m = []
        self.v = []

        for layer, param_name, _ in parameters:
            param = getattr(layer, param_name)
            self.m.append(np.zeros_like(param))
            self.v.append(np.zeros_like(param))

        self.t = 0

    def step(self):

        self.t += 1

        for i, (layer, param_name, grad_name) in enumerate(self.parameters):

            param = getattr(layer, param_name)
            grad = getattr(layer, grad_name)

            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * grad
            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * (grad ** 2)

            m_hat = self.m[i] / (1 - self.beta1 ** self.t)
            v_hat = self.v[i] / (1 - self.beta2 ** self.t)

            param -= self.lr * m_hat / (np.sqrt(v_hat) + self.eps)