def get_init_guess(norb, nelec, nroots, hdiag, orbsym, wfnsym=0):
    if isinstance(nelec, (int, numpy.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(norb), nelecb))
    airreps = numpy.zeros(strsa.size, dtype=numpy.int32)
    birreps = numpy.zeros(strsb.size, dtype=numpy.int32)
    for i in range(norb):
        airreps[numpy.bitwise_and(strsa, 1<<i) > 0] ^= orbsym[i]
        birreps[numpy.bitwise_and(strsb, 1<<i) > 0] ^= orbsym[i]
    na = len(strsa)
    nb = len(strsb)

    ci0 = []
    iroot = 0
    for addr in numpy.argsort(hdiag):
        x = numpy.zeros((na*nb))
        addra = addr // nb
        addrb = addr % nb
        if airreps[addra] ^ birreps[addrb] == wfnsym:
            x[addr] = 1
            ci0.append(x)
            iroot += 1
            if iroot >= nroots:
                break
    return ci0

def contract_2e_hubbard(u, fcivec, norb, nelec, opt=None):
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    u_aa, u_ab, u_bb = u

    strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(norb), nelecb))
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    fcivec = fcivec.reshape(na,nb)
    fcinew = numpy.zeros_like(fcivec)

    if u_aa != 0:  # u * n_alpha^+ n_alpha
        for i in range(norb):
            maska = (strsa & (1<<i)) > 0
            fcinew[maska] += u_aa * fcivec[maska]
    if u_ab != 0:  # u * (n_alpha^+ n_beta + n_beta^+ n_alpha)
        for i in range(norb):
            maska = (strsa & (1<<i)) > 0
            maskb = (strsb & (1<<i)) > 0
            fcinew[maska[:,None]&maskb] += 2*u_ab * fcivec[maska[:,None]&maskb]
    if u_bb != 0:  # u * n_beta^+ n_beta
        for i in range(norb):
            maskb = (strsb & (1<<i)) > 0
            fcinew[:,maskb] += u_bb * fcivec[:,maskb]
    return fcinew

def contract_ep(g, fcivec, nsite, nelec, nphonon):
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(nsite), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(nsite), nelecb))
    cishape = make_shape(nsite, nelec, nphonon)
    na, nb = cishape[:2]
    ci0 = fcivec.reshape(cishape)
    fcinew = numpy.zeros(cishape)
    nbar = float(neleca+nelecb) / nsite

    phonon_cre = numpy.sqrt(numpy.arange(1,nphonon+1))
    for i in range(nsite):
        maska = (strsa & (1<<i)) > 0
        maskb = (strsb & (1<<i)) > 0
        e_part = numpy.zeros((na,nb))
        e_part[maska,:] += 1
        e_part[:,maskb] += 1
        e_part[:] -= float(neleca+nelecb) / nsite
        for ip in range(nphonon):
            slices1 = slices_for_cre(i, nsite, ip)
            slices0 = slices_for    (i, nsite, ip)
            fcinew[slices1] += numpy.einsum('ij...,ij...->ij...', g*phonon_cre[ip]*e_part, ci0[slices0])
            fcinew[slices0] += numpy.einsum('ij...,ij...->ij...', g*phonon_cre[ip]*e_part, ci0[slices1])
    return fcinew.reshape(fcivec.shape)

def get_init_guess(norb, nelec, nroots, hdiag, orbsym, wfnsym=0):
    neleca, nelecb = _unpack_nelec(nelec)
    strsa = cistring.gen_strings4orblist(range(norb), neleca)
    airreps = birreps = _gen_strs_irrep(strsa, orbsym)
    if neleca != nelecb:
        strsb = cistring.gen_strings4orblist(range(norb), nelecb)
        birreps = _gen_strs_irrep(strsb, orbsym)
    return _get_init_guess(airreps, birreps, nroots, hdiag, orbsym, wfnsym)

