def fidelity(A, B):
    """
    Calculates the fidelity (pseudo-metric) between two density matrices.
    See: Nielsen & Chuang, "Quantum Computation and Quantum Information"

    Parameters
    ----------
    A : qobj
        Density matrix or state vector.
    B : qobj
        Density matrix or state vector with same dimensions as A.

    Returns
    -------
    fid : float
        Fidelity pseudo-metric between A and B.

    Examples
    --------
    >>> x = fock_dm(5,3)
    >>> y = coherent_dm(5,1)
    >>> fidelity(x,y)
    0.24104350624628332

    """
    if A.isket or A.isbra:
        A = ket2dm(A)
    if B.isket or B.isbra:
        B = ket2dm(B)

    if A.dims != B.dims:
        raise TypeError('Density matrices do not have same dimensions.')

    A = A.sqrtm()
    return float(np.real((A * (B * A)).sqrtm().tr()))

def hilbert_dist(A, B):
    """
    Returns the Hilbert-Schmidt distance between two density matrices A & B.

    Parameters
    ----------
    A : qobj
        Density matrix or state vector.
    B : qobj
        Density matrix or state vector with same dimensions as A.

    Returns
    -------
    dist : float
        Hilbert-Schmidt distance between density matrices.

    """
    if A.isket or A.isbra:
        A = ket2dm(A)
    if B.isket or B.isbra:
        B = ket2dm(B)

    if A.dims != B.dims:
        raise TypeError('A and B do not have same dimensions.')

    return (A - B).norm('fro')

def hilbert_dist(A, B):
    """
    Returns the Hilbert-Schmidt distance between two density matrices A & B.

    Parameters
    ----------
    A : qobj
        Density matrix or state vector.
    B : qobj
        Density matrix or state vector with same dimensions as A.

    Returns
    -------
    dist : float
        Hilbert-Schmidt distance between density matrices.

    Notes
    -----
    See V. Vedral and M. B. Plenio, Phys. Rev. A 57, 1619 (1998).

    """
    if A.isket or A.isbra:
        A = ket2dm(A)
    if B.isket or B.isbra:
        B = ket2dm(B)

    if A.dims != B.dims:
        raise TypeError('A and B do not have same dimensions.')

    return ((A - B)**2).tr()

def bures_angle(A, B):
    """
    Returns the Bures Angle between two density matrices A & B.

    The Bures angle ranges from 0, for states with unit fidelity, to pi/2.

    Parameters
    ----------
    A : qobj
        Density matrix or state vector.
    B : qobj
        Density matrix or state vector with same dimensions as A.

    Returns
    -------
    angle : float
        Bures angle between density matrices.
    """
    if A.isket or A.isbra:
        A = ket2dm(A)
    if B.isket or B.isbra:
        B = ket2dm(B)

    if A.dims != B.dims:
        raise TypeError('A and B do not have same dimensions.')

    return np.arccos(fidelity(A, B))

def bures_dist(A, B):
    """
    Returns the Bures distance between two density matrices A & B.

    The Bures distance ranges from 0, for states with unit fidelity,
    to sqrt(2).

    Parameters
    ----------
    A : qobj
        Density matrix or state vector.
    B : qobj
        Density matrix or state vector with same dimensions as A.

    Returns
    -------
    dist : float
        Bures distance between density matrices.
    """
    if A.isket or A.isbra:
        A = ket2dm(A)
    if B.isket or B.isbra:
        B = ket2dm(B)

    if A.dims != B.dims:
        raise TypeError('A and B do not have same dimensions.')

    dist = np.sqrt(2.0 * (1.0 - fidelity(A, B)))
    return dist

