from typing import Any, Callable

import numpy as np
from pydantic import BaseModel, ConfigDict, Field, StrictStr, field_validator, model_validator

from recpulse.dtypes import ACTIVATIONS, STR2ACTIVATION, STR2DACTIVATION

PRECISIONS = float | np.float16 | np.float32 | np.float64
TENSOR_TYPE = np.ndarray[PRECISIONS, Any]


class Dense(BaseModel):
    """Dense layer class.

    Layer that connects all inputs to all outputs and adds biases.
    """

    model_config = ConfigDict(validate_assignment=True)

    output_shape: tuple = Field(frozen=True)
    input_shape: tuple | None = None
    name: StrictStr | None = None
    activation: Callable[[Any], TENSOR_TYPE] = Field(frozen=True)
    d_activation: Callable[[Any], TENSOR_TYPE] = Field(frozen=True)
    _weights: np.ndarray | None = None

    @field_validator("output_shape", mode="after")
    @classmethod
    def validate_output_shape(cls, shape: tuple) -> tuple:
        """Output shape validator."""
        if len(shape) == 0:
            raise ValueError("The output shape must contain output neurons.")
        if len(shape) > 1:
            raise ValueError("Dense layer can only have output of shape (n, ).")
        if not isinstance(shape[0], int):
            raise ValueError("Output size must be integer.")
        if shape[0] <= 0:
            raise ValueError("Output size must be positive.")

        return shape

    @field_validator("input_shape", mode="after")
    @classmethod
    def validate_input_shape(cls, shape: tuple) -> tuple | None:
        """Input shape validator."""
        if shape is None:
            return None
        if len(shape) == 0:
            raise ValueError("The output shape must contain output neurons.")
        if len(shape) > 1:
            raise ValueError("Dense layer can only have output of shape (n, ).")
        if not isinstance(shape[0], int):
            raise ValueError("Input size must be integer.")
        if shape[0] <= 0:
            raise ValueError("Input size must be positive.")

        return shape

    @model_validator(mode="after")  # type: ignore
    def validate_model_params(self) -> None:
        """Model validator."""
        if self.input_shape is not None and self._weights is not None:
            if self.input_shape[0] + 1 != self._weights.shape[1]:
                raise ValueError("Input shape does not match weights shape.")

    def __init__(
        self,
        output_shape: tuple | int,
        input_shape: tuple | None = None,
        activation: ACTIVATIONS = "linear",
        name: StrictStr | None = None,
    ) -> None:
        """Class initializer."""
        if isinstance(output_shape, int):
            output_shape = (output_shape,)
        super().__init__(
            output_shape=output_shape,
            input_shape=input_shape,
            name=name,
            activation=STR2ACTIVATION[activation],  # type: ignore
            d_activation=STR2DACTIVATION[activation],  # type: ignore
        )

        if input_shape is not None:
            self.initialize_weights()

    def initialize_weights(
        self,
        input_shape: tuple | None = None,
        mean: float | None = None,
        standard_deviation: float | None = None,
    ) -> None:
        """Weights initializer.

        Creates a numpy array of shape (input_size + 1, self.output_size) with integers
        randomly generated by normal distribution.

        input_size + 1 is used instead of input_size to include biases for this layer.

        Inputs:
            input_size (int): size of the previous neuron layer.
            mean (float): mean of normal distribution for weights initialization.
            standard_deviation (float): standard deviation of normal distribution
                 for weights initialization.
        """

        if self._weights is not None:
            raise AttributeError("Weights are already initialized.")

        if self.input_shape is None:
            if input_shape is None:
                raise ValueError(
                    "input_shape is not specified neither in the function execution, "
                    "nor in the class's instance"
                )
            self.input_shape = input_shape

        if standard_deviation is None:
            standard_deviation = 1
        elif standard_deviation <= 0:
            raise ValueError("standard deviation must be greater than zero")

        if mean is None:
            mean = 0

        self._weights = np.random.normal(
            mean,
            standard_deviation,
            size=(self.output_shape[0], self.input_shape[0] + 1),  # type: ignore
        )

    def propagate(self, inputs: np.ndarray) -> np.ndarray:
        """Propagates input tensor.

        Inputs:
            inputs (np.array): an array of shape (self.input_size, )

        Returns:
            np.array: an array of shape (self.output_size, )
        """

        if self._weights is None:
            raise ValueError("Weights aren't initialized!")
        if self.activation is None:
            raise ValueError("Activation function isn't initialized!")

        # add 1 so that the bias doesn't need any additional code apart from matmul
        modified_input = np.concatenate((inputs, [1]), axis=0)  # type: ignore

        # get new layer
        propagated = np.dot(self._weights, modified_input)

        return self.activation(propagated)

    @property
    def initialized(self) -> bool:
        """Return True if weights are initialized, False otherwise."""
        return self._weights is not None

    def back_propagate(
        self, error: np.ndarray, inputs: np.ndarray, learning_rate: float = 0.001, tune: bool = True
    ) -> float:
        """Back propagation."""
        if self._weights is None:
            raise ValueError("Weights aren't initialized!")

        modified_input = np.concatenate((inputs, [1]), axis=0)  # type: ignore
        pred = np.dot(self._weights, modified_input)
        propagated_error = np.dot(np.dot(error, self._weights)[:-1], self.d_activation(inputs))

        if tune:
            dW = np.tensordot(np.dot(error, self.d_activation(pred)), modified_input, axes=0)
            self._weights += learning_rate * dW

        return propagated_error
