def get_lattice_Ls(cell, nimgs=None, rcut=None, dimension=None):
    '''Get the (Cartesian, unitful) lattice translation vectors for nearby images.
    The translation vectors can be used for the lattice summation.'''
    b = cell.reciprocal_vectors(norm_to=1)
    heights_inv = lib.norm(b, axis=1)

    if nimgs is None:
        if rcut is None:
            rcut = cell.rcut
# plus 1 image in rcut to handle the case atoms within the adjacent cells are
# close to each other
        rcut = rcut + min(1./heights_inv)
        nimgs = np.ceil(rcut*heights_inv)
    else:
        rcut = max((np.asarray(nimgs))/heights_inv) + min(1./heights_inv) # ~ the inradius

    if dimension is None:
        dimension = cell.dimension
    if dimension == 0:
        nimgs = [0, 0, 0]
    elif dimension == 1:
        nimgs = [nimgs[0], 0, 0]
    elif dimension == 2:
        nimgs = [nimgs[0], nimgs[1], 0]

    Ts = lib.cartesian_prod((np.arange(-nimgs[0],nimgs[0]+1),
                             np.arange(-nimgs[1],nimgs[1]+1),
                             np.arange(-nimgs[2],nimgs[2]+1)))
    Ls = np.dot(Ts, cell.lattice_vectors())
    Ls = Ls[lib.norm(Ls, axis=1)<rcut]
    return np.asarray(Ls, order='C')

 def norm_xy(w, z):
     zp = e_ia * z.reshape(e_ia.shape)
     zm = w/e_ia * z.reshape(e_ia.shape)
     x = (zp + zm) * .5
     y = (zp - zm) * .5
     norm = lib.norm(x)**2 - lib.norm(y)**2
     norm = numpy.sqrt(.5/norm)  # normalize to 0.5 for alpha spin
     return (x*norm, y*norm)

 def norm_xy(w, z):
     zp = eai * z.reshape(eai.shape)
     zm = w/eai * z.reshape(eai.shape)
     x = (zp + zm) * .5
     y = (zp - zm) * .5
     norm = 2*(lib.norm(x)**2 - lib.norm(y)**2)
     norm = 1/numpy.sqrt(norm)
     return x*norm,y*norm

def precompute_exx(cell, kpts):
    from pyscf.pbc import gto as pbcgto
    from pyscf.pbc.dft import gen_grid
    log = lib.logger.Logger(cell.stdout, cell.verbose)
    log.debug("# Precomputing Wigner-Seitz EXX kernel")
    Nk = get_monkhorst_pack_size(cell, kpts)
    log.debug("# Nk = %s", Nk)

    kcell = pbcgto.Cell()
    kcell.atom = 'H 0. 0. 0.'
    kcell.spin = 1
    kcell.unit = 'B'
    kcell.verbose = 0
    kcell.a = cell.lattice_vectors() * Nk
    Lc = 1.0/lib.norm(np.linalg.inv(kcell.a), axis=0)
    log.debug("# Lc = %s", Lc)
    Rin = Lc.min() / 2.0
    log.debug("# Rin = %s", Rin)
    # ASE:
    alpha = 5./Rin # sqrt(-ln eps) / Rc, eps ~ 10^{-11}
    log.info("WS alpha = %s", alpha)
    kcell.mesh = np.array([4*int(L*alpha*3.0) for L in Lc])  # ~ [60,60,60]
    # QE:
    #alpha = 3./Rin * np.sqrt(0.5)
    #kcell.mesh = (4*alpha*np.linalg.norm(kcell.a,axis=1)).astype(int)
    log.debug("# kcell.mesh FFT = %s", kcell.mesh)
    kcell.build(False,False)
    rs = gen_grid.gen_uniform_grids(kcell)
    kngs = len(rs)
    log.debug("# kcell kngs = %d", kngs)
    corners = np.dot(np.indices((2,2,2)).reshape((3,8)).T, kcell.a)
    #vR = np.empty(kngs)
    #for i, rv in enumerate(rs):
    #    # Minimum image convention to corners of kcell parallelepiped
    #    r = lib.norm(rv-corners, axis=1).min()
    #    if np.isclose(r, 0.):
    #        vR[i] = 2*alpha / np.sqrt(np.pi)
    #    else:
    #        vR[i] = scipy.special.erf(alpha*r) / r
    r = np.min([lib.norm(rs-c, axis=1) for c in corners], axis=0)
    vR = scipy.special.erf(alpha*r) / (r+1e-200)
    vR[r<1e-9] = 2*alpha / np.sqrt(np.pi)
    vG = (kcell.vol/kngs) * fft(vR, kcell.mesh)
    ws_exx = {'alpha': alpha,
              'kcell': kcell,
              'q'    : kcell.Gv,
              'vq'   : vG}
    log.debug("# Finished precomputing")
    return ws_exx

