# Copyright 2018-2021 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This module contains utilities and auxiliary functions which are shared
across the PennyLane submodules.
"""
# pylint: disable=protected-access,too-many-branches
from collections.abc import Iterable
import functools
import inspect
import itertools
import numbers
from operator import matmul

import numpy as np
import scipy

import pennylane as qml


def decompose_hamiltonian(H, hide_identity=False):
    r"""Decomposes a Hermitian matrix into a linear combination of Pauli operators.

    Args:
        H (array[complex]): a Hermitian matrix of dimension :math:`2^n\times 2^n`
        hide_identity (bool): does not include the :class:`~.Identity` observable within
            the tensor products of the decomposition if ``True``

    Returns:
        tuple[list[float], list[~.Observable]]: a list of coefficients and a list
        of corresponding tensor products of Pauli observables that decompose the Hamiltonian.

    **Example:**

    We can use this function to compute the Pauli operator decomposition of an arbitrary Hermitian
    matrix:

    >>> A = np.array(
    ... [[-2, -2+1j, -2, -2], [-2-1j,  0,  0, -1], [-2,  0, -2, -1], [-2, -1, -1,  0]])
    >>> coeffs, obs_list = decompose_hamiltonian(A)
    >>> coeffs
    [-1.0, -1.5, -0.5, -1.0, -1.5, -1.0, -0.5, 1.0, -0.5, -0.5]

    We can use the output coefficients and tensor Pauli terms to construct a :class:`~.Hamiltonian`:

    >>> H = qml.Hamiltonian(coeffs, obs_list)
    >>> print(H)
    (-1.0) [I0 I1]
    + (-1.5) [X1]
    + (-0.5) [Y1]
    + (-1.0) [Z1]
    + (-1.5) [X0]
    + (-1.0) [X0 X1]
    + (-0.5) [X0 Z1]
    + (1.0) [Y0 Y1]
    + (-0.5) [Z0 X1]
    + (-0.5) [Z0 Y1]

    This Hamiltonian can then be used in defining VQE problems using :class:`~.ExpvalCost`.
    """
    n = int(np.log2(len(H)))
    N = 2 ** n

    if H.shape != (N, N):
        raise ValueError(
            "The Hamiltonian should have shape (2**n, 2**n), for any qubit number n>=1"
        )

    if not np.allclose(H, H.conj().T):
        raise ValueError("The Hamiltonian is not Hermitian")

    paulis = [qml.Identity, qml.PauliX, qml.PauliY, qml.PauliZ]
    obs = []
    coeffs = []

    for term in itertools.product(paulis, repeat=n):
        matrices = [i._matrix() for i in term]
        coeff = np.trace(functools.reduce(np.kron, matrices) @ H) / N
        coeff = np.real_if_close(coeff).item()

        if not np.allclose(coeff, 0):
            coeffs.append(coeff)

            if not all(t is qml.Identity for t in term) and hide_identity:
                obs.append(
                    functools.reduce(
                        matmul,
                        [t(i) for i, t in enumerate(term) if t is not qml.Identity],
                    )
                )
            else:
                obs.append(functools.reduce(matmul, [t(i) for i, t in enumerate(term)]))

    return coeffs, obs


def sparse_hamiltonian(H, wires=None):
    r"""Computes the sparse matrix representation a Hamiltonian in the computational basis.

    Args:
        H (~.Hamiltonian): Hamiltonian operator for which the matrix representation should be
         computed
        wires (Iterable): Wire labels that indicate the order of wires according to which the matrix
         is constructed. If not profided, ``H.wires`` is used.

    Returns:
        coo_matrix: a sparse matrix in scipy coordinate list (COO) format with dimension
        :math:`(2^n, 2^n)`, where :math:`n` is the number of wires

    **Example:**

    This function can be used by passing a `qml.Hamiltonian` object as:

    >>> coeffs = [1, -0.45]
    >>> obs = [qml.PauliZ(0) @ qml.PauliZ(1), qml.PauliY(0) @ qml.PauliZ(1)]
    >>> H = qml.Hamiltonian(coeffs, obs)
    >>> H_sparse = sparse_hamiltonian(H)
    >>> H_sparse
    <4x4 sparse matrix of type '<class 'numpy.complex128'>'
        with 2 stored elements in COOrdinate format>

    The resulting sparse matrix can be either used directly or transformed into a numpy array:

    >>> H_sparse.toarray()
    array([[ 1.+0.j  ,  0.+0.j  ,  0.+0.45j,  0.+0.j  ],
           [ 0.+0.j  , -1.+0.j  ,  0.+0.j  ,  0.-0.45j],
           [ 0.-0.45j,  0.+0.j  , -1.+0.j  ,  0.+0.j  ],
           [ 0.+0.j  ,  0.+0.45j,  0.+0.j  ,  1.+0.j  ]])
    """
    if not isinstance(H, qml.Hamiltonian):
        raise TypeError("Passed Hamiltonian must be of type `qml.Hamiltonian`")

    if wires is None:
        wires = H.wires
    else:
        wires = qml.wires.Wires(wires)

    n = len(wires)
    matrix = scipy.sparse.coo_matrix((2 ** n, 2 ** n), dtype="complex128")

    coeffs = qml.math.toarray(H.data)

    for coeff, op in zip(coeffs, H.ops):

        obs = []
        for o in qml.operation.Tensor(op).obs:
            if len(o.wires) > 1:
                # todo: deal with operations created from multi-qubit operations such as Hermitian
                raise ValueError(
                    "Can only sparsify Hamiltonians whose constituent observables consist of "
                    "(tensor products of) single-qubit operators; got {}.".format(op)
                )
            obs.append(scipy.sparse.coo_matrix(o.matrix))

        mat = [scipy.sparse.eye(2, format="coo")] * n

        for i, wire in enumerate(op.wires):
            # find index of this wire in the ordering
            idx = wires.index(wire)
            mat[idx] = obs[i]

        matrix += functools.reduce(lambda i, j: scipy.sparse.kron(i, j, format="coo"), mat) * coeff

    return matrix.tocoo()


def _flatten(x):
    """Iterate recursively through an arbitrarily nested structure in depth-first order.

    See also :func:`_unflatten`.

    Args:
        x (array, Iterable, Any): each element of an array or an Iterable may itself be any of these types

    Yields:
        Any: elements of x in depth-first order
    """
    if isinstance(x, np.ndarray):
        yield from _flatten(x.flat)  # should we allow object arrays? or just "yield from x.flat"?
    elif isinstance(x, qml.wires.Wires):
        # Reursive calls to flatten `Wires` will cause infinite recursion (`Wires` atoms are `Wires`).
        # Since Wires are always flat, just yield.
        for item in x:
            yield item
    elif isinstance(x, Iterable) and not isinstance(x, (str, bytes)):
        for item in x:
            yield from _flatten(item)
    else:
        yield x


def _unflatten(flat, model):
    """Restores an arbitrary nested structure to a flattened iterable.

    See also :func:`_flatten`.

    Args:
        flat (array): 1D array of items
        model (array, Iterable, Number): model nested structure

    Raises:
        TypeError: if ``model`` contains an object of unsupported type

    Returns:
        Union[array, list, Any], array: first elements of flat arranged into the nested
        structure of model, unused elements of flat
    """
    if isinstance(model, (numbers.Number, str)):
        return flat[0], flat[1:]

    if isinstance(model, np.ndarray):
        idx = model.size
        res = np.array(flat)[:idx].reshape(model.shape)
        return res, flat[idx:]

    if isinstance(model, Iterable):
        res = []
        for x in model:
            val, flat = _unflatten(flat, x)
            res.append(val)
        return res, flat

    raise TypeError("Unsupported type in the model: {}".format(type(model)))


def unflatten(flat, model):
    """Wrapper for :func:`_unflatten`.

    Args:
        flat (array): 1D array of items
        model (array, Iterable, Number): model nested structure

    Raises:
        ValueError: if ``flat`` has more elements than ``model``
    """
    # pylint:disable=len-as-condition
    res, tail = _unflatten(np.asarray(flat), model)
    if len(tail) != 0:
        raise ValueError("Flattened iterable has more elements than the model.")
    return res


def _inv_dict(d):
    """Reverse a dictionary mapping.

    Returns multimap where the keys are the former values,
    and values are sets of the former keys.

    Args:
        d (dict[a->b]): mapping to reverse

    Returns:
        dict[b->set[a]]: reversed mapping
    """
    ret = {}
    for k, v in d.items():
        ret.setdefault(v, set()).add(k)
    return ret


def _get_default_args(func):
    """Get the default arguments of a function.

    Args:
        func (callable): a function

    Returns:
        dict[str, tuple]: mapping from argument name to (positional idx, default value)
    """
    signature = inspect.signature(func)
    return {
        k: (idx, v.default)
        for idx, (k, v) in enumerate(signature.parameters.items())
        if v.default is not inspect.Parameter.empty
    }


@functools.lru_cache()
def pauli_eigs(n):
    r"""Eigenvalues for :math:`A^{\otimes n}`, where :math:`A` is
    Pauli operator, or shares its eigenvalues.

    As an example if n==2, then the eigenvalues of a tensor product consisting
    of two matrices sharing the eigenvalues with Pauli matrices is returned.

    Args:
        n (int): the number of qubits the matrix acts on
    Returns:
        list: the eigenvalues of the specified observable
    """
    if n == 1:
        return np.array([1, -1])
    return np.concatenate([pauli_eigs(n - 1), -pauli_eigs(n - 1)])


def expand(matrix, original_wires, expanded_wires):
    r"""Expand a an operator matrix to more wires.

    Args:
        matrix (array): :math:`2^n \times 2^n` matrix where n = len(original_wires).
        original_wires (Sequence[int]): original wires of matrix
        expanded_wires (Union[Sequence[int], int]): expanded wires of matrix, can be shuffled.
            If a single int m is given, corresponds to list(range(m))

    Returns:
        array: :math:`2^m \times 2^m` matrix where m = len(expanded_wires).
    """
    if isinstance(expanded_wires, numbers.Integral):
        expanded_wires = list(range(expanded_wires))

    N = len(original_wires)
    M = len(expanded_wires)
    D = M - N

    if not set(expanded_wires).issuperset(original_wires):
        raise ValueError("Invalid target subsystems provided in 'original_wires' argument.")

    if qml.math.shape(matrix) != (2 ** N, 2 ** N):
        raise ValueError(
            "Matrix parameter must be of size (2**len(original_wires), 2**len(original_wires))"
        )

    dims = [2] * (2 * N)
    tensor = qml.math.reshape(matrix, dims)

    interface = qml.math.get_interface(tensor)

    if D > 0:
        extra_dims = [2] * (2 * D)
        identity = qml.math.reshape(qml.math.eye(2 ** D), extra_dims)

        if interface == "tensorflow":
            identity = qml.math.cast_like(identity, tensor)

        expanded_tensor = qml.math.tensordot(tensor, identity, axes=0)
        # Fix order of tensor factors
        expanded_tensor = qml.math.moveaxis(
            expanded_tensor, tuple(range(2 * N, 2 * N + D)), tuple(range(N, N + D))
        )
    else:
        expanded_tensor = tensor

    wire_indices = []
    for wire in original_wires:
        wire_indices.append(expanded_wires.index(wire))

    wire_indices = np.array(wire_indices)

    # Order tensor factors according to wires
    original_indices = np.array(range(N))
    expanded_tensor = qml.math.moveaxis(
        expanded_tensor, tuple(original_indices), tuple(wire_indices)
    )
    expanded_tensor = qml.math.moveaxis(
        expanded_tensor, tuple(original_indices + M), tuple(wire_indices + M)
    )

    return qml.math.reshape(expanded_tensor, (2 ** M, 2 ** M))


def expand_vector(vector, original_wires, expanded_wires):
    r"""Expand a vector to more wires.

    Args:
        vector (array): :math:`2^n` vector where n = len(original_wires).
        original_wires (Sequence[int]): original wires of vector
        expanded_wires (Union[Sequence[int], int]): expanded wires of vector, can be shuffled
            If a single int m is given, corresponds to list(range(m))

    Returns:
        array: :math:`2^m` vector where m = len(expanded_wires).
    """
    if isinstance(expanded_wires, numbers.Integral):
        expanded_wires = list(range(expanded_wires))

    N = len(original_wires)
    M = len(expanded_wires)
    D = M - N

    if not set(expanded_wires).issuperset(original_wires):
        raise ValueError("Invalid target subsystems provided in 'original_wires' argument.")

    if qml.math.shape(vector) != (2 ** N,):
        raise ValueError("Vector parameter must be of length 2**len(original_wires)")

    dims = [2] * N
    tensor = qml.math.reshape(vector, dims)

    if D > 0:
        extra_dims = [2] * D
        ones = qml.math.ones(2 ** D).reshape(extra_dims)
        expanded_tensor = qml.math.tensordot(tensor, ones, axes=0)
    else:
        expanded_tensor = tensor

    wire_indices = []
    for wire in original_wires:
        wire_indices.append(expanded_wires.index(wire))

    wire_indices = np.array(wire_indices)

    # Order tensor factors according to wires
    original_indices = np.array(range(N))
    expanded_tensor = qml.math.moveaxis(
        expanded_tensor, tuple(original_indices), tuple(wire_indices)
    )

    return qml.math.reshape(expanded_tensor, 2 ** M)
