# -*- coding: utf-8 -*-
#
# This module was initially contributed by Ben Criger.
#
"""
This module implements transformations between superoperator representations,
including supermatrix, Kraus, Choi and Chi (process) matrix formalisms.
"""

__all__ = [
    'kraus_to_choi', 'kraus_to_super',
    'to_choi', 'to_chi', 'to_super', 'to_kraus', 'to_stinespring',
]

import itertools
import numbers

import numpy as np
import scipy.linalg

from . import data as _data
from .superoperator import stack_columns, unstack_columns, sprepost
from .tensor import tensor
from .dimensions import flatten
from .qobj import Qobj
from .operators import identity, sigmax, sigmay, sigmaz
from .states import basis


# TODO: revisit when creation routines have dispatching.
_SINGLE_QUBIT_PAULI_BASIS = (
    identity(2).to(_data.CSR),
    sigmax().to(_data.CSR),
    sigmay().to(_data.CSR),
    sigmaz().to(_data.CSR),
)


def _superpauli_basis(nq=1):
    dims = [[[2] * nq] * 2] * 2
    nnz = 8**nq
    data = _data.csr.empty(4**nq, 4**nq, nnz)
    sci = data.as_scipy(full=True)
    ptr, ptr_inc = 0, 2**nq
    # Construct the Pauli basis by vertically stacking rows in sparse format.
    # The CSR format is much more efficient at handling row-stacking, so we
    # actually have to do a little dance through adjoint/transpose to get it
    # into the right format.
    for i, paulis in enumerate(itertools.product(_SINGLE_QUBIT_PAULI_BASIS,
                                                 repeat=nq)):
        basis = paulis[0].data
        for pauli in paulis[1:]:
            basis = _data.kron_csr(basis, pauli.data)
        basis_ket_sci = _data.column_stack_csr(basis).transpose().as_scipy()
        sci.data[ptr : ptr+ptr_inc] = basis_ket_sci.data
        sci.indices[ptr : ptr+ptr_inc] = basis_ket_sci.indices
        sci.indptr[i] = ptr
        ptr += ptr_inc
    sci.indptr[-1] = nnz
    return Qobj(data.adjoint(),
                dims=dims,
                type='super',
                superrep='super',
                isherm=False,
                isunitary=False,
                copy=False)


def _int_log_two(x):
    return int(x).bit_length() - 1


def _is_power_of_two(x):
    return isinstance(x, numbers.Integral) and x == 2**_int_log_two(x)


def _nq(dims):
    dim = np.prod(dims[0][0])
    nq = _int_log_two(dim)
    if 2 ** nq != dim:
        raise ValueError("{} is not an integer power of 2.".format(dim))
    return nq


def isqubitdims(dims):
    """
    Checks whether all entries in a dims list are integer powers of 2.

    Parameters
    ----------
    dims : nested list of ints
        Dimensions to be checked.

    Returns
    -------
    isqubitdims : bool
        True if and only if every member of the flattened dims
        list is an integer power of 2.
    """
    return all(_is_power_of_two(dim) for dim in flatten(dims))


def _to_superpauli(q_oper):
    """
    Convert a superoperator to the Pauli basis (assuming qubit dimensions).

    This is an internal function, as QuTiP does not currently have
    a way to mark that superoperators are represented in the Pauli
    basis as opposed to the column-stacking basis; a Pauli-basis
    ``type='super'`` would thus break other conversion functions.
    """
    # Ensure we start with a column-stacking-basis superoperator.
    sqobj = to_super(q_oper)
    if not isqubitdims(sqobj.dims):
        raise ValueError("Pauli basis is only defined for qubits.")
    nq = _int_log_two(sqobj.shape[0]) // 2
    B = _superpauli_basis(nq) * 2**(-0.5 * nq)
    # To do this, we have to hack a bit and force the dims to match,
    # since the superpauli_basis function makes different assumptions
    # about indices than we need here.
    B.dims = sqobj.dims
    return B.dag() @ sqobj @ B


def _choi_to_kraus(q_oper, tol=1e-9):
    """
    Takes a Choi matrix and returns a list of Kraus operators.
    TODO: Create a new class structure for quantum channels, perhaps as a
    strict sub-class of Qobj.
    """
    vals, vecs = q_oper.eigenstates()
    dims = [q_oper.dims[0][1], q_oper.dims[0][0]]
    shape = (np.prod(q_oper.dims[0][1]), np.prod(q_oper.dims[0][0]))
    return [Qobj(_data.mul(unstack_columns(vec.data, shape=shape), np.sqrt(val)),
                 dims=dims, type='oper', copy='False')
            for val, vec in zip(vals, vecs) if abs(val) >= tol]


