def single_hamiltonian(cav_dim, w_1, w_c, g_factor):
    """Return a QObj denoting a hamiltonian for one qubit coupled to a
    cavity."""
    return (w_c * qt.tensor(qt.num(cav_dim), I) +
            0.5 * w_1 * qt.tensor(qt.qeye(cav_dim), I) +
            g_factor * qt.tensor(qt.create(cav_dim), SMinus) +
            g_factor * qt.tensor(qt.destroy(cav_dim), SPlus))

    def __init__(self, N_field_levels, coupling=None, N_qubits=1):

        # basic parameters
        self.N_field_levels = N_field_levels
        self.N_qubits = N_qubits

        if coupling is None:
            self.g = 0
        else:
            self.g = coupling

        # bare operators
        self.idcavity = qt.qeye(self.N_field_levels)
        self.idqubit = qt.qeye(2)
        self.a_bare = qt.destroy(self.N_field_levels)
        self.sm_bare = qt.sigmam()
        self.sz_bare = qt.sigmaz()
        self.sx_bare = qt.sigmax()
        self.sy_bare = qt.sigmay()

        # 1 atom 1 cavity operators
        self.jc_a = qt.tensor(self.a_bare, self.idqubit)
        self.jc_sm = qt.tensor(self.idcavity, self.sm_bare)
        self.jc_sx = qt.tensor(self.idcavity, self.sx_bare)
        self.jc_sy = qt.tensor(self.idcavity, self.sy_bare)
        self.jc_sz = qt.tensor(self.idcavity, self.sz_bare)

 def construct_hamiltonian(self, number_of_spins, alpha, B):
     '''
     following example
     http://qutip.googlecode.com/svn/doc/2.0.0/html/examples/me/ex-24.html
     '''
     N = number_of_spins
     si = qeye(2)
     sx = sigmax()
     sy = sigmay()
     #constructing a list of operators sx_list and sy_list where
     #the operator sx_list[i] applies sigma_x on the ith particle and 
     #identity to the rest
     sx_list = []
     sy_list = []
     for n in range(N):
         op_list = []
         for m in range(N):
             op_list.append(si)
         op_list[n] = sx
         sx_list.append(tensor(op_list))
         op_list[n] = sy
         sy_list.append(tensor(op_list))
     #construct the hamiltonian
     H = 0
     #magnetic field term, hamiltonian is in units of J0
     for i in range(N):
         H-= B * sy_list[i]
     #ising coupling term
     for i in range(N):
         for j in range(N):
             if i < j:
                 H+= abs(i - j)**-alpha * sx_list[i] * sx_list[j]
     return H

def test_spectrum_esfft():
    """
    correlation: comparing spectrum from es and fft methods
    """

    # use JC model
    N = 4
    wc = wa = 1.0 * 2 * np.pi
    g = 0.1 * 2 * np.pi
    kappa = 0.75
    gamma = 0.25
    n_th = 0.01

    a = tensor(destroy(N), qeye(2))
    sm = tensor(qeye(N), destroy(2))
    H = wc * a.dag() * a + wa * sm.dag() * sm + \
        g * (a.dag() * sm + a * sm.dag())
    c_ops = [np.sqrt(kappa * (1 + n_th)) * a,
             np.sqrt(kappa * n_th) * a.dag(),
             np.sqrt(gamma) * sm]

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        tlist = np.linspace(0, 100, 2500)
        corr = correlation_ss(H, tlist, c_ops, a.dag(), a)
        wlist1, spec1 = spectrum_correlation_fft(tlist, corr)
        spec2 = spectrum_ss(H, wlist1, c_ops, a.dag(), a)

    assert_(max(abs(spec1 - spec2)) < 1e-3)

def test_enr_destory_full():
    "Excitation-number-restricted state-space: full state space"
    a1, a2 = enr_destroy([4, 4], 4**2)
    b1, b2 = tensor(destroy(4), identity(4)), tensor(identity(4), destroy(4))

    assert_(a1 == b1)
    assert_(a2 == b2)

