"""Finite-dimensional linear operators."""

import warnings

import numpy as np
import scipy.sparse.linalg
import scipy.sparse.linalg.interface


class LinearOperator(scipy.sparse.linalg.LinearOperator):
    """Composite base class for finite-dimensional linear operators.

    This class provides a way to define finite-dimensional linear operators without
    explicitly constructing a matrix representation. Instead it suffices to define a
    matrix-vector product and a shape attribute. This avoids unnecessary memory usage
    and can often be more convenient to derive.

    :class:`LinearOperator` instances can be multiplied, added and exponentiated. This
    happens lazily: the result of these operations is a new, composite
    :class:`LinearOperator`, that defers linear operations to the original operators and
    combines the results.

    To construct a concrete class:`LinearOperator`, either pass appropriate callables to
    the constructor of this class, or subclass it.

    A subclass must implement either one of the methods ``_matvec`` and ``_matmat``, and
    the attributes/properties ``shape`` (pair of integers) and ``dtype`` (may be
    ``None``). It may call the ``__init__`` on this class to have these attributes
    validated. Implementing ``_matvec`` automatically implements ``_matmat`` (using a
    naive algorithm) and vice-versa.

    Optionally, a subclass may implement ``_rmatvec`` or ``_adjoint`` to implement the
    Hermitian adjoint (conjugate transpose). As with ``_matvec`` and ``_matmat``,
    implementing either ``_rmatvec`` or ``_adjoint`` implements the other automatically.
    Implementing ``_adjoint`` is preferable; ``_rmatvec`` is mostly there for backwards
    compatibility.

    Parameters
    ----------
    shape : tuple
        Matrix dimensions (M, N).
    matvec : callable f(v)
        Returns :math:`A v`.
    rmatvec : callable f(v)
        Returns :math:`A^H v`, where :math:`A^H` is the conjugate transpose of
        :math:`A`.
    matmat : callable f(V)
        Returns :math:`AV`, where :math:`V` is a dense matrix with dimensions (N, K).
    dtype : dtype
        Data type of the operator.
    rmatmat : callable f(V)
        Returns :math:`A^H V`, where :math:`V` is a dense matrix with dimensions (M, K).

    See Also
    --------
    aslinop : Transform into a LinearOperator.

    Examples
    --------
    >>> import numpy as np
    >>> from probnum.linops import LinearOperator
    >>> def mv(v):
    ...     return np.array([2 * v[0] - v[1], 3 * v[1]])
    ...
    >>> A = LinearOperator(shape=(2, 2), matvec=mv)
    >>> A
    <2x2 _CustomLinearOperator with dtype=float64>
    >>> A.matvec(np.array([1., 2.]))
    array([0., 6.])
    >>> A @ np.ones(2)
    array([1., 3.])
    """

    def __new__(cls, *args, **kwargs):
        if cls is LinearOperator:
            # _CustomLinearOperator factory
            return super().__new__(_CustomLinearOperator)
        else:
            obj = super().__new__(cls)

            if (
                type(obj)._matvec == scipy.sparse.linalg.LinearOperator._matvec
                and type(obj)._matmat == scipy.sparse.linalg.LinearOperator._matmat
            ):
                warnings.warn(
                    "LinearOperator subclass should implement"
                    " at least one of _matvec and _matmat.",
                    category=RuntimeWarning,
                    stacklevel=2,
                )

            return obj

    # Overload arithmetic operators to give access to newly implemented functions (e.g.
    # todense())
    def __rmul__(self, x):
        if np.isscalar(x):
            return _ScaledLinearOperator(self, x)
        else:
            return NotImplemented

    def __pow__(self, p):
        if np.isscalar(p):
            return _PowerLinearOperator(self, p)
        else:
            return NotImplemented

    def __add__(self, x):
        if isinstance(x, LinearOperator):
            return _SumLinearOperator(self, x)
        else:
            return NotImplemented

    def __neg__(self):
        return _ScaledLinearOperator(self, -1)

    # The below methods are overloaded to allow dot products with random variables
    def dot(self, x):
        """Matrix-matrix or matrix-vector multiplication.

        Parameters
        ----------
        x : array_like
            1-d or 2-d array, representing a vector or matrix.

        Returns
        -------
        Ax : array
            1-d or 2-d array (depending on the shape of x) that represents
            the result of applying this linear operator on x.
        """
        if isinstance(x, LinearOperator):
            return _ProductLinearOperator(self, x)
        elif np.isscalar(x):
            return _ScaledLinearOperator(self, x)
        else:
            if len(x.shape) == 1 or len(x.shape) == 2 and x.shape[1] == 1:
                return self.matvec(x)
            elif len(x.shape) == 2:
                return self.matmat(x)
            else:
                raise ValueError(
                    "Expected 1-d or 2-d array, matrix or random variable, got %r." % x
                )

    def matvec(self, x):
        """Matrix-vector multiplication. Performs the operation y=A*x where A is an MxN
        linear operator and x is a 1-d array or random variable.

        Parameters
        ----------
        x : {matrix, ndarray, RandomVariable}
            An array or RandomVariable with shape (N,) or (N,1).
        Returns
        -------
        y : {matrix, ndarray}
            A matrix or ndarray or RandomVariable with shape (M,) or (M,1) depending
            on the type and shape of the x argument.
        Notes
        -----
        This matvec wraps the user-specified matvec routine or overridden
        _matvec method to ensure that y has the correct shape and type.
        """
        M, N = self.shape

        if x.shape != (N,) and x.shape != (N, 1):
            raise ValueError("Dimension mismatch.")

        y = self._matvec(x)

        if isinstance(x, np.matrix):
            y = scipy.sparse.sputils.asmatrix(y)

        if isinstance(x, (np.matrix, np.ndarray)):
            if x.ndim == 1:
                y = y.reshape(M)
            elif x.ndim == 2:
                y = y.reshape(M, 1)
            else:
                raise ValueError("Invalid shape returned by user-defined matvec().")
        # TODO: can be shortened once RandomVariable implements a reshape method
        elif y.shape[0] != M:
            raise ValueError("Invalid shape returned by user-defined matvec().")

        return y

    def transpose(self):
        """Transpose this linear operator.

        Can be abbreviated self.T instead of self.transpose().
        """
        return self._transpose()

    T = property(transpose)

    def _transpose(self):
        """Default implementation of _transpose; defers to rmatvec + conj."""
        return _TransposedLinearOperator(self)

    def todense(self):
        """Dense matrix representation of the linear operator.

        This method can be computationally very costly depending on the shape of the
        linear operator. Use with caution.

        Returns
        -------
        matrix : np.ndarray
            Matrix representation of the linear operator.
        """
        return self.matmat(np.eye(self.shape[1], dtype=self.dtype))

    def inv(self):
        """Inverse of the linear operator."""
        raise NotImplementedError

    # TODO: implement operations (eigs, cond, det, logabsdet, trace, ...)
    def rank(self):
        """Rank of the linear operator."""
        raise NotImplementedError

    def eigvals(self):
        """Eigenvalue spectrum of the linear operator."""
        raise NotImplementedError

    def cond(self, p=None):
        """Compute the condition number of the linear operator.

        The condition number of the linear operator with respect to the ``p`` norm. It
        measures how much the solution :math:`x` of the linear system :math:`Ax=b`
        changes with respect to small changes in :math:`b`.

        Parameters
        ----------
        p : {None, 1, , 2, , inf, 'fro'}, optional
            Order of the norm:

            =======  ============================
            p        norm for matrices
            =======  ============================
            None     2-norm, computed directly via singular value decomposition
            'fro'    Frobenius norm
            np.inf   max(sum(abs(x), axis=1))
            1        max(sum(abs(x), axis=0))
            2        2-norm (largest sing. value)
            =======  ============================

        Returns
        -------
        cond : {float, inf}
            The condition number of the linear operator. May be infinite.
        """
        raise NotImplementedError

    def det(self):
        """Determinant of the linear operator."""
        raise NotImplementedError

    def logabsdet(self):
        """Log absolute determinant of the linear operator."""
        raise NotImplementedError

    def trace(self):
        """Trace of the linear operator.

        Computes the trace of a square linear operator :math:`\\text{tr}(A) =
        \\sum_{i-1}^n A_ii`.

        Returns
        -------
        trace : float
            Trace of the linear operator.

        Raises
        ------
        ValueError : If :meth:`trace` is called on a non-square matrix.
        """
        if self.shape[0] != self.shape[1]:
            raise ValueError("The trace is only defined for square linear operators.")
        else:
            _identity = np.eye(self.shape[0])
            trace = 0.0
            for i in range(self.shape[0]):
                trace += np.squeeze(
                    _identity[np.newaxis, i, :]
                    @ self.matvec(_identity[i, :, np.newaxis])
                )
            return trace


