import unittest

import numpy as np

from neural_net.activation_layers.sigmoid_layer import SigmoidLayer


# noinspection PyMethodMayBeStatic
class SigmoidLayerTests(unittest.TestCase):

    def test_sigmoid_layer_1x1(self):
        ##############
        # Arrange    #
        ##############
        inputs = np.array([[1.0]])
        weights = np.array([[0.5]])
        biases = np.array([0.0])
        learning_rate = 0.001

        # Pre-activation value (z)
        # This is the intermediate value calculated as the weighted sum of inputs plus the bias.
        z = np.dot(inputs, weights) + biases

        # Ouput
        # The result of applying the activation function to the pre-activation value z
        # Sigmoid activation formula: 1 / (1 + e^-z)
        expected_output = 1 / (1 + np.exp(-z))

        # Loss gradient dL/dout
        # Represents how much the loss changes when the output changes.
        dL_dout = np.array([[1.0]])

        # Activation derivative dout/dz
        # This tells you how much the output of the activation function changes with respect to the pre-activation value z.
        # Sigmoid derivative formula: σ(z) * (1 - σ(z))
        dout_dz = expected_output * (1.0 - expected_output)

        # Gradient of the loss with respect to weights (dL/dweights)
        # This represents how much the loss changes when the weights change.
        # Formula: dL/dweights = inputs × dL/dout × σ′(z)
        expected_dl_dweights = inputs * dL_dout * dout_dz
        # Gradient of the loss with respect to the bias (dL/dbias)
        expected_dL_dbias = np.sum(dL_dout * dout_dz)

        # Gradient of the loss with respect to inputs (dL/dinputs)
        # This is the gradient of the loss with respect to the input of the neuron or layer, often needed if you want to backpropagate further.
        # Formula: dL / dinputs = dL/dout × σ′(z) × weights
        expected_dl_dinputs = dL_dout * dout_dz * weights

        # Calculate expected new weights and biases
        expected_weights = weights - learning_rate * expected_dl_dweights
        expected_biases = biases - learning_rate * expected_dL_dbias

        # Initialize SigmoidLayer
        layer = SigmoidLayer(weights.shape[0], weights.shape[1], weights=weights, biases=biases)

        ##############
        # Act        #
        ##############
        # Forward pass
        output = layer.forward(inputs)

        # Backward pass
        dl_dinputs = layer.backward(dL_dout, learning_rate)

        ##############
        # Assert     #
        ##############
        # Forward output correctness
        self.assertTrue(np.allclose(output, expected_output, atol=1e-6),
                        f"Forward output incorrect: Actual: {output}, Expected: {expected_output}")

        # Backward pass correctness
        self.assertTrue(np.allclose(dl_dinputs, expected_dl_dinputs, atol=1e-6),
                        f"Inputs derivative incorrect Actual: {dl_dinputs}, expected: {expected_dl_dinputs}")
        self.assertTrue(np.allclose(layer.weights, expected_weights, atol=1e-6),
                        f"Weight update incorrect Actual: {layer.weights}, expected: {expected_weights}")
        self.assertTrue(np.allclose(layer.biases, expected_biases, atol=1e-6),
                        f"Bias update incorrect Actual: {layer.biases}, expected: {expected_biases}")

    def test_sigmoid_layer_2x2(self):
        ##############
        # Arrange    #
        ##############
        inputs = np.array([[1.0, 2.0],
                           [3.0, 4.0]])

        weights = np.array([[0.5, 0.2],
                            [0.3, 0.7]])

        biases = np.array([0.1, -0.1])
        learning_rate = 0.001

        # Pre-activation value (z)
        # z = inputs.dot(weights) + biases
        z = np.dot(inputs, weights) + biases

        # Expected output using the sigmoid function
        expected_output = 1 / (1 + np.exp(-z))

        # Loss gradient dL/dout (assuming a gradient of 1 for simplicity)
        dL_dout = np.array([[1.0, 1.0],
                            [1.0, 1.0]])

        # Activation derivative dout/dz
        dout_dz = expected_output * (1 - expected_output)

        # Expected gradients
        expected_dl_dweights = np.dot(inputs.T, dL_dout * dout_dz)
        expected_dL_dbias = np.sum(dL_dout * dout_dz, axis=0)
        expected_dl_dinputs = np.dot(dL_dout * dout_dz, weights.T)

        # Expected updated weights and biases
        expected_weights = weights - learning_rate * expected_dl_dweights
        expected_biases = biases - learning_rate * expected_dL_dbias

        # Initialize SigmoidLayer (assuming SigmoidLayer class exists)
        layer = SigmoidLayer(weights.shape[0], weights.shape[1], weights=weights, biases=biases)

        ##############
        # Act        #
        ##############
        # Forward pass
        output = layer.forward(inputs)

        # Backward pass
        dl_dinputs = layer.backward(dL_dout, learning_rate)

        ##############
        # Assert     #
        ##############
        # Forward output correctness
        self.assertTrue(np.allclose(output, expected_output, atol=1e-6),
                        f"Forward output incorrect: Actual: {output}, Expected: {expected_output}")

        # Backward pass correctness
        self.assertTrue(np.allclose(dl_dinputs, expected_dl_dinputs, atol=1e-6),
                        f"Inputs derivative incorrect Actual: {dl_dinputs}, expected: {expected_dl_dinputs}")
        self.assertTrue(np.allclose(layer.weights, expected_weights, atol=1e-6),
                        f"Weight update incorrect Actual: {layer.weights}, expected: {expected_weights}")
        self.assertTrue(np.allclose(layer.biases, expected_biases, atol=1e-6),
                        f"Bias update incorrect Actual: {layer.biases}, expected: {expected_biases}")