def rsb_flop(rho0, W, eta, delta, theta, phi, c_op_list = [], return_op_list = []):
    ''' Return values of atom and motion populations during red sideband Rabi flop 
    for rotation angles theta. Calls numerical solution of master equation for the 
    Jaynes-Cummings Hamiltonian.
    @ var rho0: initial density matrix
    @ var W: bare Rabi frequency
    @ var delta: detuning between atom and motion
    @ var theta: list of Rabi rotation angle (i.e. theta, or g*time)
    @ var phi: phase of the input laser pulse
    @ var c_op_list: list of collapse operators for the master equation treatment
    @ var return_op_list: list of population operators the values of which will be returned 
    
    returns: time, populations of motional mode and atom
    '''
    N = shape(rho0.data)[0]/2 # assume N Fock states and two atom states
    a = tensor(destroy(N), qeye(2))
    sm = tensor( qeye(N), destroy(2))
    Wrsb = destroy(N)
    one_then_zero = ([float(x<1) for x in range(N)])
    Wrsb.data = csr_matrix( destroy(N).data.dot( np.diag( rabi_coupling(N,-1,eta) / np.sqrt(one_then_zero+np.linspace(0,N-1,N)) ) ) ) 
    Arsb = tensor(Wrsb, qeye(2))
    # use the rotating wave approxiation
    # Note that the regular a, a.dag() is used for the time evolution of the oscillator
    # Arsb is the destruction operator including the state dependent coupling strength
    H = delta * a.dag() * a + \
        (1./2.) * W * (Arsb.dag() * sm * exp(1j*phi) + Arsb * sm.dag() * exp(-1j*phi))
    
    if hasattr(theta, '__len__'):
        if len(theta)>1: # I need to be able to pass a list of length zero and not get an error
            time = theta/(eta*W)
    else:
            time = theta/(eta*W)        

    output = mesolve(H, rho0, time, c_op_list, return_op_list)
    return time, output

def testJCZeroTemperature():
    """
    brmesolve: Jaynes-Cummings model, zero temperature
    """

    N = 10
    a = tensor(destroy(N), identity(2))
    sm = tensor(identity(N), destroy(2))
    psi0 = ket2dm(tensor(basis(N, 1), basis(2, 0)))
    a_ops = [(a + a.dag())]
    e_ops = [a.dag() * a, sm.dag() * sm]

    w0 = 1.0 * 2 * np.pi
    g = 0.05 * 2 * np.pi
    kappa = 0.05
    times = np.linspace(0, 2 * 2 * np.pi / g, 1000)

    c_ops = [np.sqrt(kappa) * a]
    H = w0 * a.dag() * a + w0 * sm.dag() * sm + \
        g * (a + a.dag()) * (sm + sm.dag())

    res_me = mesolve(H, psi0, times, c_ops, e_ops)
    res_brme = brmesolve(H, psi0, times, a_ops, e_ops,
                         spectra_cb=[lambda w: kappa * (w >= 0)])

    for idx, e in enumerate(e_ops):
        diff = abs(res_me.expect[idx] - res_brme.expect[idx]).max()
        assert_(diff < 5e-2)  # accept 5% error

def collapse_operators(N, n_th_a, gamma_motion, gamma_motion_phi, gamma_atom):
    '''Collapse operators for the master equation of a single atom and a harmonic oscillator
    @ var N: size of the harmonic oscillator Hilbert space
    @ var n_th: temperature of the noise bath in quanta
    @ var gamma_motion: heating rate of the motion
    @ var gamma_motion_phi: dephasing rate of the motion
    @ var gamma_atom: decay rate of the atom
    
    returns: list of collapse operators for master equation solution of atom + harmonic oscillator
    '''
    a = tensor(destroy(N), qeye(2))
    sm = tensor(qeye(N), destroy(2))  
    c_op_list = []

    rate = gamma_motion * (1 + n_th_a)
    if rate > 0.0:
        c_op_list.append(sqrt(rate) * a)

    rate = gamma_motion * n_th_a
    if rate > 0.0:
        c_op_list.append(sqrt(rate) * a.dag())
    
    rate = gamma_motion_phi
    if rate > 0.0:
        c_op_list.append(sqrt(rate) * a.dag() * a)

    rate = gamma_atom
    if rate > 0.0:
        c_op_list.append(sqrt(rate) * sm)
    return c_op_list

