def test_QobjPermute():
    "Qobj permute"
    A = basis(5, 0)
    B = basis(5, 4)
    C = basis(5, 2)
    psi = tensor(A, B, C)
    psi2 = psi.permute([2, 0, 1])
    assert_(psi2 == tensor(C, A, B))

    A = fock_dm(5, 0)
    B = fock_dm(5, 4)
    C = fock_dm(5, 2)
    rho = tensor(A, B, C)
    rho2 = rho.permute([2, 0, 1])
    assert_(rho2 == tensor(C, A, B))

    for ii in range(3):
        A = rand_ket(5)
        B = rand_ket(5)
        C = rand_ket(5)
        psi = tensor(A, B, C)
        psi2 = psi.permute([1, 0, 2])
        assert_(psi2 == tensor(B, A, C))

    for ii in range(3):
        A = rand_dm(5)
        B = rand_dm(5)
        C = rand_dm(5)
        rho = tensor(A, B, C)
        rho2 = rho.permute([1, 0, 2])
        assert_(rho2 == tensor(B, A, C))

def test_QobjPermute():
    "Qobj permute"
    A = basis(3, 0)
    B = basis(5, 4)
    C = basis(4, 2)
    psi = tensor(A, B, C)
    psi2 = psi.permute([2, 0, 1])
    assert_(psi2 == tensor(C, A, B))
    
    psi_bra = psi.dag()
    psi2_bra = psi_bra.permute([2, 0, 1])
    assert_(psi2_bra == tensor(C, A, B).dag())

    A = fock_dm(3, 0)
    B = fock_dm(5, 4)
    C = fock_dm(4, 2)
    rho = tensor(A, B, C)
    rho2 = rho.permute([2, 0, 1])
    assert_(rho2 == tensor(C, A, B))

    for ii in range(3):
        A = rand_ket(3)
        B = rand_ket(4)
        C = rand_ket(5)
        psi = tensor(A, B, C)
        psi2 = psi.permute([1, 0, 2])
        assert_(psi2 == tensor(B, A, C))
        
        psi_bra = psi.dag()
        psi2_bra = psi_bra.permute([1, 0, 2])
        assert_(psi2_bra == tensor(B, A, C).dag())

    for ii in range(3):
        A = rand_dm(3)
        B = rand_dm(4)
        C = rand_dm(5)
        rho = tensor(A, B, C)
        rho2 = rho.permute([1, 0, 2])
        assert_(rho2 == tensor(B, A, C))
        
        rho_vec = operator_to_vector(rho)
        rho2_vec = rho_vec.permute([[1, 0, 2],[4,3,5]])
        assert_(rho2_vec == operator_to_vector(tensor(B, A, C)))
        
        rho_vec_bra = operator_to_vector(rho).dag()
        rho2_vec_bra = rho_vec_bra.permute([[1, 0, 2],[4,3,5]])
        assert_(rho2_vec_bra == operator_to_vector(tensor(B, A, C)).dag())
        
    for ii in range(3):
        super_dims = [3, 5, 4]
        U = rand_unitary(np.prod(super_dims), density=0.02, dims=[super_dims, super_dims])
        Unew = U.permute([2,1,0])
        S_tens = to_super(U)
        S_tens_new = to_super(Unew)
        assert_(S_tens_new == S_tens.permute([[2,1,0],[5,4,3]]))

