import numpy as np


class Layer:
    """Base class for layers"""

    def forward(self, x):
        raise NotImplementedError

    def backward(self, grad):
        raise NotImplementedError

    def parameters(self):
        return []



# 1. Linear Layer

class Linear(Layer):

    def __init__(self, in_features, out_features):

        self.W = np.random.randn(in_features, out_features) * np.sqrt(2 / in_features)
        self.b = np.zeros((1, out_features))

        # Initialize gradients
        self.dW = np.zeros_like(self.W)
        self.db = np.zeros_like(self.b)

    def forward(self, x):

        self.x = x
        return x @ self.W + self.b

    def backward(self, grad):

        self.dW = self.x.T @ grad
        self.db = np.sum(grad, axis=0, keepdims=True)

        dx = grad @ self.W.T

        return dx


# 2. Dropout

class Dropout(Layer):

    def __init__(self, p=0.5):
        self.p = p
        self.training = True

    def forward(self, x):

        if not self.training:
            return x

        self.mask = (np.random.rand(*x.shape) > self.p) / (1 - self.p)

        return x * self.mask

    def backward(self, grad):

        if not self.training:
            return grad

        return grad * self.mask



# 3. Softmax

class Softmax(Layer):

    def forward(self, x):

        x_shifted = x - np.max(x, axis=1, keepdims=True)
        exp_x = np.exp(x_shifted)

        self.out = exp_x / np.sum(exp_x, axis=1, keepdims=True)

        return self.out

    def backward(self, grad):

        batch_size, n_classes = grad.shape
        dx = np.zeros_like(grad)

        for i in range(batch_size):

            y = self.out[i].reshape(-1, 1)

            jacobian = np.diagflat(y) - y @ y.T

            dx[i] = jacobian @ grad[i]

        return dx



# 4. Batch Normalization

class BatchNorm(Layer):

    def __init__(self, num_features, eps=1e-5, momentum=0.9):

        self.gamma = np.ones((1, num_features))
        self.beta = np.zeros((1, num_features))

        self.dgamma = np.zeros_like(self.gamma)
        self.dbeta = np.zeros_like(self.beta)

        self.eps = eps
        self.momentum = momentum

        self.running_mean = np.zeros((1, num_features))
        self.running_var = np.ones((1, num_features))

        self.training = True

    def forward(self, x):

        if self.training:

            self.x = x

            self.mean = np.mean(x, axis=0, keepdims=True)
            self.var = np.var(x, axis=0, keepdims=True)

            self.x_hat = (x - self.mean) / np.sqrt(self.var + self.eps)

            out = self.gamma * self.x_hat + self.beta

            self.running_mean = self.momentum * self.running_mean + (1 - self.momentum) * self.mean
            self.running_var = self.momentum * self.running_var + (1 - self.momentum) * self.var

            return out

        else:

            x_hat = (x - self.running_mean) / np.sqrt(self.running_var + self.eps)

            return self.gamma * x_hat + self.beta

    def backward(self, grad):

        m = grad.shape[0]

        dbeta = np.sum(grad, axis=0, keepdims=True)
        dgamma = np.sum(grad * self.x_hat, axis=0, keepdims=True)

        dx_hat = grad * self.gamma

        dvar = np.sum(dx_hat * (self.x - self.mean) * -0.5 * (self.var + self.eps)**(-3/2), axis=0, keepdims=True)

        dmean = np.sum(dx_hat * -1 / np.sqrt(self.var + self.eps), axis=0, keepdims=True) + \
                dvar * np.mean(-2 * (self.x - self.mean), axis=0, keepdims=True)

        dx = dx_hat / np.sqrt(self.var + self.eps) + \
             dvar * 2 * (self.x - self.mean) / m + \
             dmean / m

        self.dgamma = dgamma
        self.dbeta = dbeta

        return dx