import numpy as np
from activations import ActivationFunc, Identity


class NNLayer:
    def __init__(self,
                 in_size: int,
                 out_size: int,
                 lr: float,
                 activation_func: ActivationFunc):
        limit = np.sqrt(6 / (in_size + out_size))
        self.W = np.random.uniform(-limit, limit, (in_size, out_size))
        self.B = np.zeros((out_size))
        self.lr = lr
        self.input = None
        self.output = None
        self.output_linear = None
        self.activation_func = activation_func

    def __str__(self):
        return f'[ {self.W.shape[0]} => {self.W.shape[1]}\tlr:{self.lr}\tactivation:{self.activation_func.__class__.__name__} ]' # noqa

    def forward(self, v: np.ndarray) -> np.ndarray:
        self.input = v
        self.output_linear = self.input @ self.W + self.B
        self.output = self.activation_func(
                self.output_linear
            )
        return self.output

    def backprop(self, error: np.ndarray) -> np.ndarray:
        error *= self.activation_func.d(self.output_linear)
        ret = self.W @ error
        dW = np.outer(self.input, error) * self.lr
        dB = error * self.lr
        self.W -= dW
        self.B -= dB
        return ret


class SampleLayer:
    def __init__(self,
                 in_size: int,
                 lr: float,
                 activation_func: ActivationFunc):
        self.input = None
        self.mean_nn = NNLayer(
            in_size,
            in_size,
            lr,
            activation_func)
        self.std_nn = NNLayer(
            in_size,
            in_size,
            lr,
            activation_func)

    def DKL(self):
        return -0.5 * np.mean(1 + self.logvar - self.mean ** 2 - np.exp(self.logvar)) # noqa

    def forward(self, v: np.ndarray) -> np.ndarray:
        self.input = v
        self.mean = self.mean_nn.forward(v)
        self.logvar = np.clip(self.std_nn.forward(v), -10, 10)
        self.std = np.exp(0.5 * self.logvar)
        self.eps = np.random.normal(0, 1, self.mean.shape)
        return 0.5 * self.eps * self.std + self.mean

    def backprop(self, error: np.ndarray) -> np.ndarray:
        dmean = error + self.mean
        dstd = error * self.eps + 0.5 * (np.exp(self.logvar) - 1)
        mean_error = self.mean_nn.backprop(dmean)
        logvar_error = self.std_nn.backprop(dstd * self.std)
        return mean_error + logvar_error


class DeepNNLayer:
    def __init__(self,
                 layers: list[int],
                 lr: float,
                 activation_func: ActivationFunc):
        self.layers: list[NNLayer] = []
        for i in range(len(layers) - 1):
            self.layers.append(
                    NNLayer(
                        layers[i],
                        layers[i+1],
                        lr,
                        activation_func if i != len(layers) - 2 else Identity()
                    )
                )
        self.in_size = layers[0]
        self.out_size = layers[-1]

    def __str__(self):
        return '\n'.join([str(layer) for layer in self.layers])

    def forward(self, v: np.ndarray) -> np.ndarray:
        for layer in self.layers:
            v = layer.forward(v)
        return v

    def backprop(self, error: np.ndarray) -> np.ndarray:
        for layer in self.layers[::-1]:
            error = layer.backprop(error)
        return error