def test_csr_kron():
    "spmath: zcsr_kron"
    for kk in range(10):
        ra = np.random.randint(2,100)
        rb = np.random.randint(2,100)
        A = rand_herm(ra,0.5).data
        B = rand_herm(rb,0.5).data
        As = A.tocsr(1)
        Bs = B.tocsr(1)
        C = sp.kron(As,Bs, format='csr')
        D = zcsr_kron(A, B)
        assert_almost_equal(C.data, D.data)
        assert_equal(C.indices, D.indices)
        assert_equal(C.indptr, D.indptr)
        
    for kk in range(10):
        ra = np.random.randint(2,100)
        rb = np.random.randint(2,100)
        A = rand_ket(ra,0.5).data
        B = rand_herm(rb,0.5).data
        As = A.tocsr(1)
        Bs = B.tocsr(1)
        C = sp.kron(As,Bs, format='csr')
        D = zcsr_kron(A, B)
        assert_almost_equal(C.data, D.data)
        assert_equal(C.indices, D.indices)
        assert_equal(C.indptr, D.indptr)
    
    for kk in range(10):
        ra = np.random.randint(2,100)
        rb = np.random.randint(2,100)
        A = rand_dm(ra,0.5).data
        B = rand_herm(rb,0.5).data
        As = A.tocsr(1)
        Bs = B.tocsr(1)
        C = sp.kron(As,Bs, format='csr')
        D = zcsr_kron(A, B)
        assert_almost_equal(C.data, D.data)
        assert_equal(C.indices, D.indices)
        assert_equal(C.indptr, D.indptr)
        
    for kk in range(10):
        ra = np.random.randint(2,100)
        rb = np.random.randint(2,100)
        A = rand_ket(ra,0.5).data
        B = rand_ket(rb,0.5).data
        As = A.tocsr(1)
        Bs = B.tocsr(1)
        C = sp.kron(As,Bs, format='csr')
        D = zcsr_kron(A, B)
        assert_almost_equal(C.data, D.data)
        assert_equal(C.indices, D.indices)
        assert_equal(C.indptr, D.indptr)

def test_csr_kron():
    "Sparse: Test CSR Kron"
    for kk in range(10):
        ra = np.random.randint(2,100)
        rb = np.random.randint(2,100)
        A = rand_herm(ra,0.5).data
        B = rand_herm(rb,0.5).data
        C = sp.kron(A,B, format='csr')
        D = _csr_kron(A.data,A.indices,A.indptr, A.shape[0], A.shape[1],
                    B.data,B.indices,B.indptr, B.shape[0], B.shape[1])
        assert_equal(C.data, D.data)
        assert_equal(C.indices, D.indices)
        assert_equal(C.indptr, D.indptr)
        
    for kk in range(10):
        ra = np.random.randint(2,100)
        rb = np.random.randint(2,100)
        A = rand_ket(ra,0.5).data
        B = rand_herm(rb,0.5).data
        C = sp.kron(A,B, format='csr')
        D = _csr_kron(A.data,A.indices,A.indptr, A.shape[0], A.shape[1],
                    B.data,B.indices,B.indptr, B.shape[0], B.shape[1])
        assert_equal(C.data, D.data)
        assert_equal(C.indices, D.indices)
        assert_equal(C.indptr, D.indptr)
    
    for kk in range(10):
        ra = np.random.randint(2,100)
        rb = np.random.randint(2,100)
        A = rand_dm(ra,0.5).data
        B = rand_herm(rb,0.5).data
        C = sp.kron(A,B, format='csr')
        D = _csr_kron(A.data,A.indices,A.indptr, A.shape[0], A.shape[1],
                    B.data,B.indices,B.indptr, B.shape[0], B.shape[1])
        assert_equal(C.data, D.data)
        assert_equal(C.indices, D.indices)
        assert_equal(C.indptr, D.indptr)
        
    for kk in range(10):
        ra = np.random.randint(2,100)
        rb = np.random.randint(2,100)
        A = rand_ket(ra,0.5).data
        B = rand_ket(rb,0.5).data
        C = sp.kron(A,B, format='csr')
        D = _csr_kron(A.data,A.indices,A.indptr, A.shape[0], A.shape[1],
                    B.data,B.indices,B.indptr, B.shape[0], B.shape[1])
        assert_equal(C.data, D.data)
        assert_equal(C.indices, D.indices)
        assert_equal(C.indptr, D.indptr)

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 test_Qobj_Spmv():
    "Qobj mult ndarray right"
    A = rand_herm(5)
    b = rand_ket(5).full()
    C = A*b
    D = A.full().dot(b)
    assert_(np.all((C-D)<1e-14))

def test_projection():
    """
    Test Qobj: Projection operator
    """
    for kk in range(10):
        N = 5
        K = tensor(rand_ket(N,0.75),rand_ket(N,0.75))
        B = K.dag()
        
        ans = K*K.dag()
        
        out1 = K.proj()
        out2 = B.proj()
        
        assert_(out1==ans)
        assert_(out2==ans)