def transform_ci_for_orbital_rotation(ci, norb, nelec, u):
    '''Transform CI coefficients to the representation in new one-particle basis.
    Solving CI problem for Hamiltonian h1, h2 defined in old basis,
    CI_old = fci.kernel(h1, h2, ...)
    Given orbital rotation u, the CI problem can be either solved by
    transforming the Hamiltonian, or transforming the coefficients.
    CI_new = fci.kernel(u^T*h1*u, ...) = transform_ci_for_orbital_rotation(CI_old, u)

    Args:
        u : 2D array or a list of 2D array
            the orbital rotation to transform the old one-particle basis to new
            one-particle basis
    '''
    neleca, nelecb = _unpack(nelec)
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(norb), nelecb))
    one_particle_strs = numpy.asarray([1<<i for i in range(norb)])
    na = len(strsa)
    nb = len(strsb)

    if isinstance(u, numpy.ndarray) and u.ndim == 2:
        ua = ub = u
    else:
        ua, ub = u

    # Unitary transformation array trans_ci is the overlap between two sets of CI basis.
    occ_masks = (strsa[:,None] & one_particle_strs) != 0
    trans_ci_a = numpy.zeros((na,na))
    #for i in range(na): # for old basis
    #    for j in range(na):
    #        uij = u[occ_masks[i]][:,occ_masks[j]]
    #        trans_ci_a[i,j] = numpy.linalg.det(uij)
    occ_idx_all_strs = numpy.where(occ_masks)[1]
    for i in range(na):
        ui = ua[occ_masks[i]].T.copy()
        minors = numpy.take(ui, occ_idx_all_strs, axis=0).reshape(na,neleca,neleca)
        trans_ci_a[i,:] = numpy.linalg.det(minors)

    if neleca == nelecb and numpy.allclose(ua, ub):
        trans_ci_b = trans_ci_a
    else:
        occ_masks = (strsb[:,None] & one_particle_strs) != 0
        trans_ci_b = numpy.zeros((nb,nb))
        #for i in range(nb):
        #    for j in range(nb):
        #        uij = u[occ_masks[i]][:,occ_masks[j]]
        #        trans_ci_b[i,j] = numpy.linalg.det(uij)
        occ_idx_all_strs = numpy.where(occ_masks)[1]
        for i in range(nb):
            ui = ub[occ_masks[i]].T.copy()
            minors = numpy.take(ui, occ_idx_all_strs, axis=0).reshape(nb,nelecb,nelecb)
            trans_ci_b[i,:] = numpy.linalg.det(minors)

    # Transform old basis to new basis for all alpha-electron excitations
    ci = lib.dot(trans_ci_a.T, ci.reshape(na,nb))
    # Transform old basis to new basis for all beta-electron excitations
    ci = lib.dot(ci.reshape(na,nb), trans_ci_b)
    return ci

 def test_strings4orblist(self):
     ref = ['0b1010', '0b100010', '0b101000', '0b10000010', '0b10001000',
            '0b10100000']
     self.assertEqual(cistring.gen_strings4orblist([1,3,5,7], 2),
                      [int(x,2) for x in ref])
     ref = ['0b11', '0b101', '0b110', '0b1001', '0b1010', '0b1100',
            '0b10001', '0b10010', '0b10100', '0b11000']
     self.assertEqual(cistring.gen_strings4orblist(range(5), 2),
                      [int(x,2) for x in ref])

def contract_2e_hubbard(u, fcivec, nsite, nelec, nphonon):
    neleca, nelecb = _unpack_nelec(nelec)
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(nsite), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(nsite), nelecb))
    cishape = make_shape(nsite, nelec, nphonon)
    ci0 = fcivec.reshape(cishape)
    fcinew = numpy.zeros(cishape)

    for i in range(nsite):
        maska = (strsa & (1<<i)) > 0
        maskb = (strsb & (1<<i)) > 0
        fcinew[maska[:,None]&maskb] += u * ci0[maska[:,None]&maskb]
    return fcinew.reshape(fcivec.shape)

    def test_strings4orblist(self):
        ref = ['0b1010', '0b100010', '0b101000', '0b10000010', '0b10001000',
               '0b10100000']
        for i, x in enumerate(cistring.gen_strings4orblist([1,3,5,7], 2)):
            self.assertEqual(bin(x), ref[i])
        ref = ['0b11', '0b101', '0b110', '0b1001', '0b1010', '0b1100',
               '0b10001', '0b10010', '0b10100', '0b11000']
        for i, x in enumerate(cistring.gen_strings4orblist(range(5), 2)):
            self.assertEqual(bin(x), ref[i])

        strs = cistring.gen_strings4orblist(range(8), 4)
        occlst = cistring._gen_occslst(range(8), 4)
        self.assertAlmostEqual(abs(occlst - cistring._strs2occslst(strs, 8)).sum(), 0, 12)
        self.assertAlmostEqual(abs(strs - cistring._occslst2strs(occlst)).sum(), 0, 12)