def fidelity(A, B):
    """
    Calculates the fidelity (pseudo-metric) between two density matrices.
    See: Nielsen & Chuang, "Quantum Computation and Quantum Information"

    Parameters
    ----------
    A : qobj
        Density matrix or state vector.
    B : qobj
        Density matrix or state vector with same dimensions as A.

    Returns
    -------
    fid : float
        Fidelity pseudo-metric between A and B.

    Examples
    --------
    >>> x = fock_dm(5,3)
    >>> y = coherent_dm(5,1)
    >>> fidelity(x,y)
    0.24104350624628332

    """
    if A.isket or A.isbra:
        # Take advantage of the fact that the density operator for A
        # is a projector to avoid a sqrtm call.
        sqrtmA = ket2dm(A)
        # Check whether we have to turn B into a density operator, too.
        if B.isket or B.isbra:
            B = ket2dm(B)
    else:
        if B.isket or B.isbra:
            # Swap the order so that we can take a more numerically
            # stable square root of B.
            return fidelity(B, A)
        # If we made it here, both A and B are operators, so
        # we have to take the sqrtm of one of them.
        sqrtmA = A.sqrtm()

    if sqrtmA.dims != B.dims:
        raise TypeError('Density matrices do not have same dimensions.')

    # We don't actually need the whole matrix here, just the trace
    # of its square root, so let's just get its eigenenergies instead.
    # We also truncate negative eigenvalues to avoid nan propagation;
    # even for positive semidefinite matrices, small negative eigenvalues
    # can be reported.
    eig_vals = (sqrtmA * B * sqrtmA).eigenenergies()
    return float(np.real(np.sqrt(eig_vals[eig_vals > 0]).sum()))

def test_composite_vec():
    """
    Composite: Tests compositing states and density operators.
    """
    k1 = rand_ket(5)
    k2 = rand_ket(7)
    r1 = operator_to_vector(ket2dm(k1))
    r2 = operator_to_vector(ket2dm(k2))

    r3 = operator_to_vector(rand_dm(3))
    r4 = operator_to_vector(rand_dm(4))

    assert_(composite(k1, k2) == tensor(k1, k2))
    assert_(composite(r3, r4) == super_tensor(r3, r4))
    assert_(composite(k1, r4) == super_tensor(r1, r4))
    assert_(composite(r3, k2) == super_tensor(r3, r2))

def _subsystem_apply_reference(state, channel, mask):
    if isket(state):
        state = ket2dm(state)

    if isoper(channel):
        full_oper = tensor([channel if mask[j]
                            else qeye(state.dims[0][j])
                            for j in range(len(state.dims[0]))])
        return full_oper * state * full_oper.dag()
    else:
        # Go to Choi, then Kraus
        # chan_mat = array(channel.data.todense())
        choi_matrix = super_to_choi(channel)
        vals, vecs = eig(choi_matrix.full())
        vecs = list(map(array, zip(*vecs)))
        kraus_list = [sqrt(vals[j]) * vec2mat(vecs[j])
                      for j in range(len(vals))]
        # Kraus operators to be padded with identities:
        k_qubit_kraus_list = product(kraus_list, repeat=sum(mask))
        rho_out = Qobj(inpt=zeros(state.shape), dims=state.dims)
        for operator_iter in k_qubit_kraus_list:
            operator_iter = iter(operator_iter)
            op_iter_list = [next(operator_iter).conj().T if mask[j]
                            else qeye(state.dims[0][j])
                            for j in range(len(state.dims[0]))]
            full_oper = tensor(list(map(Qobj, op_iter_list)))
            rho_out = rho_out + full_oper * state * full_oper.dag()
        return Qobj(rho_out)

def _correlation_me_4op_2t(H, rho0, tlist, taulist, c_ops,
                           a_op, b_op, c_op, d_op, reverse=False,
                           args=None, options=Odeoptions()):
    """
    Calculate the four-operator two-time correlation function on the form
    <A(t)B(t+tau)C(t+tau)D(t)>.

    See, Gardiner, Quantum Noise, Section 5.2.1
    """

    if debug:
        print(inspect.stack()[0][3])

    if rho0 is None:
        rho0 = steadystate(H, c_ops)
    elif rho0 and isket(rho0):
        rho0 = ket2dm(rho0)

    C_mat = np.zeros([np.size(tlist), np.size(taulist)], dtype=complex)

    rho_t = mesolve(
        H, rho0, tlist, c_ops, [], args=args, options=options).states

    for t_idx, rho in enumerate(rho_t):
        C_mat[t_idx, :] = mesolve(H, d_op * rho * a_op, taulist,
                                  c_ops, [b_op * c_op],
                                  args=args, options=options).expect[0]

    return C_mat