def get_ewald_params(cell, precision=1e-8, gs=None):
    r'''Choose a reasonable value of Ewald 'eta' and 'cut' parameters.

    Choice is based on largest G vector and desired relative precision.

    The relative error in the G-space sum is given by (keeping only
    exponential factors)
        precision ~ e^{(-Gmax^2)/(4 \eta^2)}
    which determines eta. Then, real-space cutoff is determined by (exp.
    factors only)
        precision ~ erfc(eta*rcut) / rcut ~ e^{(-eta**2 rcut*2)}

    Returns:
        ew_eta, ew_cut : float
            The Ewald 'eta' and 'cut' parameters.
    '''
    if gs is None:
        gs = cell.gs

    #  See Martin, p. 85 
    _h = cell.lattice_vectors()
    Gmax = min([ 2.*np.pi*gs[i]/lib.norm(_h[i,:]) for i in range(3) ])

    log_precision = np.log(precision)
    ew_eta = float(np.sqrt(-Gmax**2/(4*log_precision)))

    rcut = np.sqrt(-log_precision)/ew_eta
    ew_cut = cell.get_bounding_sphere(rcut)
    return ew_eta, ew_cut

def get_ewald_params(cell, precision=1e-8, gs=None):
    r'''Choose a reasonable value of Ewald 'eta' and 'cut' parameters.

    Choice is based on largest G vector and desired relative precision.

    The relative error in the G-space sum is given by (keeping only
    exponential factors)

        precision ~ e^{(-Gmax^2)/(4 \eta^2)}

    which determines eta. Then, real-space cutoff is determined by (exp.
    factors only)

        precision ~ erfc(eta*rcut) / rcut ~ e^{(-eta**2 rcut*2)}

    Returns:
        ew_eta, ew_cut : float
            The Ewald 'eta' and 'cut' parameters.
    '''
    if gs is None:
        gs = cell.gs

    if cell.dimension == 3:
        Gmax = min(np.asarray(cell.gs) * lib.norm(cell.reciprocal_vectors(), axis=1))
        log_precision = np.log(precision*.1)
        ew_eta = np.sqrt(-Gmax**2/(4*log_precision))
        ew_cut = np.sqrt(-log_precision)/ew_eta
    else:
# Non-uniform PW grids are used for low-dimensional ewald summation.  The cutoff
# estimation for long range part based on exp(G^2/(4*eta^2)) does not work for
# non-uniform grids.  Smooth model density is preferred.
        ew_cut = cell.rcut
        ew_eta = np.sqrt(max(np.log(ew_cut**2/precision)/ew_cut**2, .1))
    return ew_eta, ew_cut

def mesh_to_cutoff(a, mesh):
    '''
    Convert #grid points to KE cutoff
    '''
    b = 2 * np.pi * np.linalg.inv(a.T)
    Gmax = lib.norm(b, axis=1) * np.asarray(mesh) * .5
    return Gmax**2/2

def estimate_eta(cell, cutoff=1e-12):
    '''The exponent of the smooth gaussian model density, requiring that at
    boundary, density ~ 4pi rmax^2 exp(-eta*rmax^2) ~ 1e-12
    '''
    rmax = max(lib.norm(cell.lattice_vectors(), axis=0))
    eta = max(-numpy.log(cutoff/(4*numpy.pi*rmax**2))/rmax**2, .1)
    return eta

def cosmo_fock_o1(cosmo, dm):
    mol = cosmo.mol
    nao = dm.shape[0]
    # phi
    cosmo.loadsegs()
    coords = cosmo.cosurf[:cosmo.nps*3].reshape(-1,3)
    fakemol = _make_fakemol(coords)
    j3c = df.incore.aux_e2(mol, fakemol, intor='cint3c2e_sph', aosym='s2ij')
    tril_dm = lib.pack_tril(dm) * 2
    diagidx = numpy.arange(nao)
    diagidx = diagidx*(diagidx+1)//2 + diagidx
    tril_dm[diagidx] *= .5
    cosmo.phi = -numpy.einsum('x,xk->k', tril_dm, j3c)
    for ia in range(mol.natm):
        cosmo.phi += mol.atom_charge(ia)/lib.norm(mol.atom_coord(ia)-coords, axis=1)
    cosmo.savesegs()
    # qk
    cosmo.charges()
    # vpot
    cosmo.loadsegs()
#X    fakemol = _make_fakemol(cosmo.cosurf[:cosmo.nps*3].reshape(-1,3))
#X    j3c = df.incore.aux_e2(mol, fakemol, intor='cint3c2e_sph', aosym='s2ij')
    fock = lib.unpack_tril(numpy.einsum('xk,k->x', j3c, -cosmo.qcos[:cosmo.nps]))
    fepsi = cosmo.fepsi() 
    fock = fepsi*fock
    return fock