# Individual conversions from Kraus operators are public because the output
# list of Kraus operators is not itself a quantum object.

def kraus_to_choi(kraus_list):
    """
    Take a list of Kraus operators and returns the Choi matrix for the channel
    represented by the Kraus operators in `kraus_list`
    """
    kraus_mat_list = [k.full() for k in kraus_list]
    op_rng = list(range(kraus_mat_list[0].shape[1]))
    choi_blocks = np.array(
        [[sum(op[:, c_ix, None] @ np.conj(op[None, :, r_ix])
              for op in kraus_mat_list)
          for r_ix in op_rng]
         for c_ix in op_rng]
    )
    return Qobj(np.hstack(np.hstack(choi_blocks)),
                dims=[kraus_list[0].dims[::-1]]*2,
                type='super',
                superrep='choi',
                copy=False)


def kraus_to_super(kraus_list):
    """
    Convert a list of Kraus operators to a superoperator.
    """
    return to_super(kraus_to_choi(kraus_list))


def _super_tofrom_choi(q_oper):
    """
    We exploit that the basis transformation between Choi and supermatrix
    representations squares to the identity, so that if we munge Qobj.type,
    we can use the same function.
    """
    if not q_oper.issuper:
        raise ValueError("needs to be a superoperator")
    if q_oper.superrep not in ('super', 'choi'):
        raise ValueError("operator is not in super or choi format")
    data = q_oper.data.to_array()
    dims = q_oper.dims
    new_dims = [[dims[1][1], dims[0][1]], [dims[1][0], dims[0][0]]]
    d0 = np.prod(flatten(new_dims[0]))
    d1 = np.prod(flatten(new_dims[1]))
    s0 = np.prod(dims[0][0])
    s1 = np.prod(dims[1][1])
    data = data.reshape([s0, s1, s0, s1]).transpose(3, 1, 2, 0).reshape(d0, d1)
    return Qobj(data,
                dims=new_dims,
                type='super',
                superrep='super' if q_oper.superrep == 'choi' else 'choi',
                copy=False)


def _choi_to_chi(q_oper):
    """
    Converts a Choi matrix to a Chi matrix in the Pauli basis.

    NOTE: this is only supported for qubits right now. Need to extend to
    Heisenberg-Weyl for other subsystem dimensions.
    """
    nq = _nq(q_oper.dims)
    B = _superpauli_basis(nq).data
    return Qobj(_data.matmul(_data.matmul(B.adjoint(), q_oper.data), B),
                dims=q_oper.dims,
                type='super',
                superrep='chi',
                copy=False)


def _chi_to_choi(q_oper):
    """
    Converts a Chi matrix to a Choi matrix.

    NOTE: this is only supported for qubits right now. Need to extend to
    Heisenberg-Weyl for other subsystem dimensions.
    """
    nq = _nq(q_oper.dims)
    B = _superpauli_basis(nq).data
    # The Chi matrix has tr(chi) == d², so we need to divide out
    # by that to get back to the Choi form.
    return Qobj(_data.mul(
                    _data.matmul(_data.matmul(B, q_oper.data), B.adjoint()),
                    1 / q_oper.shape[0]
                ),
                dims=q_oper.dims,
                type='super',
                superrep='choi',
                copy=False)


def _svd_u_to_kraus(U, S, d, dK, indims, outdims):
    """
    Given a partial isometry U and a vector of square-roots of singular values
    S obtained from a SVD, produces the Kraus operators represented by U.

    Returns
    -------
    Ks : list of Qobj
        Quantum objects represnting each of the Kraus operators.
    """
    # We use U * S since S is 1-index, such that this is equivalent to
    # U . diag(S), but easier to write down.
    data = np.array(U * S).reshape((d, d, dK), order='F').transpose((2, 0, 1))
    return [
        Qobj(x,
             dims=[outdims, indims],
             type='oper',
             copy=False)
        for x in data
    ]