def test_call():
    """
    Test Qobj: Call
    """
    # Make test objects.
    psi = rand_ket(3)
    rho = rand_dm_ginibre(3)
    U = rand_unitary(3)
    S = rand_super_bcsz(3)

    # Case 0: oper(ket).
    assert U(psi) == U * psi

    # Case 1: oper(oper). Should raise TypeError.
    with expect_exception(TypeError):
        U(rho)

    # Case 2: super(ket).
    assert S(psi) == vector_to_operator(S * operator_to_vector(ket2dm(psi)))

    # Case 3: super(oper).
    assert S(rho) == vector_to_operator(S * operator_to_vector(rho))

    # Case 4: super(super). Should raise TypeError.
    with expect_exception(TypeError):
        S(S)

def _correlation_es_2t(H, state0, tlist, taulist, c_ops, a_op, b_op, c_op):
    """
    Internal function for calculating the three-operator two-time
    correlation function:
    <A(t)B(t+tau)C(t)>
    using an exponential series solver.
    """

    # the solvers only work for positive time differences and the correlators
    # require positive tau
    if state0 is None:
        rho0 = steadystate(H, c_ops)
        tlist = [0]
    elif isket(state0):
        rho0 = ket2dm(state0)
    else:
        rho0 = state0

    if debug:
        print(inspect.stack()[0][3])

    # contruct the Liouvillian
    L = liouvillian(H, c_ops)

    corr_mat = np.zeros([np.size(tlist), np.size(taulist)], dtype=complex)
    solES_t = ode2es(L, rho0)

    # evaluate the correlation function
    for t_idx in range(len(tlist)):
        rho_t = esval(solES_t, [tlist[t_idx]])
        solES_tau = ode2es(L, c_op * rho_t * a_op)
        corr_mat[t_idx, :] = esval(expect(b_op, solES_tau), taulist)

    return corr_mat

def _correlation_me_2op_2t(H, rho0, tlist, taulist, c_ops, a_op, b_op,
                           reverse=False, args=None, options=Odeoptions()):
    """
    Internal function for calculating correlation functions using the master
    equation solver. See :func:`correlation` for usage.
    """

    if debug:
        print(inspect.stack()[0][3])

    if rho0 is None:
        rho0 = steadystate(H, c_ops)
    elif rho0 and isket(rho0):
        rho0 = ket2dm(rho0)

    C_mat = np.zeros([np.size(tlist), np.size(taulist)], dtype=complex)

    rho_t_list = mesolve(
        H, rho0, tlist, c_ops, [], args=args, options=options).states

    if reverse:
        # <A(t)B(t+tau)>
        for t_idx, rho_t in enumerate(rho_t_list):
            C_mat[t_idx, :] = mesolve(H, rho_t * a_op, taulist,
                                      c_ops, [b_op], args=args,
                                      options=options).expect[0]
    else:
        # <A(t+tau)B(t)>
        for t_idx, rho_t in enumerate(rho_t_list):
            C_mat[t_idx, :] = mesolve(H, b_op * rho_t, taulist,
                                      c_ops, [a_op], args=args,
                                      options=options).expect[0]

    return C_mat