def contract_2e_hubbard(u, fcivec, norb, nelec, opt=None):
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    u_aa, u_ab, u_bb = u

    strsa = cistring.gen_strings4orblist(range(norb), neleca)
    strsb = cistring.gen_strings4orblist(range(norb), nelecb)
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    fcivec = fcivec.reshape(na,nb)
    t1a = numpy.zeros((norb,na,nb))
    t1b = numpy.zeros((norb,na,nb))
    fcinew = numpy.zeros_like(fcivec)

    for addr, s in enumerate(strsa):
        for i in range(norb):
            if s & (1<<i):
                t1a[i,addr] += fcivec[addr]
    for addr, s in enumerate(strsb):
        for i in range(norb):
            if s & (1<<i):
                t1b[i,:,addr] += fcivec[:,addr]

    if u_aa != 0:
        # u * n_alpha^+ n_alpha
        for addr, s in enumerate(strsa):
            for i in range(norb):
                if s & (1<<i):
                    fcinew[addr] += t1a[i,addr] * u_aa
    if u_ab != 0:
        # u * n_alpha^+ n_beta
        for addr, s in enumerate(strsa):
            for i in range(norb):
                if s & (1<<i):
                    fcinew[addr] += t1b[i,addr] * u_ab
        # u * n_beta^+ n_alpha
        for addr, s in enumerate(strsb):
            for i in range(norb):
                if s & (1<<i):
                    fcinew[:,addr] += t1a[i,:,addr] * u_ab
    if u_bb != 0:
        # u * n_beta^+ n_beta
        for addr, s in enumerate(strsb):
            for i in range(norb):
                if s & (1<<i):
                    fcinew[:,addr] += t1b[i,:,addr] * u_bb
    return fcinew

    def _parse_fci_vector(self, ci_vecmat):
        """ Translate the PySCF FCI matrix into a dictionary of configurations and weights

        Args:
            ci_vecmat (np.ndarray): ci vector from a PySCF FCI calculation

        Returns:
            Mapping[str, float]: dictionary of configuration weights (normalized) organized by
                configuration label. Configurations labeled by their active space orbital
                occupations: 0 (unoccupied), a (alpha electron only), b (beta electron only), or '2'
                (doubly occupied)

        Example:
            >>> import numpy as np
            >>> model = PySCFPotential(active_orbitals=2, active_electrons=2)
            >>> model._parse_fci_vector(np.array([[1.0, 2.0],[3.0, 4.0]]))
            {'20': 1.0,
             'ba': 2.0,
             'ab': 3.0,
             '02': 4.0}
        """
        from pyscf.fci import cistring
        conf_bin = cistring.gen_strings4orblist(list(range(self.params.active_orbitals)),
                                                old_div(self.params.active_electrons,2))
        civecs = {}
        for i, ca in enumerate(conf_bin):
            for j, cb in enumerate(conf_bin):
                astring = bin(ca)[2:].zfill(self.params.active_orbitals)
                bstring = bin(cb)[2:].zfill(self.params.active_orbitals)
                s = ''.join(reversed([self._OCCMAP[a, b] for a, b in zip(astring, bstring)]))
                civecs[s] = ci_vecmat[i, j]
        return civecs

 def test_contract1(self):
     myci = select_ci.SelectCI()
     nelec = (4,3)
     strsa = cistring.gen_strings4orblist(range(norb), nelec[0])
     strsb = cistring.gen_strings4orblist(range(norb), nelec[1])
     ci0 = select_ci._as_SCIvector(numpy.random.random((len(strsa),len(strsb))), (strsa,strsb))
     h2 = ao2mo.restore(1, eri, norb)
     c1 = myci.contract_2e(h2, ci0, norb, nelec)
     c2 = direct_spin1.contract_2e(h2, ci0, norb, nelec)
     self.assertAlmostEqual(abs(c1-c2).sum(), 0, 9)
     dm1_1 = myci.make_rdm1(c1, norb, nelec)
     dm1_2 = direct_spin1.make_rdm1(c2, norb, nelec)
     self.assertAlmostEqual(abs(dm1_1 - dm1_2).sum(), 0, 9)
     dm2_1 = myci.make_rdm2(c1, norb, nelec)
     dm2_2 = direct_spin1.make_rdm12(c2, norb, nelec)[1]
     self.assertAlmostEqual(abs(dm2_1 - dm2_2).sum(), 0, 9)

