# -*- coding: utf-8 -*-

# This code is part of Qiskit.
#
# (C) Copyright IBM 2020.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.
# pylint: disable=invalid-name

"""Utilities for type handling/conversion, primarily dealing with
reshaping arrays, and handling qiskit types that wrap arrays.
"""

from typing import Union, List
from collections.abc import Iterable
import numpy as np
from scipy.sparse import csr_matrix, issparse, spmatrix
from scipy.sparse import kron as sparse_kron
from scipy.sparse import identity as sparse_identity

from qiskit.quantum_info.operators import Operator
from qiskit_dynamics.array import Array
from qiskit_dynamics.dispatch import requires_backend

try:
    from jax.experimental import sparse as jsparse
except ImportError:
    pass


class StateTypeConverter:
    """Contains descriptions of two type specifications for DE solvers/methods,
    with functions for converting states and rhs functions between
    representations.

    A type specification is a `dict` describing a specific expected type,
    e.g. an array of a given
    shape. Currently only handled types `Array`s, specified via:
        - {'type': 'array', 'shape': tuple}

    While this class stores exact type specifications, it can be
    instantiated with a concrete type and a more general type.
    This facilitates the situation in which a solver requires a 1d `Array`,
    which is specified by the type:
        - {'type': 'array', 'ndim': 1}
    """

    def __init__(self, inner_type_spec, outer_type_spec=None, order="F"):
        """Instantiate with the inner and return types for the state.

        Args:
            inner_type_spec (dict): inner type
            outer_type_spec (dict): outer type
            order (str): order argument to be used in array reshaping.
                        Defaults to `'F'`, which corresponds to column
                        stacking convention when switching between 2d and 1d
                        arrays.
        """

        self.inner_type_spec = inner_type_spec
        self.outer_type_spec = self.inner_type_spec if outer_type_spec is None else outer_type_spec
        self.order = order

    @classmethod
    def from_instances(cls, inner_y, outer_y=None, order="F"):
        """Instantiate from concrete instances. Type of instances must
        be supported by `type_spec_from_instance`. If `outer_y is None` the
        outer type is set to the inner type.

        Args:
            inner_y (array): concrete representative of inner type
            outer_y (array): concrete representative of outer type
            order (str): order argument to be used in array reshaping.

        Returns:
            StateTypeConverter: type converter as specified by args
        """
        inner_type_spec = type_spec_from_instance(inner_y)

        outer_type_spec = None
        if outer_y is not None:
            outer_type_spec = type_spec_from_instance(outer_y)

        return cls(inner_type_spec, outer_type_spec, order)

    @classmethod
    def from_outer_instance_inner_type_spec(cls, outer_y, inner_type_spec=None, order="F"):
        """Instantiate from concrete instance of the outer type,
        and an inner type-spec. The inner type spec can be either
        be fully specified, or be more general (i.e. to
        facilitate the situation in which a solver needs a 1d array).

        Accepted general data types:
            - {'type': 'array'}
            - {'type': 'array', 'ndim': 1}

        Args:
            outer_y (array): concrete outer data type
            inner_type_spec (dict): inner, potentially general, type spec
            order (str): order argument to be used in array reshaping.

        Returns:
            StateTypeConverter: type converter as specified by args

        Raises:
            Exception: if inner_type_spec is not properly specified or is
            not a handled type
        """

        # if no inner_type_spec given just instantiate both inner
        # and outer to the outer_y
        if inner_type_spec is None:
            return cls.from_instances(outer_y, order=order)

        inner_type = inner_type_spec.get("type")
        if inner_type is None:
            raise Exception("inner_type_spec needs a 'type' key.")

        if inner_type == "array":
            outer_y_as_array = Array(outer_y)

            # if a specific shape is given attempt to instantiate from a
            # reshaped outer_y
            shape = inner_type_spec.get("shape")
            if shape is not None:
                return cls.from_instances(
                    outer_y_as_array.reshape(shape, order=order), outer_y, order=order
                )

            # handle the case that ndim == 1 is given
            ndim = inner_type_spec.get("ndim")
            if ndim == 1:
                return cls.from_instances(
                    outer_y_as_array.flatten(order=order), outer_y, order=order
                )

            # if neither shape nor ndim is given, assume it can be an array
            # of any shape
            return cls.from_instances(outer_y_as_array, outer_y, order=order)

        raise Exception("inner_type_spec not a handled type.")

    def inner_to_outer(self, y):
        """Convert a state of inner type to one of outer type."""
        return convert_state(y, self.outer_type_spec, self.order)

    def outer_to_inner(self, y):
        """Convert a state of outer type to one of inner type."""
        return convert_state(y, self.inner_type_spec, self.order)

    def rhs_outer_to_inner(self, rhs):
        """Convert an rhs function f(t, y) to work on inner type."""

        # if inner and outer type specs are the same do nothing
        if self.inner_type_spec == self.outer_type_spec:
            return rhs

        def new_rhs(t, y):
            outer_y = self.inner_to_outer(y)
            rhs_val = rhs(t, outer_y)
            return self.outer_to_inner(rhs_val)

        return new_rhs

    def generator_outer_to_inner(self, generator):
        """Convert generator from outer to inner type."""
        # if inner and outer type specs are the same do nothing
        if self.inner_type_spec == self.outer_type_spec:
            return generator

        # raise exceptions based on assumptions
        if (self.inner_type_spec["type"] != "array") or (self.outer_type_spec["type"] != "array"):
            raise Exception(
                """RHS generator transformation only
                               valid for state types Array."""
            )
        if len(self.inner_type_spec["shape"]) != 1:
            raise Exception(
                """RHS generator transformation only valid
                               if inner_type is 1d."""
            )

        if self.order == "C":
            # create identity of size the second dimension of the
            # outer type
            ident = Array(np.eye(self.outer_type_spec["shape"][1]))

            def new_generator(t):
                return Array(np.kron(generator(t), ident))

            return new_generator
        elif self.order == "F":
            # create identity of size the second dimension of the
            # outer type
            ident = Array(np.eye(self.outer_type_spec["shape"][1]))

            def new_generator(t):
                return Array(np.kron(ident, generator(t)))

            return new_generator
        else:
            raise Exception("""Unsupported order for generator conversion.""")