def carrier_flop(rho0, W, eta, delta, theta, phi, c_op_list = [], return_op_list = []):
    ''' Return values of atom and motion populations during carrier Rabi flop 
    for rotation angles theta. Calls numerical solution of master equation.
    @ var rho0: initial density matrix
    @ var W: bare Rabi frequency
    @ var eta: Lamb-Dicke parameter
    @ var delta: detuning between atom and motion
    @ var theta: list of Rabi rotation angles (i.e. theta, or g*time)
    @ var phi: phase of the input laser pulse
    @ var c_op_list: list of collapse operators for the master equation treatment
    @ var return_op_list: list of population operators the values of which will be returned 
    
    returns: time, populations of motional mode and atom
    '''    
    N = shape(rho0.data)[0]/2 # assume N Fock states and two atom states
    a = tensor(destroy(N), qeye(2))
    Wc = qeye(N)
    Wc.data = csr_matrix( qeye(N).data.dot( np.diag(rabi_coupling(N,0,eta) ) ) )    
    sm = tensor( Wc, destroy(2))

    # use the rotating wave approxiation
    H = delta * a.dag() * a + \
         (1./2.)* W * (sm.dag()*exp(1j*phi) + sm*exp(-1j*phi))
    
    if hasattr(theta, '__len__'):
        if len(theta)>1: # I need to be able to pass a list of length zero and not get an error
            time = theta/W
    else:
            time = theta/W        

    output = mesolve(H, rho0, time, c_op_list, return_op_list)
    return time, output

def M_ops(N,j,i):
    k = Kspace(N)
    op = k[i]*k[j].dag()
    if i==j+1 or i==j-1:
        return qutip.tensor(B_ops(N,j,i),op)
    if i==j and (j==0 or j==N-1):
        return qutip.tensor(B_ops(N,j,i),op)

    def compute(N, wc, wa, glist, use_rwa):

        # Pre-compute operators for the hamiltonian
        a  = tensor(destroy(N), qeye(2))
        sm = tensor(qeye(N), destroy(2))
        nc = a.dag() * a
        na = sm.dag() * sm

        idx = 0
        na_expt = zeros(shape(glist))
        nc_expt = zeros(shape(glist))
        for g in glist:

            # recalculate the hamiltonian for each value of g
            if use_rwa:
                H = wc * nc + wa * na + g * (a.dag() * sm + a * sm.dag())
            else:
                H = wc * nc + wa * na + g * (a.dag() + a) * (sm + sm.dag())

            # find the groundstate of the composite system
            evals, ekets = H.eigenstates()
            psi_gnd = ekets[0]
            na_expt[idx] = expect(na, psi_gnd)
            nc_expt[idx] = expect(nc, psi_gnd)

            idx += 1

        return nc_expt, na_expt, ket2dm(psi_gnd)