def _generalized_kraus(q_oper, threshold=1e-10):
    # TODO: document!
    # TODO: use this to generalize to_kraus to the case where U != V.
    #       This is critical for non-CP maps, as appear in (for example)
    #       diamond norm differences between two CP maps.
    if not q_oper.issuper or q_oper.superrep != "choi":
        raise ValueError("".join([
            "Expected a Choi matrix, got a ", repr(q_oper.type),
            " (superrep ", repr(q_oper.superrep), ")."
        ]))

    # Remember the shape of the underlying space,
    # as we'll need this to make Kraus operators later.
    dL, dR = int(np.sqrt(q_oper.shape[0])), int(np.sqrt(q_oper.shape[1]))
    # Also remember the dims breakout.
    out_dims, in_dims = q_oper.dims
    out_left, out_right = out_dims
    in_left, in_right = in_dims

    # Find the SVD.
    U, S, V = scipy.linalg.svd(q_oper.full())

    # Truncate away the zero singular values, up to a threshold.
    nonzero_idxs = S > threshold
    dK = nonzero_idxs.sum()
    U = np.array(U)[:, nonzero_idxs]
    # We also want S to be a single index array, which np.matrix
    # doesn't allow for. This is stripped by calling array() on it.
    S = np.sqrt(np.array(S)[nonzero_idxs])
    # Since NumPy returns V and not V+, we need to take the dagger
    # to get back to quantum info notation for Stinespring pairs.
    V = np.array(V.conj().T)[:, nonzero_idxs]

    # Next, we convert each of U and V into Kraus operators.
    # Finally, we want the Kraus index to be left-most so that we
    # can map over it when making Qobjs.
    # FIXME: does not preserve dims!
    kU = _svd_u_to_kraus(U, S, dL, dK, out_right, out_left)
    kV = _svd_u_to_kraus(V, S, dL, dK, in_right, in_left)

    return kU, kV


def _choi_to_stinespring(q_oper, threshold=1e-10):
    # TODO: document!
    kU, kV = _generalized_kraus(q_oper, threshold=threshold)

    assert len(kU) == len(kV)
    dK = len(kU)
    dL = kU[0].shape[0]
    dR = kV[0].shape[1]
    # Also remember the dims breakout.
    out_dims, in_dims = q_oper.dims
    out_left, out_right = out_dims
    in_left, in_right = in_dims

    A = Qobj(_data.zeros(dK * dL, dL),
             dims=[out_left + [dK], out_right + [1]],
             type='oper',
             isherm=True,
             isunitary=False,
             copy=False)
    B = Qobj(_data.zeros(dK * dR, dR),
             dims=[in_left + [dK], in_right + [1]],
             type='oper',
             isherm=True,
             isunitary=False,
             copy=False)

    for idx_kraus, (KL, KR) in enumerate(zip(kU, kV)):
        A += tensor(KL, basis(dK, idx_kraus))
        B += tensor(KR, basis(dK, idx_kraus))

    # There is no input (right) Kraus index, so strip that off.
    del A.dims[1][-1]
    del B.dims[1][-1]

    return A, B


def to_choi(q_oper):
    """
    Converts a Qobj representing a quantum map to the Choi representation,
    such that the trace of the returned operator is equal to the dimension
    of the system.

    Parameters
    ----------
    q_oper : Qobj
        Superoperator to be converted to Choi representation. If
        ``q_oper`` is ``type="oper"``, then it is taken to act by conjugation,
        such that ``to_choi(A) == to_choi(sprepost(A, A.dag()))``.

    Returns
    -------
    choi : Qobj
        A quantum object representing the same map as ``q_oper``, such that
        ``choi.superrep == "choi"``.

    Raises
    ------
    TypeError: if the given quantum object is not a map, or cannot be converted
        to Choi representation.
    """
    if q_oper.type == 'super':
        if q_oper.superrep == 'choi':
            return q_oper
        if q_oper.superrep == 'super':
            return _super_tofrom_choi(q_oper)
        if q_oper.superrep == 'chi':
            return _chi_to_choi(q_oper)
        else:
            raise TypeError(q_oper.superrep)
    elif q_oper.type == 'oper':
        return _super_tofrom_choi(sprepost(q_oper, q_oper.dag()))
    else:
        raise TypeError(
            "Conversion of Qobj with type = {0.type} "
            "and superrep = {0.choi} to Choi not supported.".format(q_oper)
        )