def test_wigner_compare_methods_ket():
    "wigner: compare wigner methods for random state vectors"

    xvec = np.linspace(-5.0, 5.0, 100)
    yvec = xvec

    X, Y = np.meshgrid(xvec, yvec)

    # a = X + 1j * Y  # consistent with g=2 option to wigner function

    dx = xvec[1] - xvec[0]
    dy = yvec[1] - yvec[0]

    N = 15

    for n in range(10):
        # try ten different random density matrices

        psi = rand_ket(N, 0.5 + rand() / 2)

        # calculate the wigner function using qutip and analytic formula
        W_qutip1 = wigner(psi, xvec, yvec, g=2)
        W_qutip2 = wigner(psi, xvec, yvec, g=2, method='laguerre')

        # check difference
        assert_(np.sum(abs(W_qutip1 - W_qutip1)) < 1e-4)

        # check normalization
        assert_(np.sum(W_qutip1) * dx * dy - 1.0 < 1e-8)
        assert_(np.sum(W_qutip2) * dx * dy - 1.0 < 1e-8)

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 testExpandGate2toNSwap(self):
        """
        gates: expand 2 to N (using swap)
        """

        a, b = np.random.rand(), np.random.rand()
        k1 = (a * basis(2, 0) + b * basis(2, 1)).unit()

        c, d = np.random.rand(), np.random.rand()
        k2 = (c * basis(2, 0) + d * basis(2, 1)).unit()

        N = 6
        kets = [rand_ket(2) for k in range(N)]

        for m in range(N):
            for n in set(range(N)) - {m}:

                psi_in = tensor([k1 if k == m else k2 if k == n else kets[k]
                                 for k in range(N)])
                psi_out = tensor([k2 if k == m else k1 if k == n else kets[k]
                                  for k in range(N)])

                targets = [m, n]
                G = swap(N, targets)

                psi_out = G * psi_in

                assert_((psi_out - G * psi_in).norm() < 1e-12)

    def testExpandGate1toN(self):
        """
        gates: expand 1 to N
        """
        N = 7

        for g in [rx, ry, rz, phasegate]:

            theta = np.random.rand() * 2 * 3.1415
            a, b = np.random.rand(), np.random.rand()
            psi1 = (a * basis(2, 0) + b * basis(2, 1)).unit()
            psi2 = g(theta) * psi1

            psi_rand_list = [rand_ket(2) for k in range(N)]

            for m in range(N):

                psi_in = tensor([psi1 if k == m else psi_rand_list[k]
                                 for k in range(N)])
                psi_out = tensor([psi2 if k == m else psi_rand_list[k]
                                  for k in range(N)])

                G = g(theta, N, m)
                psi_res = G * psi_in

                assert_((psi_out - psi_res).norm() < 1e-12)

def test_overlap():
    """
    Test Qobj: Overlap (inner product)
    """
    for kk in range(10):
        N = 10
        A = rand_ket(N,0.75)
        Ad = A.dag()
        B = rand_ket(N,0.75)
        Bd = B.dag()
        
        ans = (A.dag()*B).tr()
        
        assert_almost_equal(A.overlap(B), ans)
        assert_almost_equal(Ad.overlap(B), ans)
        assert_almost_equal(Ad.overlap(Bd), ans)
        assert_almost_equal(A.overlap(Bd), np.conj(ans))

def test_tracedist3():
    """
    Metrics: Trace dist. & Fidelity mixed/pure inequality
    """
    for k in range(10):
        ket = rand_ket(25, 0.25)
        rho1 = ket*ket.dag()
        rho2 = rand_dm(25, 0.25)
        F = fidelity(rho1, rho2)
        D = tracedist(rho1, rho2)
        assert_(1-F**2 <= D)

def test_wigner_fft_comparse_ket():
    "Wigner: Compare Wigner fft and iterative for rand. ket"
    N = 20
    xvec = np.linspace(-10, 10, 128)
    for i in range(3):
        rho = rand_ket(N)

        Wfft, yvec = wigner(rho, xvec, xvec, method='fft')
        W = wigner(rho, xvec, yvec, method='iterative')

        Wdiff = abs(W - Wfft)
        assert_equal(np.sum(abs(Wdiff)) < 1e-7, True)