def convert_state(y: Array, type_spec: dict, order="F"):
    """Convert the de state y into the type specified by type_spec.
    Accepted values of type_spec are given at the beginning of the file.

    Args:
        y: the state to convert.
        type_spec (dict): the type description to convert to.
        order (str): order argument for any array reshaping function.

    Returns:
        Array: converted state.
    """

    new_y = None

    if type_spec["type"] == "array":
        # default array data type to complex
        new_y = Array(y, dtype=type_spec.get("dtype", "complex"))

        shape = type_spec.get("shape")
        if shape is not None:
            new_y = new_y.reshape(shape, order=order)

    return new_y


def type_spec_from_instance(y):
    """Determine type spec from an instance."""
    type_spec = {}
    if isinstance(y, (Array, np.ndarray)):
        type_spec["type"] = "array"
        type_spec["shape"] = y.shape

    return type_spec


def vec_commutator(
    A: Union[Array, csr_matrix, List[csr_matrix]]
) -> Union[Array, csr_matrix, List[csr_matrix]]:
    r"""Linear algebraic vectorization of the linear map X -> -i[A, X]
    in column-stacking convention. In column-stacking convention we have

    .. math::
        vec(ABC) = C^T \otimes A vec(B),

    so for the commutator we have

    .. math::
        -i[A, \cdot] = -i(A \cdot - \cdot A \mapsto id \otimes A - A^T \otimes id)

    Note: this function is also "vectorized" in the programming sense for dense arrays.

    Args:
        A: Either a 2d array representing the matrix A described above,
           a 3d array representing a list of matrices, a sparse matrix, or a
           list of sparse matrices.

    Returns:
        Array: vectorized version of the map.
    """

    if issparse(A):
        # single, sparse matrix
        sp_iden = sparse_identity(A.shape[-1], format="csr")
        return -1j * (sparse_kron(sp_iden, A) - sparse_kron(A.T, sp_iden))
    if isinstance(A, list) and issparse(A[0]):
        # taken to be 1d array of 2d sparse matrices
        sp_iden = sparse_identity(A[0].shape[-1], format="csr")
        out = [-1j * (sparse_kron(sp_iden, mat) - sparse_kron(mat.T, sp_iden)) for mat in A]
        return out

    A = to_array(A)
    iden = Array(np.eye(A.shape[-1]))
    axes = list(range(A.ndim))
    axes[-1] = axes[-2]
    axes[-2] += 1
    return -1j * (np.kron(iden, A) - np.kron(A.transpose(axes), iden))


def vec_dissipator(
    L: Union[Array, csr_matrix, List[csr_matrix]]
) -> Union[Array, csr_matrix, List[csr_matrix]]:
    r"""Linear algebraic vectorization of the linear map
    X -> L X L^\dagger - 0.5 * (L^\dagger L X + X L^\dagger L)
    in column stacking convention.

    This gives

    .. math::
        \overline{L} \otimes L - 0.5(id \otimes L^\dagger L +
            (L^\dagger L)^T \otimes id)

    Note: this function is also "vectorized" in the programming sense for dense matrices.
    """

    if issparse(L):
        sp_iden = sparse_identity(L[0].shape[-1], format="csr")
        return sparse_kron(L.conj(), L) - 0.5 * (
            sparse_kron(sp_iden, L.conj().T * L) + sparse_kron(L.T * L.conj(), sp_iden)
        )
    if isinstance(L, list) and issparse(L[0]):
        # taken to be 1d array of 2d sparse matrices
        sp_iden = sparse_identity(L[0].shape[-1], format="csr")
        out = [
            sparse_kron(mat.conj(), mat)
            - 0.5
            * (sparse_kron(sp_iden, mat.conj().T * mat) + sparse_kron(mat.T * mat.conj(), sp_iden))
            for mat in L
        ]
        return out

    iden = Array(np.eye(L.shape[-1]))
    axes = list(range(L.ndim))

    L = to_array(L)
    axes[-1] = axes[-2]
    axes[-2] += 1
    Lconj = L.conj()
    LdagL = Lconj.transpose(axes) @ L
    LdagLtrans = LdagL.transpose(axes)

    return np.kron(Lconj, iden) @ np.kron(iden, L) - 0.5 * (
        np.kron(iden, LdagL) + np.kron(LdagLtrans, iden)
    )