class _CustomLinearOperator(
    scipy.sparse.linalg.interface._CustomLinearOperator, LinearOperator
):
    """Linear operator defined in terms of user-specified operations."""

    def __init__(
        self, shape, matvec, rmatvec=None, matmat=None, rmatmat=None, dtype=None
    ):
        super().__init__(
            shape=shape,
            matvec=matvec,
            rmatvec=rmatvec,
            matmat=matmat,
            rmatmat=rmatmat,
            dtype=dtype,
        )


# TODO: inheritance from _TransposedLinearOperator causes dependency on scipy>=1.4,
# maybe implement our own instead?
class _TransposedLinearOperator(
    scipy.sparse.linalg.interface._TransposedLinearOperator, LinearOperator
):
    """Transposition of a linear operator."""

    def __init__(self, A):
        self.A = A
        super().__init__(A=A)

    def todense(self):
        return self.A.todense().T

    def inv(self):
        return self.A.inv().T


class _SumLinearOperator(
    scipy.sparse.linalg.interface._SumLinearOperator, LinearOperator
):
    """Sum of two linear operators."""

    def __init__(self, A, B):
        self.A = A
        self.B = B
        super().__init__(A=A, B=B)

    def todense(self):
        return self.A.todense() + self.B.todense()

    def inv(self):
        return self.A.inv() + self.B.inv()

    def trace(self):
        return self.A.trace() + self.B.trace()