def test_matelem():
    """
    Test Qobj: Compute matrix elements
    """
    for kk in range(10):
        N = 20
        H = rand_herm(N,0.2)

        L = rand_ket(N,0.3)
        Ld = L.dag()
        R = rand_ket(N,0.3)
    
        ans = (Ld*H*R).tr()
    
        #bra-ket
        out1 = H.matrix_element(Ld,R)
        #ket-ket
        out2 = H.matrix_element(Ld,R)
    
        assert_(abs(ans-out1) < 1e-14)
        assert_(abs(ans-out2) < 1e-14)

    def testExpandGate3toN(self):
        """
        gates: expand 3 to N (using toffoli, fredkin, and random 3 qubit gate)
        """

        a, b = np.random.rand(), np.random.rand()
        psi1 = (a * basis(2, 0) + b * basis(2, 1)).unit()

        c, d = np.random.rand(), np.random.rand()
        psi2 = (c * basis(2, 0) + d * basis(2, 1)).unit()

        e, f = np.random.rand(), np.random.rand()
        psi3 = (e * basis(2, 0) + f * basis(2, 1)).unit()

        N = 4
        psi_rand_list = [rand_ket(2) for k in range(N)]

        _rand_gate_U = tensor([rand_herm(2, density=1) for k in range(3)])

        def _rand_3qubit_gate(N=None, controls=None, k=None):
            if N is None:
                return _rand_gate_U
            else:
                return gate_expand_3toN(_rand_gate_U, N, controls, k)

        for g in [fredkin, toffoli, _rand_3qubit_gate]:

            psi_ref_in = tensor(psi1, psi2, psi3)
            psi_ref_out = g() * psi_ref_in

            for m in range(N):
                for n in set(range(N)) - {m}:
                    for k in set(range(N)) - {m, n}:

                        psi_list = [psi_rand_list[p] for p in range(N)]
                        psi_list[m] = psi1
                        psi_list[n] = psi2
                        psi_list[k] = psi3
                        psi_in = tensor(psi_list)

                        if g == fredkin:
                            targets = [n, k]
                            G = g(N, control=m, targets=targets)
                        else:
                            controls = [m, n]
                            G = g(N, controls, k)

                        psi_out = G * psi_in

                        o1 = psi_out.overlap(psi_in)
                        o2 = psi_ref_out.overlap(psi_ref_in)
                        assert_(abs(o1 - o2) < 1e-12)

    def testExpandGate2toN(self):
        """
        gates: expand 2 to N (using cnot, iswap, sqrtswap)
        """

        a, b = np.random.rand(), np.random.rand()
        k1 = (a * basis(2, 0) + b * basis(2, 1)).unit()

        c, d = np.random.rand(), np.random.rand()
        k2 = (c * basis(2, 0) + d * basis(2, 1)).unit()

        psi_ref_in = tensor(k1, k2)

        N = 6
        psi_rand_list = [rand_ket(2) for k in range(N)]

        for g in [cnot, iswap, sqrtswap]:

            psi_ref_out = g() * psi_ref_in
            rho_ref_out = ket2dm(psi_ref_out)

            for m in range(N):
                for n in set(range(N)) - {m}:

                    psi_list = [psi_rand_list[k] for k in range(N)]
                    psi_list[m] = k1
                    psi_list[n] = k2
                    psi_in = tensor(psi_list)

                    if g == cnot:
                        G = g(N, m, n)
                    else:
                        G = g(N, [m, n])

                    psi_out = G * psi_in

                    o1 = psi_out.overlap(psi_in)
                    o2 = psi_ref_out.overlap(psi_ref_in)
                    assert_(abs(o1 - o2) < 1e-12)

                    p = [0, 1] if m < n else [1, 0]
                    rho_out = psi_out.ptrace([m, n]).permute(p)

                    assert_((rho_ref_out - rho_out).norm() < 1e-12)