def test_spectrum_espi():
    """
    correlation: comparing spectrum from es and pi methods
    """

    # use JC model
    N = 4
    wc = wa = 1.0 * 2 * np.pi
    g = 0.1 * 2 * np.pi
    kappa = 0.75
    gamma = 0.25
    n_th = 0.01

    a = tensor(destroy(N), qeye(2))
    sm = tensor(qeye(N), destroy(2))
    H = wc * a.dag() * a + wa * sm.dag() * sm + \
        g * (a.dag() * sm + a * sm.dag())
    c_ops = [np.sqrt(kappa * (1 + n_th)) * a,
             np.sqrt(kappa * n_th) * a.dag(),
             np.sqrt(gamma) * sm]

    wlist = 2 * pi * np.linspace(0.5, 1.5, 100)
    spec1 = spectrum(H, wlist, c_ops, a.dag(), a, solver='es')
    spec2 = spectrum(H, wlist, c_ops, a.dag(), a, solver='pi')

    assert_(max(abs(spec1 - spec2)) < 1e-3)

    def to_matrix(self, fd):
        n = num(fd)
        a = destroy(fd)
        ic = qeye(fd)
        sz = sigmaz()
        sm = sigmam()
        iq = qeye(2)

        ms = {
            "id": tensor(iq, ic),
            "a*ad" : tensor(iq, n),
            "a+hc" : tensor(iq, a),
            "sz" : tensor(sz, ic),
            "sm+hc" : tensor(sm, ic)
        }

        H0 = 0
        H1s = []
        for (p1, p2), v in self.coefs.items():
            h = ms[p1] * ms[p2]
            try:
                term = float(v) * h
                if not term.isherm:
                    term += term.dag()
                H0 += term
            except ValueError:
                H1s.append([h, v])
                if not h.isherm:
                    replacement = lambda m: '(-' + m.group() + ')'
                    conj_v = re.sub('[1-9]+j', replacement, v)
                    H1s.append([h.dag(), conj_v])
        if H1s:
            return [H0] + H1s
        else:
            return H0

    def construct_hamiltonian(self, number_of_spins, alpha):
        """
following example
http://qutip.googlecode.com/svn/doc/2.0.0/html/examples/me/ex-24.html
returns H0 - hamiltonian without the B field
and y_list - list of sigma_y operators
"""
        N = number_of_spins
        si = qeye(2)
        sx = sigmax()
        sy = sigmay()
        # constructing a list of operators sx_list and sy_list where
        # the operator sx_list[i] applies sigma_x on the ith particle and
        # identity to the rest
        sx_list = []
        sy_list = []
        for n in range(N):
            op_list = []
            for m in range(N):
                op_list.append(si)
            op_list[n] = sx
            sx_list.append(tensor(op_list))
            op_list[n] = sy
            sy_list.append(tensor(op_list))
        # construct the hamiltonian
        H0 = 0
        # ising coupling term, time independent
        for i in range(N):
            for j in range(N):
                if i < j:
                    H0 -= abs(i - j) ** -alpha * sx_list[i] * sx_list[j]
        H1 = 0
        for i in range(N):
            H1 -= sy_list[i]
        return H0, H1

def test_spectrum_espi_legacy():
    """
    correlation: legacy spectrum from es and pi methods
    """

    # use JC model
    N = 4
    wc = wa = 1.0 * 2 * np.pi
    g = 0.1 * 2 * np.pi
    kappa = 0.75
    gamma = 0.25
    n_th = 0.01

    a = tensor(destroy(N), qeye(2))
    sm = tensor(qeye(N), destroy(2))
    H = wc * a.dag() * a + wa * sm.dag() * sm + \
        g * (a.dag() * sm + a * sm.dag())
    c_ops = [np.sqrt(kappa * (1 + n_th)) * a,
             np.sqrt(kappa * n_th) * a.dag(),
             np.sqrt(gamma) * sm]

    wlist = 2 * np.pi * np.linspace(0.5, 1.5, 100)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        spec1 = spectrum_ss(H, wlist, c_ops, a.dag(), a)
        spec2 = spectrum_pi(H, wlist, c_ops, a.dag(), a)

    assert_(max(abs(spec1 - spec2)) < 1e-3)

