# This code is part of isQ.
# (C) Copyright ArcLight Quantum 2023.
# This code is licensed under the MIT License.

"""
A Neural Network abstract class for all.
This part code is inspired by IBM qiskit.
For more information, check out:https://qiskit.org/
"""

from __future__ import annotations

from abc import ABC, abstractmethod

import numpy as np

try:
    from sparse import SparseArray
except Exception:

    class SparseArray:
        """Empty SparseArray class
        Replacement if sparse.SparseArray is not present.
        """


class NeuralNetworkError(Exception):
    """Neural network error."""


class NeuralNetwork(ABC):
    """Abstract Neural Network class providing forward and backward pass and
    handling batched inputs. This is to be implemented by other (quantum)
    neural networks.
    """

    def __init__(
        self,
        num_inputs: int,
        num_weights: int,
        sparse: bool,
        output_shape: int | tuple[int, ...],
        input_gradients: bool = True,
    ) -> None:
        """Initialize quantum neural network.

        Args:
            num_inputs: The number of input features.
            num_weights: The number of trainable weights.
            sparse: Determines whether the output is a sparse array or not.
            output_shape: The shape of the output.
            input_gradients: Determines whether to compute gradients with
                             respect to input data.

        Raises:
            NeuralNetworkError: Invalid parameter values.

        """
        if num_inputs < 0:
            raise NeuralNetworkError(
                f"Number of inputs cannot be negative: {num_inputs}!",
            )

        self._num_inputs = num_inputs

        if num_weights < 0:
            raise NeuralNetworkError(
                f"Number of weights cannot be negative: {num_weights}!"
            )

        self._num_weights = num_weights

        self._sparse = sparse

        # output shape may be derived later, so check it only if it is not None
        if output_shape is not None:
            self._output_shape = self._validate_output_shape(output_shape)

        self._input_gradients = input_gradients

    @property
    def num_inputs(self) -> int:
        """Returns the number of input features."""
        return self._num_inputs

    @property
    def num_weights(self) -> int:
        """Returns the number of trainable weights."""
        return self._num_weights

    @property
    def sparse(self) -> bool:
        """Returns whether the output is sparse or not."""
        return self._sparse

    @property
    def output_shape(self) -> tuple[int, ...]:
        """Returns the output shape."""
        return self._output_shape

    @property
    def input_gradients(self) -> bool:
        """Returns whether gradients with respect to input data are computed by
        this neural network in the ``backward`` method or not. By default such
        gradients are not computed."""
        return self._input_gradients

    @input_gradients.setter
    def input_gradients(self, input_gradients: bool) -> None:
        """Turn on/off computation of gradients with respect to input data."""
        self._input_gradients = input_gradients

    def _validate_output_shape(self, output_shape):
        if isinstance(output_shape, int):
            output_shape = (output_shape,)
        if not np.all([s > 0 for s in output_shape]):
            raise NeuralNetworkError(
                f"Invalid output shape, all components must be > 0, "
                f"but got: {output_shape}."
            )
        return output_shape

    def _validate_input(
        self, input_data: float | list[float] | np.ndarray | None
    ) -> tuple[np.ndarray | None, tuple[int, ...] | None]:
        if input_data is None:
            return None, None
        input_ = np.array(input_data)
        shape = input_.shape
        if len(shape) == 0:
            # there's a single value in the input.
            input_ = input_.reshape((1, 1))
            return input_, shape

        if shape[-1] != self._num_inputs:
            raise NeuralNetworkError(
                f"Input data has incorrect shape, last dimension "
                f"is not equal to the number of inputs: "
                f"{self._num_inputs}, but got: {shape[-1]}."
            )

        if len(shape) == 1:
            # add an empty dimension for samples (batch dimension)
            input_ = input_.reshape((1, -1))
        elif len(shape) > 2:
            # flatten lower dimensions, keep num_inputs as a last dimension
            input_ = input_.reshape((np.product(input_.shape[:-1]), -1))

        return input_, shape

    def _preprocess_forward(
        self,
        input_data: np.ndarray | None,
        weights: np.ndarray | None,
    ) -> tuple[np.ndarray | None, int | None]:
        """Pre-processing during forward pass of the network for the
        primitive-based networks.
        """
        if input_data is not None:
            num_samples = input_data.shape[0]
            if weights is not None:
                weights = np.broadcast_to(weights, (num_samples, len(weights)))
                parameters = np.concatenate((input_data, weights), axis=1)
            else:
                parameters = input_data
        else:
            if weights is not None:
                num_samples = 1
                parameters = np.broadcast_to(
                    weights,
                    (
                        num_samples,
                        len(weights),
                    ),
                )
            else:
                return None, None
        return parameters, num_samples

    def _validate_weights(
        self, weights: float | list[float] | np.ndarray | None
    ) -> np.ndarray | None:
        if weights is None:
            return None
        weights_ = np.array(weights)
        return weights_.reshape(self._num_weights)

    def _validate_forward_output(
        self, output_data: np.ndarray, original_shape: tuple[int, ...]
    ) -> np.ndarray:
        if original_shape and len(original_shape) >= 2:
            output_data = output_data.reshape(
                (*original_shape[:-1], *self._output_shape)
            )

        return output_data

    def _validate_backward_output(
        self,
        input_grad: np.ndarray,
        weight_grad: np.ndarray,
        original_shape: tuple[int, ...],
    ) -> tuple[np.ndarray | SparseArray, np.ndarray | SparseArray]:
        if input_grad is not None and np.prod(input_grad.shape) == 0:
            input_grad = None

        if input_grad is not None and original_shape and len(original_shape) >= 2:
            input_grad = input_grad.reshape(
                (*original_shape[:-1], *self._output_shape, self._num_inputs)
            )

        if weight_grad is not None and np.prod(weight_grad.shape) == 0:
            weight_grad = None

        if weight_grad is not None and original_shape and len(original_shape) >= 2:
            weight_grad = weight_grad.reshape(
                (*original_shape[:-1], *self._output_shape, self._num_weights)
            )

        return input_grad, weight_grad

    def forward(
        self,
        input_data: float | list[float] | np.ndarray | None,
        weights: float | list[float] | np.ndarray | None,
    ) -> np.ndarray | SparseArray:
        """Forward pass of the network.

        Args:
            input_data: input data of the shape (num_inputs). In case of a
                single scalar input it is directly cast to and interpreted like
                a one-element array.
            weights: trainable weights of the shape (num_weights). In case of a
                single scalar weightit is directly cast to and interpreted like
                a one-element array.

        Returns:
            The result of the neural network of the shape (output_shape).

        """
        input_, shape = self._validate_input(input_data)
        weights_ = self._validate_weights(weights)
        output_data = self._forward(input_, weights_)
        return self._validate_forward_output(output_data, shape)

    @abstractmethod
    def _forward(
        self, input_data: np.ndarray | None, weights: np.ndarray | None
    ) -> np.ndarray | SparseArray:
        """Forward calculation."""

    def backward(
        self,
        input_data: float | list[float] | np.ndarray | None,
        weights: float | list[float] | np.ndarray | None,
    ) -> tuple[np.ndarray | SparseArray | None, np.ndarray | SparseArray | None]:
        input_, shape = self._validate_input(input_data)
        weights_ = self._validate_weights(weights)
        input_grad, weight_grad = self._backward(input_, weights_)
        (
            input_grad_reshaped,
            weight_grad_reshaped,
        ) = self._validate_backward_output(input_grad, weight_grad, shape)

        return input_grad_reshaped, weight_grad_reshaped

    @abstractmethod
    def _backward(
        self, input_data: np.ndarray | None, weights: np.ndarray | None
    ) -> tuple[np.ndarray | SparseArray | None, np.ndarray | SparseArray | None]:
        """Backward calculation."""