def test_zcsr_transpose():
    "spmath: zcsr_transpose"
    for k in range(50):
        ra = np.random.randint(2,100)
        A = rand_ket(ra,0.5).data
        B = A.T.tocsr()
        C = A.trans()
        x = np.all(B.data == C.data)
        y = np.all(B.indices == C.indices)
        z = np.all(B.indptr == C.indptr)
        assert_(x*y*z)
    
    for k in range(50):
        ra = np.random.randint(2,100)
        A = rand_herm(5,1.0/ra).data
        B = A.T.tocsr()
        C = A.trans()
        x = np.all(B.data == C.data)
        y = np.all(B.indices == C.indices)
        z = np.all(B.indptr == C.indptr)
        assert_(x*y*z)
        
    for k in range(50):
        ra = np.random.randint(2,100)
        A = rand_dm(ra,1.0/ra).data
        B = A.T.tocsr()
        C = A.trans()
        x = np.all(B.data == C.data)
        y = np.all(B.indices == C.indices)
        z = np.all(B.indptr == C.indptr)
        assert_(x*y*z)
        
    for k in range(50):
        ra = np.random.randint(2,100)
        A = rand_unitary(ra,1.0/ra).data
        B = A.T.tocsr()
        C = A.trans()
        x = np.all(B.data == C.data)
        y = np.all(B.indices == C.indices)
        z = np.all(B.indptr == C.indptr)
        assert_(x*y*z)

def system_bench(func, dims):
    from qutip.random_objects import rand_ket
    ratio = 0
    ratio_old = 0
    nnz_old = 0
    for N in dims:
        L = func(N).data
        vec = rand_ket(L.shape[0],0.25).full().ravel()
        nnz = L.nnz
        out = np.zeros_like(vec)
        ser = _min_timer(_spmvpy, L.data, L.indices, L.indptr, vec, 1, out)
        out = np.zeros_like(vec)
        par = _min_timer(_spmvpy_openmp, L.data, L.indices, L.indptr, vec, 1, out, 2)
        ratio = ser/par
        if ratio > 1:
            break
        nnz_old = nnz 
        ratio_old = ratio  
    if ratio > 1:
        rate = (ratio-ratio_old)/(nnz-nnz_old)
        return int((1.0-ratio_old)/rate+nnz_old)
    else:
        return -1

def test_zcsr_mult():
    "spmath: zcsr_mult"
    for k in range(50):
        A = rand_ket(10,0.5).data
        B = rand_herm(10,0.5).data
        
        C = A.tocsr(1)
        D = B.tocsr(1)
        
        ans1 = B*A
        ans2 = D*C
        ans2.sort_indices()
        x = np.all(ans1.data == ans2.data)
        y = np.all(ans1.indices == ans2.indices)
        z = np.all(ans1.indptr == ans2.indptr)
        assert_(x*y*z)
        
    for k in range(50):
        A = rand_ket(10,0.5).data
        B = rand_ket(10,0.5).dag().data
        
        C = A.tocsr(1)
        D = B.tocsr(1)
        
        ans1 = B*A
        ans2 = D*C
        ans2.sort_indices()
        x = np.all(ans1.data == ans2.data)
        y = np.all(ans1.indices == ans2.indices)
        z = np.all(ans1.indptr == ans2.indptr)
        assert_(x*y*z)
        
        ans1 = A*B
        ans2 = C*D
        ans2.sort_indices()
        x = np.all(ans1.data == ans2.data)
        y = np.all(ans1.indices == ans2.indices)
        z = np.all(ans1.indptr == ans2.indptr)
        assert_(x*y*z)
        
    for k in range(50):
        A = rand_dm(10,0.5).data
        B = rand_dm(10,0.5).data
        
        C = A.tocsr(1)
        D = B.tocsr(1)
        
        ans1 = B*A
        ans2 = D*C
        ans2.sort_indices()
        x = np.all(ans1.data == ans2.data)
        y = np.all(ans1.indices == ans2.indices)
        z = np.all(ans1.indptr == ans2.indptr)
        assert_(x*y*z)
        
    for k in range(50):
        A = rand_dm(10,0.5).data
        B = rand_herm(10,0.5).data
        
        C = A.tocsr(1)
        D = B.tocsr(1)
        
        ans1 = B*A
        ans2 = D*C
        ans2.sort_indices()
        x = np.all(ans1.data == ans2.data)
        y = np.all(ans1.indices == ans2.indices)
        z = np.all(ans1.indptr == ans2.indptr)
        assert_(x*y*z)