def _correlation_es_2op_1t(H, rho0, tlist, c_ops, a_op, b_op, reverse=False,
                           args=None, options=Odeoptions()):
    """
    Internal function for calculating correlation functions using the
    exponential series solver. See :func:`correlation_ss` usage.
    """

    if debug:
        print(inspect.stack()[0][3])

    # contruct the Liouvillian
    L = liouvillian(H, c_ops)

    # find the steady state
    if rho0 is None:
        rho0 = steadystate(L)
    elif rho0 and isket(rho0):
        rho0 = ket2dm(rho0)

    # evaluate the correlation function
    if reverse:
        # <A(t)B(t+tau)>
        solC_tau = ode2es(L, rho0 * a_op)
        return esval(expect(b_op, solC_tau), tlist)
    else:
        # default: <A(t+tau)B(t)>
        solC_tau = ode2es(L, b_op * rho0)
        return esval(expect(a_op, solC_tau), tlist)

    def update(self, rho):
        """
        Calculate the probability function for the given state of an harmonic
        oscillator (as density matrix)
        """

        if isket(rho):
            rho = ket2dm(rho)

        self.data = np.zeros(len(self.xvecs[0]), dtype=complex)
        M, N = rho.shape

        for m in range(M):
            k_m = pow(self.omega / pi, 0.25) / \
                sqrt(2 ** m * factorial(m)) * \
                exp(-self.xvecs[0] ** 2 / 2.0) * \
                np.polyval(hermite(m), self.xvecs[0])

            for n in range(N):
                k_n = pow(self.omega / pi, 0.25) / \
                    sqrt(2 ** n * factorial(n)) * \
                    exp(-self.xvecs[0] ** 2 / 2.0) * \
                    np.polyval(hermite(n), self.xvecs[0])

                self.data += np.conjugate(k_n) * k_m * rho.data[m, n]

def _mesolve_func_td(H_func, rho0, tlist, c_op_list, expt_ops, H_args, opt):
    """!
    Evolve the density matrix using an ODE solver with time dependent
    Hamiltonian.
    """

    n_op = len(c_op_list)
        
    #
    # check initial state
    #
    if isket(rho0):
        # if initial state is a ket and no collapse operator where given,
        # fallback on the unitary schrodinger equation solver
        if n_op == 0:
            return _wfsolve_list_td(H_func, rho0, tlist, expt_ops, H_args, opt)

        # Got a wave function as initial state: convert to density matrix.
        rho0 = ket2dm(rho0)

    #
    # construct liouvillian
    #
    L = 0
    for m in range(0, n_op):
        cdc = c_op_list[m].dag() * c_op_list[m]
        if L == 0:
            L = spre(c_op_list[m])*spost(c_op_list[m].dag())-0.5*spre(cdc)-0.5*spost(cdc)            
        else:
            L += spre(c_op_list[m])*spost(c_op_list[m].dag())-0.5*spre(cdc)-0.5*spost(cdc)

    if n_op > 0:
        L_func_and_args = [H_func, L.data]
    else:
        n,m = rho0.shape
        L_func_and_args = [H_func, sp.lil_matrix((n**2,m**2)).tocsr()]

    for arg in H_args:
        if isinstance(arg,Qobj):
            L_func_and_args.append((-1j*(spre(arg) - spost(arg))).data)
        else:
            L_func_and_args.append(arg)

    #
    # setup integrator
    #
    initial_vector = mat2vec(rho0.full())
    r = scipy.integrate.ode(_ode_rho_func_td)
    r.set_integrator('zvode', method=opt.method, order=opt.order,
                              atol=opt.atol, rtol=opt.rtol, nsteps=opt.nsteps,
                              first_step=opt.first_step, min_step=opt.min_step,
                              max_step=opt.max_step)
    r.set_initial_value(initial_vector, tlist[0])
    r.set_f_params(L_func_and_args)

    #
    # call generic ODE code
    #
    return _generic_ode_solve(r, rho0, tlist, expt_ops, opt, vec2mat)