def to_fci(civec, norb, nelec, root=0):
    assert(norb <= 64)
    neleca, nelecb = nelec
    strsa = cistring.gen_strings4orblist(range(norb), neleca)
    stradic = dict(zip(strsa,range(strsa.__len__())))
    strsb = cistring.gen_strings4orblist(range(norb), nelecb)
    strbdic = dict(zip(strsb,range(strsb.__len__())))
    na = len(stradic)
    nb = len(strbdic)
    ndet = len(civec[root])
    fcivec = numpy.zeros((na,nb))
    for idet, (stra, strb) in enumerate(civec[root]._strs.reshape(ndet,2,-1)):
        ka = stradic[stra[0]]
        kb = strbdic[strb[0]]
        fcivec[ka,kb] = civec[root][idet]
    return fcivec

    def test_rdm(self):
        norb, nelec = 10, 4
        strs = cistring.gen_strings4orblist(range(norb), nelec)
        numpy.random.seed(11)
        mask = numpy.random.random(len(strs)) > .6
        strsa = strs[mask]
        mask = numpy.random.random(len(strs)) > .7
        strsb = strs[mask]
        ci_strs = (strsa, strsb)
        ci_coeff = select_ci._as_SCIvector(numpy.random.random((len(strsa),len(strsb))), ci_strs)
        ci0 = select_ci.to_fci(ci_coeff, norb, (nelec,nelec))
        dm1ref, dm2ref = direct_spin1.make_rdm12s(ci0, norb, (nelec,nelec))
        dm1 = select_ci.make_rdm1s(ci_coeff, norb, (nelec,nelec))
        self.assertAlmostEqual(abs(dm1[0]-dm1ref[0]).sum(), 0, 9)
        self.assertAlmostEqual(abs(dm1[1]-dm1ref[1]).sum(), 0, 9)
        dm2 = select_ci.make_rdm2s(ci_coeff, norb, (nelec,nelec))
        self.assertAlmostEqual(abs(dm2[0]-dm2ref[0]).sum(), 0, 9)
        self.assertAlmostEqual(abs(dm2[1]-dm2ref[1]).sum(), 0, 9)
        self.assertAlmostEqual(abs(dm2[2]-dm2ref[2]).sum(), 0, 9)

        ci1_coeff = select_ci._as_SCIvector(numpy.random.random((len(strsa),len(strsb))), ci_strs)
        ci1 = select_ci.to_fci(ci1_coeff, norb, (nelec,nelec))
        dm1ref, dm2ref = direct_spin1.trans_rdm12s(ci1, ci0, norb, (nelec,nelec))
        dm1 = select_ci.trans_rdm1s(ci1_coeff, ci_coeff, norb, (nelec,nelec))
        self.assertAlmostEqual(abs(dm1[0]-dm1ref[0]).sum(), 0, 9)
        self.assertAlmostEqual(abs(dm1[1]-dm1ref[1]).sum(), 0, 9)

 def test_parity(self):
     strs = cistring.gen_strings4orblist(range(5), 3)
     links = cistring.gen_linkstr_index(range(5), 3)
     parity = []
     for addr0, link in enumerate(links):
         parity.append([cistring.parity(strs[addr0], strs[addr1])
                        for addr1 in link[:,2]])
     self.assertEqual(parity, links[:,:,3].tolist())

def contract_2e_hubbard(u, fcivec, nsite, nelec, nphonon):
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(nsite), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(nsite), nelecb))
    cishape = make_shape(nsite, nelec, nphonon)
    ci0 = fcivec.reshape(cishape)
    fcinew = numpy.zeros(cishape)

    for i in range(nsite):
        maska = (strsa & (1<<i)) > 0
        maskb = (strsb & (1<<i)) > 0
        fcinew[maska[:,None]&maskb] += u * ci0[maska[:,None]&maskb]
    return fcinew.reshape(fcivec.shape)

def from_fci(fcivec, ci_strs, norb, nelec):
    neleca, nelecb = nelec
    strsa = cistring.gen_strings4orblist(range(norb), neleca)
    stradic = dict(zip(strsa,range(strsa.__len__())))
    strsb = cistring.gen_strings4orblist(range(norb), nelecb)
    strbdic = dict(zip(strsb,range(strsb.__len__())))
    na = len(stradic)
    nb = len(strbdic)
    fcivec = fcivec.reshape(na,nb)
    ta = [excitation_level(s, neleca) for s in strsa.reshape(-1,1)]
    tb = [excitation_level(s, nelecb) for s in strsb.reshape(-1,1)]
    ndet = len(ci_strs)
    civec = numpy.zeros(ndet)
    for idet, (stra, strb) in enumerate(ci_strs.reshape(ndet,2,-1)):
        ka = stradic[stra[0]]
        kb = strbdic[strb[0]]
        civec[idet] = fcivec[ka,kb]
    return as_SCIvector(civec, ci_strs)