def isinstance_qutip_qobj(obj):
    """Check if the object is a qutip Qobj.

    Args:
        obj (any): Any object for testing.

    Returns:
        Bool: True if obj is qutip Qobj
    """
    if (
        type(obj).__name__ == "Qobj"
        and hasattr(obj, "_data")
        and type(obj._data).__name__ == "fast_csr_matrix"
    ):
        return True
    return False


# pylint: disable=too-many-return-statements
def to_array(op: Union[Operator, Array, List[Operator], List[Array], spmatrix], no_iter=False):
    """Convert an operator or list of operators to an Array.
    Args:
        op: Either an Operator to be converted to an array, a list of Operators
            to be converted to a 3d array, or an array (which simply gets
            returned)
        no_iter (Bool): Boolean determining whether to recursively unroll `Iterables`.
            If recurring, this should be True to avoid making each element of the
            input array into a separate Array.
    Returns:
        Array: Array version of input
    """
    if op is None:
        return op

    if isinstance(op, np.ndarray) and op.dtype != "O":
        if Array.default_backend() in [None, "numpy"]:
            return op
        else:
            return Array(op)

    if isinstance(op, Array):
        return op

    if issparse(op):
        return Array(op.toarray())

    if type(op).__name__ == "BCOO":
        return Array(op.todense())

    if isinstance(op, Iterable) and not no_iter:
        op = Array([to_array(sub_op, no_iter=True) for sub_op in op])
    elif isinstance(op, Iterable) and no_iter:
        return op
    else:
        op = Array(op)

    if op.backend == "numpy":
        return op.data
    else:
        return op


# pylint: disable=too-many-return-statements
def to_csr(
    op: Union[Operator, Array, List[Operator], List[Array], spmatrix], no_iter=False
) -> csr_matrix:
    """Convert an operator or list of operators to a sparse matrix.
    Args:
        op: Either an Operator to be converted to an sparse matrix, a list of Operators
            to be converted to a 3d sparse matrix, or a sparse matrix (which simply gets
            returned)
        no_iter (Bool): Boolean determining whether to recursively unroll `Iterables`.
            If recurring, this should be True to avoid making each element of the
            input into a separate csr_matrix.
    Returns:
        csr_matrix: Sparse matrix version of input
    """
    if op is None:
        return op

    if isinstance(op, csr_matrix):
        return op
    if isinstance_qutip_qobj(op):
        return op.data
    if isinstance(op, np.ndarray) and op.dtype == "O":
        op = list(op)
    if isinstance(op, (Array, np.ndarray)) and op.ndim < 3:
        return csr_matrix(op)

    if isinstance(op, Iterable) and not no_iter:
        return [to_csr(item, no_iter=True) for item in op]
    else:
        return csr_matrix(op)


@requires_backend("jax")
def to_BCOO(op: Union[Operator, Array, List[Operator], List[Array], spmatrix, "BCOO"]) -> "BCOO":
    """Convert input op or list of ops to a jax BCOO sparse array.

    Calls ``to_array`` to handle general conversion to a numpy or jax array, then
    builds the BCOO sparse array from the result.

    Args:
        op: Operator or list of operators to convert.

    Returns:
        BCOO: BCOO sparse version of the operator.
    """
    if op is None:
        return op

    if type(op).__name__ == "BCOO":
        return op

    return jsparse.BCOO.fromdense(to_array(op).data)


def to_numeric_matrix_type(
    op: Union[Operator, Array, spmatrix, List[Operator], List[Array], List[spmatrix]]
):
    """Given an operator, array, sparse matrix, or a list of operators, arrays, or sparse matrices,
    attempts to leave them in their original form, only converting the operator to an array,
    and converting lists as necessary. Summarized below:
    - operator is converted to array
    - spmatrix and Array are unchanged
    - lists of Arrays and sparse matrices are passed to their respective to_ functions
    - anything else is passed to to_array
    Args:
        op: An operator, array, sparse matrix, or list of operators, arrays or sparse matrices.
    Returns:
        Array: Array version of input
        csr_matrix: Sparse matrix version of input
    """

    if op is None:
        return op

    elif isinstance_qutip_qobj(op):
        return to_csr(op.data)

    elif isinstance(op, Array):
        return op
    elif isinstance(op, spmatrix):
        return op
    elif type(op).__name__ == "BCOO":
        return op
    elif isinstance(op, Operator):
        return to_array(op)

    elif isinstance(op, Iterable) and isinstance(op[0], spmatrix):
        return to_csr(op)

    elif isinstance(op, Iterable) and isinstance_qutip_qobj(op[0]):
        return to_csr(op)

    else:
        return to_array(op)