def _mesolve_const(H, rho0, tlist, c_op_list, e_ops, args, opt,
                   progress_bar):
    """
    Evolve the density matrix using an ODE solver, for constant hamiltonian
    and collapse operators.
    """

    if debug:
        print(inspect.stack()[0][3])

    #
    # check initial state
    #
    if isket(rho0):
        # if initial state is a ket and no collapse operator where given,
        # fall back on the unitary schrodinger equation solver
        if len(c_op_list) == 0 and isoper(H):
            return _sesolve_const(H, rho0, tlist, e_ops, args, opt,
                                  progress_bar)

        # Got a wave function as initial state: convert to density matrix.
        rho0 = ket2dm(rho0)

    #
    # construct liouvillian
    #
    if opt.tidy:
        H = H.tidyup(opt.atol)

    L = liouvillian(H, c_op_list)
    

    #
    # setup integrator
    #
    initial_vector = mat2vec(rho0.full()).ravel('F')
    if issuper(rho0):
        r = scipy.integrate.ode(_ode_super_func)
        r.set_f_params(L.data)
    else:
        if opt.use_openmp and L.data.nnz >= qset.openmp_thresh:
            r = scipy.integrate.ode(cy_ode_rhs_openmp)
            r.set_f_params(L.data.data, L.data.indices, L.data.indptr, 
                            opt.openmp_threads)
        else:
            r = scipy.integrate.ode(cy_ode_rhs)
            r.set_f_params(L.data.data, L.data.indices, L.data.indptr)
        # r = scipy.integrate.ode(_ode_rho_test)
        # r.set_f_params(L.data)
    r.set_integrator('zvode', method=opt.method, order=opt.order,
                     atol=opt.atol, rtol=opt.rtol, nsteps=opt.nsteps,
                     first_step=opt.first_step, min_step=opt.min_step,
                     max_step=opt.max_step)
    r.set_initial_value(initial_vector, tlist[0])

    #
    # call generic ODE code
    #
    return _generic_ode_solve(r, rho0, tlist, e_ops, opt, progress_bar)

def plot_wigner_fock_distribution(rho, fig=None, axes=None, figsize=(8, 4),
                                  cmap=None, alpha_max=7.5, colorbar=False,
                                  method='iterative'):
    """
    Plot the Fock distribution and the Wigner function for a density matrix
    (or ket) that describes an oscillator mode.

    Parameters
    ----------
    rho : :class:`qutip.qobj.Qobj`
        The density matrix (or ket) of the state to visualize.

    fig : a matplotlib Figure instance
        The Figure canvas in which the plot will be drawn.

    axes : a list of two matplotlib axes instances
        The axes context in which the plot will be drawn.

    figsize : (width, height)
        The size of the matplotlib figure (in inches) if it is to be created
        (that is, if no 'fig' and 'ax' arguments are passed).

    cmap : a matplotlib cmap instance
        The colormap.

    alpha_max : float
        The span of the x and y coordinates (both [-alpha_max, alpha_max]).

    colorbar : bool
        Whether (True) or not (False) a colorbar should be attached to the
        Wigner function graph.

    method : string {'iterative', 'laguerre', 'fft'}
        The method used for calculating the wigner function. See the
        documentation for qutip.wigner for details.

    Returns
    -------
    fig, ax : tuple
        A tuple of the matplotlib figure and axes instances used to produce
        the figure.
    """

    if not fig and not axes:
        fig, axes = plt.subplots(1, 2, figsize=figsize)

    if isket(rho):
        rho = ket2dm(rho)

    plot_fock_distribution(rho, fig=fig, ax=axes[0])
    plot_wigner(rho, fig=fig, ax=axes[1], figsize=figsize, cmap=cmap,
                alpha_max=alpha_max, colorbar=colorbar, method=method)

    return fig, axes