def get_init_guess(norb, nelec, nroots, hdiag, orbsym, wfnsym=0):
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(norb), nelecb))
    airreps = numpy.zeros(strsa.size, dtype=numpy.int32)
    birreps = numpy.zeros(strsb.size, dtype=numpy.int32)
    for i in range(norb):
        airreps[numpy.bitwise_and(strsa, 1<<i) > 0] ^= orbsym[i]
        birreps[numpy.bitwise_and(strsb, 1<<i) > 0] ^= orbsym[i]
    na = len(strsa)
    nb = len(strsb)

    init_strs = []
    iroot = 0
    for addr in numpy.argsort(hdiag):
        addra = addr // nb
        addrb = addr % nb
        if airreps[addra] ^ birreps[addrb] == wfnsym:
            if (addrb,addra) not in init_strs:
                init_strs.append((addra,addrb))
                iroot += 1
                if iroot >= nroots:
                    break
    ci0 = []
    for addra,addrb in init_strs:
        x = numpy.zeros((na,nb))
        if addra == addrb == 0:
            x[addra,addrb] = 1
        else:
            x[addra,addrb] = x[addrb,addra] = numpy.sqrt(.5)
        ci0.append(x.ravel())

    # Add noise
    #ci0[0][0 ] += 1e-5
    #ci0[0][-1] -= 1e-5
    if len(ci0) == 0:
        raise RuntimeError('No determinant matches the target symmetry %s' %
                           wfnsym)
    return ci0

 def test_des_linkstr(self):
     norb, nelec = 10, 4
     strs = cistring.gen_strings4orblist(range(norb), nelec)
     numpy.random.seed(11)
     mask = numpy.random.random(len(strs)) > .6
     strs = strs[mask]
     c_index0 = gen_cre_linkstr(strs, norb, nelec)
     c_index1 = select_ci.gen_cre_linkstr(strs, norb, nelec)
     c_index1[:,:,1] = 0
     self.assertTrue(numpy.all(c_index0 == c_index1))

def reorder(ci, nelec, orbidxa, orbidxb=None):
    '''reorder the CI coefficients wrt the reordering of orbitals (The relation
    of the reordered orbitals and original orbitals is  new = old[idx]).  Eg.
    orbidx = [2,0,1] to map   old orbital  a b c  ->   new orbital  c a b
    old-strings   0b011, 0b101, 0b110 ==  (1,2), (1,3), (2,3)
    orb-strings   (3,1), (3,2), (1,2) ==  0B101, 0B110, 0B011    <= by gen_strings4orblist
    then argsort to translate the string representation to the address
    [2(=0B011), 0(=0B101), 1(=0B110)]
    '''
    if isinstance(nelec, (int, numpy.integer)):
        neleca = nelecb = nelec // 2
    else:
        neleca, nelecb = nelec
    if orbidxb is None:
        orbidxb = orbidxa
    guide_stringsa = cistring.gen_strings4orblist(orbidxa, neleca)
    guide_stringsb = cistring.gen_strings4orblist(orbidxb, nelecb)
    old_det_idxa = numpy.argsort(guide_stringsa)
    old_det_idxb = numpy.argsort(guide_stringsb)
    return ci.take(old_det_idxa, axis=0).take(old_det_idxb, axis=1)

 def test_des_des_linkstr(self):
     norb, nelec = 10, 4
     strs = cistring.gen_strings4orblist(range(norb), nelec)
     numpy.random.seed(11)
     mask = numpy.random.random(len(strs)) > .4
     strs = strs[mask]
     dd_index0 = des_des_linkstr(strs, norb, nelec)
     dd_index1 = select_ci.des_des_linkstr(strs, norb, nelec)
     self.assertTrue(numpy.all(dd_index0 == dd_index1))
     dd_index0 = des_des_linkstr_tril(strs, norb, nelec)
     dd_index1 = select_ci.des_des_linkstr_tril(strs, norb, nelec)
     dd_index1[:,:,1] = 0
     self.assertTrue(numpy.all(dd_index0 == dd_index1))