class _ProductLinearOperator(
    scipy.sparse.linalg.interface._ProductLinearOperator, LinearOperator
):
    """(Operator) Product of two linear operators."""

    def __init__(self, A, B):
        self.A = A
        self.B = B
        super().__init__(A=A, B=B)

    def todense(self):
        return self.A.todense() @ self.B.todense()


class _ScaledLinearOperator(
    scipy.sparse.linalg.interface._ScaledLinearOperator, LinearOperator
):
    """Linear operator scaled with a scalar."""

    def __init__(self, A, alpha):
        super().__init__(A=A, alpha=alpha)

    def todense(self):
        A, alpha = self.args
        return alpha * A.todense()

    def inv(self):
        A, alpha = self.args
        return _ScaledLinearOperator(A.inv(), 1 / alpha)

    def trace(self):
        A, alpha = self.args
        return alpha * A.trace()


class _PowerLinearOperator(
    scipy.sparse.linalg.interface._PowerLinearOperator, LinearOperator
):
    """Linear operator raised to a non-negative integer power."""

    def __init__(self, A, p):
        super().__init__(A=A, p=p)


class Diagonal(LinearOperator):
    """A linear operator representing the diagonal from another linear operator.

    Parameters
    ----------
    Op : LinearOperator
        Linear operator of which to represent the diagonal.
    """

    # TODO: should this be an operator itself or a function of a LinearOperator?
    #   - a function allows subclasses (e.g. MatrixMult) to implement more efficient
    # versions than n products e_i A e_i
    def __init__(self, Op):
        # pylint: disable=super-init-not-called
        raise NotImplementedError


class ScalarMult(LinearOperator):
    """A linear operator representing scalar multiplication.

    Parameters
    ----------
    shape : tuple
        Matrix dimensions (M, N).
    scalar : float
        Scalar to multiply by.
    """

    def __init__(self, shape, scalar):
        self.scalar = scalar
        super().__init__(shape=shape, dtype=float)

    def _matvec(self, x):
        return self.scalar * x

    def _matmat(self, X):
        return self.scalar * X

    def todense(self):
        return np.eye(self.shape[0]) * self.scalar

    def inv(self):
        return ScalarMult(shape=self.shape, scalar=1 / self.scalar)

    # Properties
    def rank(self):
        return np.minimum(self.shape[0], self.shape[1])

    def eigvals(self):
        return np.ones(self.shape[0]) * self.scalar

    def cond(self, p=None):
        return 1

    def det(self):
        return self.scalar ** self.shape[0]

    def logabsdet(self):
        return np.log(np.abs(self.scalar))

    def trace(self):
        return self.scalar * self.shape[0]


class Identity(ScalarMult):
    """The identity operator.

    Parameters
    ----------
    shape : int or tuple
        Shape of the identity operator.
    """

    def __init__(self, shape):
        # Check shape
        if np.isscalar(shape):
            _shape = (shape, shape)
        elif shape[0] != shape[1]:
            raise ValueError("The identity operator must be square.")
        else:
            _shape = shape
        # Initiator of super class
        super().__init__(shape=_shape, scalar=1.0)

    def todense(self):
        return np.eye(self.shape[0])

    def inv(self):
        return self

    # Properties
    def rank(self):
        return self.shape[0]

    def eigvals(self):
        return np.ones(self.shape[0])

    def cond(self, p=None):
        return 1

    def det(self):
        return 1.0

    def logabsdet(self):
        return 0.0

    def trace(self):
        return self.shape[0]


class MatrixMult(scipy.sparse.linalg.interface.MatrixLinearOperator, LinearOperator):
    """A linear operator defined via a matrix.

    Parameters
    ----------
    A : array-like or scipy.sparse.spmatrix
        The explicit matrix.
    """

    def __init__(self, A):
        super().__init__(A=A)

    def _matvec(self, x):
        return self.A @ x  # Needed to call __matmul__ instead of np.dot or np.matmul

    def _matmat(self, X):
        return self.A @ X

    def todense(self):
        if isinstance(self.A, scipy.sparse.spmatrix):
            return self.A.todense()
        else:
            return np.asarray(self.A)

    def inv(self):
        if isinstance(self.A, scipy.sparse.spmatrix):
            invmat = scipy.sparse.linalg.inv(self.A)
        else:
            invmat = np.linalg.inv(self.A)
        return MatrixMult(invmat)

    # Arithmetic operations
    # TODO: perform arithmetic operations between MatrixMult operators explicitly

    # Properties
    def rank(self):
        return np.linalg.matrix_rank(self.A)

    def eigvals(self):
        return np.linalg.eigvals(self.A)

    def cond(self, p=None):
        return np.linalg.cond(self.A, p=p)

    def det(self):
        return np.linalg.det(self.A)

    def logabsdet(self):
        _sign, logdet = np.linalg.slogdet(self.A)
        return logdet

    def trace(self):
        if self.shape[0] != self.shape[1]:
            raise ValueError("The trace is only defined for square linear operators.")
        else:
            return np.trace(self.A)
