__all__ = ["Diagonal"]

from typing import Optional, Union

import numpy as np

from pylops import LinearOperator
from pylops.utils._internal import _value_or_sized_to_tuple
from pylops.utils.backend import get_array_module, to_cupy_conditional
from pylops.utils.decorators import reshaped
from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray


class Diagonal(LinearOperator):
    r"""Diagonal operator.

    Applies element-wise multiplication of the input vector with the vector
    ``diag`` in forward and with its complex conjugate in adjoint mode.

    This operator can also broadcast; in this case the input vector is
    reshaped into its dimensions ``dims`` and the element-wise multiplication
    with ``diag`` is perfomed along ``axis``. Note that the
    vector ``diag`` will need to have size equal to ``dims[axis]``.

    Parameters
    ----------
    diag : :obj:`numpy.ndarray`
        Vector to be used for element-wise multiplication.
    dims : :obj:`list`, optional
        Number of samples for each dimension
        (``None`` if only one dimension is available)
    axis : :obj:`int`, optional
        .. versionadded:: 2.0.0

        Axis along which multiplication is applied.
    dtype : :obj:`str`, optional
        Type of elements in input array.
    name : :obj:`str`, optional
        .. versionadded:: 2.0.0

        Name of operator (to be used by :func:`pylops.utils.describe.describe`)

    Attributes
    ----------
    shape : :obj:`tuple`
        Operator shape
    explicit : :obj:`bool`
        Operator contains a matrix that can be solved explicitly (``True``) or
        not (``False``)

    Notes
    -----
    Element-wise multiplication between the model :math:`\mathbf{x}` and/or
    data :math:`\mathbf{y}` vectors and the array :math:`\mathbf{d}`
    can be expressed as

    .. math::

        y_i = d_i x_i  \quad \forall i=1,2,\ldots,N

    This is equivalent to a matrix-vector multiplication with a matrix
    containing the vector :math:`\mathbf{d}` along its main diagonal.

    For real-valued ``diag``, the Diagonal operator is self-adjoint as the
    adjoint of a diagonal matrix is the diagonal matrix itself. For
    complex-valued ``diag``, the adjoint is equivalent to the element-wise
    multiplication with the complex conjugate elements of ``diag``.

    """

    def __init__(
        self,
        diag: NDArray,
        dims: Optional[Union[int, InputDimsLike]] = None,
        axis: int = -1,
        dtype: DTypeLike = "float64",
        name: str = "D",
    ) -> None:
        self.diag = diag
        self.axis = axis
        origdims = dims
        dims = self.diag.shape if dims is None else _value_or_sized_to_tuple(dims)
        super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dims, name=name)

        ncp = get_array_module(diag)
        self.complex = True if ncp.iscomplexobj(self.diag) else False
        if origdims is not None:
            diagdims = np.ones_like(self.dims)
            diagdims[axis] = self.dims[axis]
            self.diag = self.diag.reshape(diagdims)

    @reshaped
    def _matvec(self, x: NDArray) -> NDArray:
        if type(self.diag) is not type(x):
            self.diag = to_cupy_conditional(x, self.diag)
        y = self.diag * x
        return y

    @reshaped
    def _rmatvec(self, x: NDArray) -> NDArray:
        if type(self.diag) is not type(x):
            self.diag = to_cupy_conditional(x, self.diag)
        if self.complex:
            diagadj = self.diag.conj()
        else:
            diagadj = self.diag
        y = diagadj * x
        return y

    def matrix(self) -> NDArray:
        """Return diagonal matrix as dense :obj:`numpy.ndarray`

        Returns
        -------
        densemat : :obj:`numpy.ndarray`
            Dense matrix.

        """
        ncp = get_array_module(self.diag)
        densemat = ncp.diag(self.diag.squeeze())
        return densemat

    def todense(self) -> NDArray:
        """Fast implementation of todense based on known structure of the
        operator

        Returns
        -------
        densemat : :obj:`numpy.ndarray`
            Dense matrix.
        """
        if self.diag.ndim == 1:
            return self.matrix()
        else:
            dims = list(self.dims)
            dims[self.axis] = 1
            matrix = np.diag(np.tile(self.diag, dims).ravel())
            return matrix