def tracedist(A, B, sparse=False, tol=0):
    """
    Calculates the trace distance between two density matrices..
    See: Nielsen & Chuang, "Quantum Computation and Quantum Information"

    Parameters
    ----------!=
    A : qobj
        Density matrix or state vector.
    B : qobj
        Density matrix or state vector with same dimensions as A.
    tol : float
        Tolerance used by sparse eigensolver, if used. (0=Machine precision)
    sparse : {False, True}
        Use sparse eigensolver.

    Returns
    -------
    tracedist : float
        Trace distance between A and B.

    Examples
    --------
    >>> x=fock_dm(5,3)
    >>> y=coherent_dm(5,1)
    >>> tracedist(x,y)
    0.9705143161472971

    """
    if A.isket or A.isbra:
        A = ket2dm(A)
    if B.isket or B.isbra:
        B = ket2dm(B)

    if A.dims != B.dims:
        raise TypeError("A and B do not have same dimensions.")

    diff = A - B
    diff = diff.dag() * diff
    vals = sp_eigs(diff.data, diff.isherm, vecs=False, sparse=sparse, tol=tol)
    return float(np.real(0.5 * np.sum(np.sqrt(np.abs(vals)))))

def plot_fock_distribution(rho, offset=0, fig=None, ax=None,
                           figsize=(8, 6), title=None, unit_y_range=True):
    """
    Plot the Fock distribution for a density matrix (or ket) that describes
    an oscillator mode.

    Parameters
    ----------
    rho : :class:`qutip.qobj.Qobj`
        The density matrix (or ket) of the state to visualize.

    fig : a matplotlib Figure instance
        The Figure canvas in which the plot will be drawn.

    ax : a matplotlib axes instance
        The axes context in which the plot will be drawn.

    title : string
        An optional title for the figure.

    figsize : (width, height)
        The size of the matplotlib figure (in inches) if it is to be created
        (that is, if no 'fig' and 'ax' arguments are passed).

    Returns
    -------
    fig, ax : tuple
        A tuple of the matplotlib figure and axes instances used to produce
        the figure.
    """

    if not fig and not ax:
        fig, ax = plt.subplots(1, 1, figsize=figsize)

    if isket(rho):
        rho = ket2dm(rho)

    N = rho.shape[0]

    ax.bar(np.arange(offset, offset + N) - .4, np.real(rho.diag()),
           color="green", alpha=0.6, width=0.8)
    if unit_y_range:
        ax.set_ylim(0, 1)

    ax.set_xlim(-.5 + offset, N + offset)
    ax.set_xlabel('Fock number', fontsize=12)
    ax.set_ylabel('Occupation probability', fontsize=12)

    if title:
        ax.set_title(title)

    return fig, ax

def spin_q_function(rho, theta, phi):
    """Husimi Q-function for spins.

    Parameters
    ----------

    state : qobj
        A state vector or density matrix for a spin-j quantum system.

    theta : array_like
        theta-coordinates at which to calculate the Q function.

    phi : array_like
        phi-coordinates at which to calculate the Q function.

    Returns
    -------

    Q, THETA, PHI : 2d-array
        Values representing the spin Q function at the values specified
        by THETA and PHI.

    """

    if rho.type == 'bra':
        rho = rho.dag()

    if rho.type == 'ket':
        rho = ket2dm(rho)

    J = rho.shape[0]
    j = (J - 1) / 2

    THETA, PHI = meshgrid(theta, phi)

    Q = np.zeros_like(THETA, dtype=complex)

    for m1 in arange(-j, j+1):

        Q += binom(2*j, j+m1) * cos(THETA/2) ** (2*(j-m1)) * sin(THETA/2) ** (2*(j+m1)) * \
            rho.data[int(j-m1), int(j-m1)]

        for m2 in arange(m1+1, j+1):

            Q += (sqrt(binom(2*j, j+m1)) * sqrt(binom(2*j, j+m2)) *
                  cos(THETA/2) ** (2*j-m1-m2) * sin(THETA/2) ** (2*j+m1+m2)) * \
                (exp(1j * (m2-m1) * PHI) * rho.data[int(j-m1), int(j-m2)] +
                 exp(1j * (m1-m2) * PHI) * rho.data[int(j-m2), int(j-m1)])

    return Q.real, THETA, PHI