def test_mc_dtypes2():
    "Monte-carlo: check for correct dtypes (average_states=False)"
    # set system parameters
    kappa = 2.0  # mirror coupling
    gamma = 0.2  # spontaneous emission rate
    g = 1  # atom/cavity coupling strength
    wc = 0  # cavity frequency
    w0 = 0  # atom frequency
    wl = 0  # driving frequency
    E = 0.5  # driving amplitude
    N = 5  # number of cavity energy levels (0->3 Fock states)
    tlist = np.linspace(0, 10, 5)  # times for expectation values
    # construct Hamiltonian
    ida = qeye(N)
    idatom = qeye(2)
    a = tensor(destroy(N), idatom)
    sm = tensor(ida, sigmam())
    H = (w0 - wl) * sm.dag() * sm + (wc - wl) * a.dag() * a + \
        1j * g * (a.dag() * sm - sm.dag() * a) + E * (a.dag() + a)
    # collapse operators
    C1 = np.sqrt(2 * kappa) * a
    C2 = np.sqrt(gamma) * sm
    C1dC1 = C1.dag() * C1
    C2dC2 = C2.dag() * C2
    # intial state
    psi0 = tensor(basis(N, 0), basis(2, 1))
    opts = Options(average_expect=False)
    data = mcsolve(
        H, psi0, tlist, [C1, C2], [C1dC1, C2dC2, a], ntraj=5, options=opts)
    assert_equal(isinstance(data.expect[0][0][1], float), True)
    assert_equal(isinstance(data.expect[0][1][1], float), True)
    assert_equal(isinstance(data.expect[0][2][1], complex), True)

def test_diagHamiltonian2():
    """
    Diagonalization of composite systems
    """

    H1 = scipy.rand() * sigmax() + scipy.rand() * sigmay() +\
        scipy.rand() * sigmaz()
    H2 = scipy.rand() * sigmax() + scipy.rand() * sigmay() +\
        scipy.rand() * sigmaz()

    H = tensor(H1, H2)

    evals, ekets = H.eigenstates()

    for n in range(len(evals)):
        # assert that max(H * ket - e * ket) is small
        assert_equal(amax(
            abs((H * ekets[n] - evals[n] * ekets[n]).full())) < 1e-10, True)

    N1 = 10
    N2 = 2

    a1 = tensor(destroy(N1), qeye(N2))
    a2 = tensor(qeye(N1), destroy(N2))
    H = scipy.rand() * a1.dag() * a1 + scipy.rand() * a2.dag() * a2 + \
        scipy.rand() * (a1 + a1.dag()) * (a2 + a2.dag())
    evals, ekets = H.eigenstates()

    for n in range(len(evals)):
        # assert that max(H * ket - e * ket) is small
        assert_equal(amax(
            abs((H * ekets[n] - evals[n] * ekets[n]).full())) < 1e-10, True)

def direct_hamiltonian(w_1, w_2, g_factor):
    """Return a QObj denoting a hamiltonian for two qubits interacting with the
    cavity mode eliminated."""
    return (
            0.5 * w_1 * qt.tensor(SZ, I) +
            0.5 * w_2 * qt.tensor(I, SZ) +
            g_factor * qt.tensor(SPlus, SMinus) +
            g_factor * qt.tensor(SMinus, SPlus))

def population_operators(N):
    '''Population operators for the master equation 
    @ var N: size of the oscillator Hilbert space
    
    returns: list of population operators for the harmonic oscillator and the atom
    '''
    a = tensor(destroy(N), qeye(2))
    sm = tensor( qeye(N), destroy(2))
    return [a.dag()*a, sm.dag()*sm]

def Ren_gate(carbon_nr, B_field=304.22):
    '''create a Ren gate for given carbon number, only interacting with the electron spin '''

    U0, U1, U0id, U1id = nuclear_Ren_matrix(carbon_nr, B_field)

    Ren = qutip.tensor(rho0,U0)+qutip.tensor(rho1,U1)
    Ren_id = qutip.tensor(rho0,U0id)+qutip.tensor(rho1,U1id)

    return Ren, Ren_id