def make_mask(cell, coords, relativity=0, shls_slice=None, verbose=None):
    '''Mask to indicate whether a shell is zero on grid.
    The resultant mask array is an extension to the mask array used in
    molecular code (see also pyscf.dft.numint.make_mask function).
    For given shell ID and block ID, the value of the extended mask array
    means the number of images in Ls that does not vanish.
    '''
    coords = np.asarray(coords, order='F')
    natm = ctypes.c_int(cell._atm.shape[0])
    nbas = ctypes.c_int(cell.nbas)
    ngrids = len(coords)
    if shls_slice is None:
        shls_slice = (0, cell.nbas)
    assert(shls_slice == (0, cell.nbas))

    Ls = cell.get_lattice_Ls(dimension=3)
    Ls = Ls[np.argsort(lib.norm(Ls, axis=1))]

    non0tab = np.empty(((ngrids+BLKSIZE-1)//BLKSIZE, cell.nbas),
                          dtype=np.uint8)
    libpbc.PBCnr_ao_screen(non0tab.ctypes.data_as(ctypes.c_void_p),
                           coords.ctypes.data_as(ctypes.c_void_p),
                           ctypes.c_int(ngrids),
                           Ls.ctypes.data_as(ctypes.c_void_p),
                           ctypes.c_int(len(Ls)),
                           cell._atm.ctypes.data_as(ctypes.c_void_p), natm,
                           cell._bas.ctypes.data_as(ctypes.c_void_p), nbas,
                           cell._env.ctypes.data_as(ctypes.c_void_p))
    return non0tab

def make_psi(mol, dm, r_vdw, lmax):
    grids = dft.gen_grid.Grids(mol)
    atom_grids_tab = grids.gen_atomic_grids(mol)
    grids.build()

    ao = dft.numint.eval_ao(mol, grids.coords)
    den = dft.numint.eval_rho(mol, ao, dm)
    den *= grids.weights
    natm = mol.natm
    nlm = (lmax+1)**2
    psi = numpy.empty((natm,nlm))
    i1 = 0
    for ia in range(natm):
        xnj, w = atom_grids_tab[mol.atom_symbol(ia)]
        i0, i1 = i1, i1 + w.size
        r = lib.norm(xnj, axis=1)
        snj = xnj/r.reshape(-1,1)
        Ys = sph.real_sph_vec(snj, lmax, True)
        p1 = 0
        for l in range(lmax+1):
            fac = 4*numpy.pi/(l*2+1)
            p0, p1 = p1, p1 + (l*2+1)
            rr = numpy.zeros_like(r)
            rr[r<=r_vdw[ia]] = r[r<=r_vdw[ia]]**l / r_vdw[ia]**(l+1)
            rr[r> r_vdw[ia]] = r_vdw[ia]**l / r[r>r_vdw[ia]]**(l+1)
            psi[ia,p0:p1] = -fac * numpy.einsum('n,n,mn->m', den[i0:i1], rr, Ys[l])
        psi[ia,0] += numpy.sqrt(4*numpy.pi)/r_vdw[ia] * mol.atom_charge(ia)
    return psi

 def makedm(mos, occs):
     where = [np.argmin(lib.norm(chk_kpts-kpt, axis=1)) for kpt in kpts]
     moa, mob = mos
     occa, occb = occs
     mos = ([fproj(moa[w], chk_kpts[w]-kpts[i]) for i,w in enumerate(where)],
            [fproj(mob[w], chk_kpts[w]-kpts[i]) for i,w in enumerate(where)])
     occs = (occa[where],occb[where])
     return make_rdm1(mos, occs)

        def energy_nuc(self):
# nuclei lattice interaction
            nuc = self.mol.energy_nuc()
            for j in range(self.mol.natm):
                q2, r2 = self.mol.atom_charge(j), self.mol.atom_coord(j)
                r = lib.norm(r2-coords, axis=1)
                nuc += q2*(charges/r).sum()
            return nuc

        def energy_tot(self, dm=None, h1e=None, vhf=None):
# nuclei lattice interaction
            nuc = 0.0
            for j in range(self.mol.natm):
                q2, r2 = self.mol.atom_charge(j), self.mol.atom_coord(j)
                r = lib.norm(r2-coords, axis=1)
                nuc += q2*(charges/r).sum()
            return method_class.energy_tot(self, dm, h1e, vhf) + nuc

 def makedm(mos, occs):
     where = [np.argmin(lib.norm(chk_kpts-kpt, axis=1)) for kpt in kpts]
     moa, mob = mos
     occa, occb = occs
     dkpts = [chk_kpts[w]-kpts[i] for i,w in enumerate(where)]
     mos = (fproj([moa[w] for w in where], dkpts),
            fproj([mob[w] for w in where], dkpts))
     occs = ([occa[i] for i in where], [occb[i] for i in where])
     return make_rdm1(mos, occs)

def enlarge_space(myci, civec_strs, eri, norb, nelec):
    if isinstance(civec_strs, (tuple, list)):
        nelec, (strsa, strsb) = _unpack(civec_strs[0], nelec)[1:]
        ci_coeff = lib.asarray(civec_strs)
    else:
        ci_coeff, nelec, (strsa, strsb) = _unpack(civec_strs, nelec)
    na = len(strsa)
    nb = len(strsb)
    ci0 = ci_coeff.reshape(-1,na,nb)
    civec_a_max = lib.norm(ci0, axis=2).max(axis=0)
    civec_b_max = lib.norm(ci0, axis=1).max(axis=0)
    ci_aidx = numpy.where(civec_a_max > myci.ci_coeff_cutoff)[0]
    ci_bidx = numpy.where(civec_b_max > myci.ci_coeff_cutoff)[0]
    civec_a_max = civec_a_max[ci_aidx]
    civec_b_max = civec_b_max[ci_bidx]
    strsa = strsa[ci_aidx]
    strsb = strsb[ci_bidx]

    eri = ao2mo.restore(1, eri, norb)
    eri_pq_max = abs(eri.reshape(norb**2,-1)).max(axis=1).reshape(norb,norb)

    strsa_add = select_strs(myci, eri, eri_pq_max, civec_a_max, strsa, norb, nelec[0])
    strsb_add = select_strs(myci, eri, eri_pq_max, civec_b_max, strsb, norb, nelec[1])
    strsa = numpy.append(strsa, strsa_add)
    strsb = numpy.append(strsb, strsb_add)
    aidx = numpy.argsort(strsa)
    bidx = numpy.argsort(strsb)
    ci_strs = (strsa[aidx], strsb[bidx])
    aidx = numpy.where(aidx < len(ci_aidx))[0]
    bidx = numpy.where(bidx < len(ci_bidx))[0]
    ma = len(strsa)
    mb = len(strsb)

    cs = []
    for i in range(ci0.shape[0]):
        ci1 = numpy.zeros((ma,mb))
        tmp = lib.take_2d(ci0[i], ci_aidx, ci_bidx)
        lib.takebak_2d(ci1, tmp, aidx, bidx)
        cs.append(_as_SCIvector(ci1, ci_strs))

    if not isinstance(civec_strs, (tuple, list)) and civec_strs.ndim < 3:
        cs = cs[0]
    return cs

def get_vlocG(cell):
    '''Local PP kernel in G space: Vloc(G) for G!=0, 0 for G=0.

    Returns:
        (natm, ngs) ndarray
    '''
    Gvnorm = lib.norm(cell.Gv, axis=1)
    vlocG = get_gth_vlocG(cell, Gvnorm)
    vlocG[:,0] = 0.
    return vlocG

def get_bounding_sphere(cell, rcut):
    '''Finds all the lattice points within a sphere of radius rcut.  

    Defines a parallelipiped given by -N_x <= n_x <= N_x, with x in [1,3]
    See Martin p. 85

    Args:
        rcut : number
            real space cut-off for interaction

    Returns:
        cut : ndarray of 3 ints defining N_x
    '''
    Gmat = invh = scipy.linalg.inv(cell.lattice_vectors())
    n1 = np.ceil(lib.norm(Gmat[0,:])*rcut).astype(int)
    n2 = np.ceil(lib.norm(Gmat[1,:])*rcut).astype(int)
    n3 = np.ceil(lib.norm(Gmat[2,:])*rcut).astype(int)
    cut = np.array([n1, n2, n3])
    return cut

def cart2polar(rvec):
    # The rows of rvec are the 3-component vectors
    # i.e. rvec is N x 3
    x,y,z = rvec.T
    r = lib.norm(rvec, axis=1)
    # theta is the polar angle, 0 < theta < pi
    # catch possible 0/0
    theta = np.arccos(z/(r+1e-200))
    # phi is the azimuthal angle, 0 < phi < 2pi (or -pi < phi < pi)
    phi = np.arctan2(y,x)
    return r, theta, phi

 def grad_nuc(self, mol=None, atmlst=None):
     if mol is None: mol = method.mol
     g_qm = method.grad_nuc(mol, atmlst)
     g_mm = numpy.empty((mol.natm,3))
     for i in range(mol.natm):
         q1 = mol.atom_charge(i)
         r1 = mol.atom_coord(i)
         r = lib.norm(r1-coords, axis=1)
         g_mm[i] = -q1 * numpy.einsum('i,ix,i->x', charges, r1-coords, 1/r**3)
     if atmlst is not None:
         g_mm = g_mm[atmlst]
     return g_qm + g_mm