def to_chi(q_oper):
    """
    Converts a Qobj representing a quantum map to a representation as a chi
    (process) matrix in the Pauli basis, such that the trace of the returned
    operator is equal to the dimension of the system.

    Parameters
    ----------
    q_oper : Qobj
        Superoperator to be converted to Chi representation. If
        ``q_oper`` is ``type="oper"``, then it is taken to act by conjugation,
        such that ``to_chi(A) == to_chi(sprepost(A, A.dag()))``.

    Returns
    -------
    chi : Qobj
        A quantum object representing the same map as ``q_oper``, such that
        ``chi.superrep == "chi"``.

    Raises
    ------
    TypeError: if the given quantum object is not a map, or cannot be converted
        to Chi representation.
    """
    if q_oper.type == 'super':
        if q_oper.superrep == 'chi':
            return q_oper
        elif q_oper.superrep == 'choi':
            return _choi_to_chi(q_oper)
        elif q_oper.superrep == 'super':
            return _choi_to_chi(to_choi(q_oper))
        else:
            raise TypeError(q_oper.superrep)
    elif q_oper.type == 'oper':
        return to_chi(sprepost(q_oper, q_oper.dag()))
    else:
        raise TypeError(
            "Conversion of Qobj with type = {0.type} "
            "and superrep = {0.choi} to Choi not supported.".format(q_oper)
        )


def to_super(q_oper):
    """
    Converts a Qobj representing a quantum map to the supermatrix (Liouville)
    representation.

    Parameters
    ----------
    q_oper : Qobj
        Superoperator to be converted to supermatrix representation. If
        ``q_oper`` is ``type="oper"``, then it is taken to act by conjugation,
        such that ``to_super(A) == sprepost(A, A.dag())``.

    Returns
    -------
    superop : Qobj
        A quantum object representing the same map as ``q_oper``, such that
        ``superop.superrep == "super"``.

    Raises
    ------
    TypeError
        If the given quantum object is not a map, or cannot be converted
        to supermatrix representation.
    """
    if q_oper.type == 'super':
        if q_oper.superrep == "super":
            return q_oper
        elif q_oper.superrep == 'choi':
            return _super_tofrom_choi(q_oper)
        elif q_oper.superrep == 'chi':
            return to_super(to_choi(q_oper))
        else:
            raise ValueError(
                "Unrecognized superrep '{}'.".format(q_oper.superrep))
    elif q_oper.type == 'oper':  # Assume unitary
        return sprepost(q_oper, q_oper.dag())
    else:
        raise TypeError(
            "Conversion of Qobj with type = {0.type} "
            "and superrep = {0.superrep} to supermatrix not "
            "supported.".format(q_oper)
        )


def to_kraus(q_oper, tol=1e-9):
    """
    Converts a Qobj representing a quantum map to a list of quantum objects,
    each representing an operator in the Kraus decomposition of the given map.

    Parameters
    ----------
    q_oper : Qobj
        Superoperator to be converted to Kraus representation. If
        ``q_oper`` is ``type="oper"``, then it is taken to act by conjugation,
        such that ``to_kraus(A) == to_kraus(sprepost(A, A.dag())) == [A]``.

    tol : Float
        Optional threshold parameter for eigenvalues/Kraus ops to be discarded.
        The default is to=1e-9.

    Returns
    -------
    kraus_ops : list of Qobj
        A list of quantum objects, each representing a Kraus operator in the
        decomposition of ``q_oper``.

    Raises
    ------
    TypeError: if the given quantum object is not a map, or cannot be
        decomposed into Kraus operators.
    """
    if q_oper.issuper:
        if q_oper.superrep != 'choi':
            q_oper = to_choi(q_oper)
        return _choi_to_kraus(q_oper, tol)
    elif q_oper.isoper:  # Assume unitary
        return [q_oper]
    raise TypeError(
        "Conversion of Qobj with type={0.type} "
        "and superrep={0.superrep} to Kraus decomposition not "
        "supported.".format(q_oper)
    )


def to_stinespring(q_oper, threshold=1e-10):
    r"""
    Converts a Qobj representing a quantum map $\Lambda$ to a pair of partial
    isometries $A$ and $B$ such that $\Lambda(X) = \Tr_2(A X B^\dagger)$ for
    all inputs $X$, where the partial trace is taken over a a new index on the
    output dimensions of $A$ and $B$.

    For completely positive inputs, $A$ will always equal $B$ up to precision
    errors.

    Parameters
    ----------
    q_oper : Qobj
        Superoperator to be converted to a Stinespring pair.

    Returns
    -------
    A, B : Qobj
        Quantum objects representing each of the Stinespring matrices for the
        input Qobj.
    """
    return _choi_to_stinespring(to_choi(q_oper), threshold)
