#
# Author: Qiming Sun <[email protected]>
#

'''
XC functional, the interface to xcfun (https://github.com/dftlibs/xcfun)
U. Ekstrom et al, J. Chem. Theory Comput., 6, 1971
'''

import copy
import ctypes
import math
import numpy
from pyscf import lib

_itrf = lib.load_library('libxcfun_itrf')

XC = XC_CODES = {
'SLATERX'       :  0,  # Slater LDA exchange
'PW86X'         :  1,  # PW86 exchange
'VWN3C'         :  2,  # VWN3 LDA Correlation functional
'VWN5C'         :  3,  # VWN5 LDA Correlation functional
'PBEC'          :  4,  # PBE correlation functional
'PBEX'          :  5,  # PBE Exchange Functional
'BECKEX'        :  6,  # Becke 88 exchange
'BECKECORRX'    :  7,  # Becke 88 exchange correction
'BECKESRX'      :  8,  # Short range Becke 88 exchange
'BECKECAMX'     :  9,  # CAM Becke 88 exchange
'BRX'           : 10,  # Becke-Roussells exchange with jp dependence
'BRC'           : 11,  # Becke-Roussells correlation with jp dependence
'BRXC'          : 12,  # Becke-Roussells correlation with jp dependence

#!/usr/bin/env python

import os
import ctypes
import unittest
import numpy
from pyscf import lib
from pyscf import scf
from pyscf import gto
from pyscf import ao2mo

libcvhf2 = lib.load_library("libcvhf")

mol = gto.Mole()
mol.verbose = 0
mol.output = None  #'out_h2o'
mol.atom = [["O", (0.0, 0.0, 0.0)], [1, (0.0, -0.757, 0.587)], [1, (0.0, 0.757, 0.587)]]

mol.basis = "cc-pvdz"

mol.build()
rhf = scf.RHF(mol)
rhf.scf()


nao = mol.nao_nr()
npair = nao * (nao + 1) // 2
c_atm = numpy.array(mol._atm, dtype=numpy.int32)
c_bas = numpy.array(mol._bas, dtype=numpy.int32)
c_env = numpy.array(mol._env)
natm = ctypes.c_int(c_atm.shape[0])

#!/usr/bin/env python

import os
import ctypes
import _ctypes
import unittest
import numpy
from pyscf import lib
from pyscf import scf
from pyscf import gto
from pyscf import ao2mo

libcvhf2 = lib.load_library('libcvhf')

mol = gto.Mole()
mol.verbose = 0
mol.output = None#'out_h2o'
mol.atom = [
    ['O' , (0. , 0.     , 0.)],
    [1   , (0. , -0.757 , 0.587)],
    [1   , (0. , 0.757  , 0.587)] ]

mol.basis = 'cc-pvdz'

mol.build()
rhf = scf.RHF(mol)
rhf.scf()


nao = mol.nao_nr()
npair = nao*(nao+1)//2

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import ctypes
import unittest
import numpy
from pyscf import lib
from pyscf import scf
from pyscf import gto
from pyscf import ao2mo

# FIXME

libri1 = lib.load_library('libri')

mol = gto.Mole()
mol.verbose = 0
mol.output = None#'out_h2o'
mol.atom = [
    ['O' , (0. , 0.     , 0.)],
    [1   , (0. , -0.757 , 0.587)],
    [1   , (0. , 0.757  , 0.587)] ]

mol.basis = {'H': 'cc-pvdz',
             'O': 'cc-pvdz',}
mol.build()
rhf = scf.RHF(mol)
rhf.scf()

# Author: Qiming Sun <[email protected]>
#

import time
import ctypes
import tempfile
import numpy
import h5py
import pyscf.lib as lib
from pyscf.lib import logger
import pyscf.ao2mo
import pyscf.cc.ccsd_slow as ccsd
from pyscf.cc import ccsd_rdm
from pyscf import grad

libcc = lib.load_library('libcc')

def IX_intermediates(cc, t1, t2, l1, l2, eris=None, d1=None, d2=None):
    if eris is None:
# Note eris are in Chemist's notation
        eris = ccsd._ERIS(cc)
    if d1 is None:
        doo, dvv = ccsd_rdm.gamma1_intermediates(cc, t1, t2, l1, l2)
    else:
        doo, dvv = d1
    if d2 is None:
# Note gamma2 are in Chemist's notation
        d2 = ccsd_rdm.gamma2_intermediates(cc, t1, t2, l1, l2)
    dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = d2
    dvvov = dovvv.transpose(2,3,0,1)
    nocc, nvir = t1.shape

import struct
import time
import tempfile
from subprocess import check_call, check_output, STDOUT, CalledProcessError
import numpy
from pyscf import lib
from pyscf import tools
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf import mcscf
from pyscf.dmrgscf import dmrg_sym
from pyscf import __config__

# Libraries
import pyscf.lib
libunpack = lib.load_library('libicmpspt')

# Settings
try:
    from pyscf.dmrgscf import settings
except ImportError:
    settings = lambda: None
    settings.BLOCKEXE = getattr(__config__, 'dmrgscf_BLOCKEXE', None)
    settings.BLOCKEXE_COMPRESS_NEVPT = \
            getattr(__config__, 'dmrgscf_BLOCKEXE_COMPRESS_NEVPT', None)
    settings.BLOCKSCRATCHDIR = getattr(__config__, 'dmrgscf_BLOCKSCRATCHDIR', None)
    settings.BLOCKRUNTIMEDIR = getattr(__config__, 'dmrgscf_BLOCKRUNTIMEDIR', None)
    settings.MPIPREFIX = getattr(__config__, 'dmrgscf_MPIPREFIX', None)
    settings.BLOCKVERSION = getattr(__config__, 'dmrgscf_BLOCKVERSION', None)
    if (settings.BLOCKEXE is None or settings.BLOCKSCRATCHDIR is None):
        import sys

#!/usr/bin/env python

import os
import ctypes
import numpy
from pyscf import lib

libcc = lib.load_library("libcc")


# NOTE requisite on data continuous


def unpack_tril(tril, out=None):
    assert tril.flags.c_contiguous
    count = tril.shape[0]
    nd = int(numpy.sqrt(tril.shape[1] * 2))
    if out is None:
        out = numpy.empty((count, nd, nd))
    libcc.CCunpack_tril(
        tril.ctypes.data_as(ctypes.c_void_p), out.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(count), ctypes.c_int(nd)
    )
    return out


def pack_tril(mat, out=None):
    assert mat.flags.c_contiguous
    count, nd = mat.shape[:2]
    if out is None:
        out = numpy.empty((count, nd * (nd + 1) // 2))
    libcc.CCpack_tril(

Note the 1-particle density matrix has the same convention as the mean-field
1-particle density matrix (see McWeeney's book Eq 5.4.20), which is
        dm[p,q] = < q^+ p >
The contraction between 1-particle Hamiltonian and 1-pdm is
        E = einsum('pq,qp', h1, 1pdm)
Different conventions are used in the high order density matrices:
        dm[p,q,r,s,...] = < p^+ r^+ ... s q >
'''

import ctypes
import numpy
from pyscf import lib
from pyscf.fci import cistring
from pyscf.fci.addons import _unpack_nelec

librdm = lib.load_library('libfci')

def reorder_rdm(rdm1, rdm2, inplace=False):
    nmo = rdm1.shape[0]
    if not inplace:
        rdm2 = rdm2.copy()
    for k in range(nmo):
        rdm2[:,k,k,:] -= rdm1.T
    #return rdm1, rdm2
    rdm2 = lib.transpose_sum(rdm2.reshape(nmo*nmo,-1), inplace=True) * .5
    return rdm1, rdm2.reshape(nmo,nmo,nmo,nmo)

# dm[p,q] = <|q^+ p|>
def make_rdm1_ms0(fname, cibra, ciket, norb, nelec, link_index=None):
    assert(cibra is not None and ciket is not None)
    cibra = numpy.asarray(cibra, order='C')

import scipy.linalg
import h5py
from pyscf import lib
from pyscf import gto
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.ao2mo import _ao2mo
from pyscf.scf import _vhf
from pyscf.df import incore
from pyscf.df import _ri

#
# for auxe1 (P|ij)
#

libri = lib.load_library("libri")


def cholesky_eri(
    mol,
    erifile,
    auxbasis="weigend+etb",
    dataname="eri_mo",
    tmpdir=None,
    int3c="cint3c2e_sph",
    aosym="s2ij",
    int2c="cint2c2e_sph",
    comp=1,
    ioblk_size=256,
    auxmol=None,
    verbose=0,

Reference for Lebedev-Laikov grid:
  V. I. Lebedev, and D. N. Laikov "A quadrature formula for the sphere of the
  131st algebraic order of accuracy", Doklady Mathematics, 59, 477-481 (1999)
'''


import ctypes
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import gto
from pyscf.dft import radi
from pyscf import __config__

libdft = lib.load_library('libdft')
BLKSIZE = 128  # needs to be the same to lib/gto/grid_ao_drv.c

# ~= (L+1)**2/3
LEBEDEV_ORDER = {
      0:    1,
      3:    6,
      5:   14,
      7:   26,
      9:   38,
     11:   50,
     13:   74,
     15:   86,
     17:  110,
     19:  146,
     21:  170,

(JCTC 2016, 12, 3674-3680)

Simple usage::

'''

import numpy
import time
import ctypes
from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.fci import cistring
from pyscf.fci import direct_spin1

libhci = lib.load_library('libhci')

def contract_2e_ctypes(h1_h2, civec, norb, nelec, hdiag=None, **kwargs):
    h1, eri = h1_h2
    strs = civec._strs
    ndet = len(strs)
    if hdiag is None:
        hdiag = make_hdiag(h1, eri, strs, norb, nelec)
    ci1 = numpy.zeros_like(civec)

    h1 = numpy.asarray(h1, order='C')
    eri = numpy.asarray(eri, order='C')
    strs = numpy.asarray(strs, order='C')
    civec = numpy.asarray(civec, order='C')
    hdiag = numpy.asarray(hdiag, order='C')
    ci1 = numpy.asarray(ci1, order='C')

# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import ctypes
import unittest
import numpy
from pyscf import lib, gto

libcgto = lib.load_library('libdft')
BLKSIZE = 64

mol = gto.M(atom='''
            O 0.5 0.5 0.
            H  1.  1.2 0.
            H  0.  0.  1.3
            ''',
            basis='ccpvqz')

def eval_gto(mol, eval_name, coords,
             comp=1, shls_slice=None, non0tab=None, ao_loc=None, out=None):
    atm = numpy.asarray(mol._atm, dtype=numpy.int32, order='C')
    bas = numpy.asarray(mol._bas, dtype=numpy.int32, order='C')
    env = numpy.asarray(mol._env, dtype=numpy.double, order='C')
    coords = numpy.asarray(coords, dtype=numpy.double, order='F')

import os
import ctypes
import time
import tempfile
from functools import reduce
import numpy
import h5py
from pyscf import lib
from pyscf.lib import logger
from pyscf import fci
from pyscf.mcscf import mc_ao2mo
from pyscf import ao2mo
from pyscf.ao2mo import _ao2mo

libmc = lib.load_library('libmcscf')

NUMERICAL_ZERO = 1e-14
# Ref JCP, 117, 9138

# h1e is the CAS space effective 1e hamiltonian
# h2e is the CAS space 2e integrals in  notation # a' -> p # b' -> q # c' -> r
# d' -> s

def make_a16(h1e, h2e, dms, civec, norb, nelec, link_index=None):
    dm3 = dms['3']
    #dm4 = dms['4']
    if 'f3ca' in dms and 'f3ac' in dms:
        f3ca = dms['f3ca']
        f3ac = dms['f3ac']
    else:

#!/usr/bin/env python
#
# Author: Qiming Sun <[email protected]>
#

import ctypes
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.fci import direct_spin1
from pyscf.fci import direct_spin1_symm
from pyscf.fci import select_ci
from pyscf.fci import addons

libfci = lib.load_library('libfci')

def reorder4irrep_minors(eri, norb, link_index, orbsym):
    if orbsym is None:
        return eri, link_index, numpy.array(norb, dtype=numpy.int32)
    orbsym = numpy.asarray(orbsym) % 10
    trilirrep = (orbsym[:,None]^orbsym)[numpy.tril_indices(norb,-1)]
    dimirrep = numpy.array(numpy.bincount(trilirrep), dtype=numpy.int32)
    order = numpy.argsort(trilirrep)
    rank = order.argsort()
    eri = lib.take_2d(eri, order, order)
    link_index_irrep = link_index.copy()
    link_index_irrep[:,:,0] = rank[link_index[:,:,0]]
    return numpy.asarray(eri, order='C'), link_index_irrep, dimirrep

def contract_2e(eri, civec_strs, norb, nelec, link_index=None, orbsym=None):

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <[email protected]>
#

import warnings
import ctypes
import numpy
from pyscf import lib
from pyscf.gto.moleintor import make_loc

BLKSIZE = 128 # needs to be the same to lib/gto/grid_ao_drv.c

libcgto = lib.load_library('libcgto')

def eval_gto(mol, eval_name, coords,
             comp=None, shls_slice=None, non0tab=None, ao_loc=None, out=None):
    r'''Evaluate AO function value on the given grids,

    Args:
        eval_name : str

            ==================  ======  =======================
            Function            comp    Expression
            ==================  ======  =======================
            "GTOval_sph"        1       |AO>
            "GTOval_ip_sph"     3       nabla |AO>
            "GTOval_ig_sph"     3       (#C(0 1) g) |AO>
            "GTOval_ipig_sph"   3       (#C(0 1) nabla g) |AO>

import ctypes
import numpy as np
import scipy.linalg
import scipy.optimize
import pyscf.lib.parameters as param
from pyscf import lib
from pyscf.lib import logger
from pyscf.gto import mole
from pyscf.gto import moleintor
from pyscf.gto.mole import _symbol, _rm_digit, _std_symbol, _charge
from pyscf.gto.mole import conc_env
from pyscf.pbc.gto import basis
from pyscf.pbc.gto import pseudo
from pyscf.pbc.tools import pbc as pbctools

libpbc = lib.load_library('libpbc')

def M(**kwargs):
    r'''This is a shortcut to build up Cell object.

    Examples:

    >>> from pyscf.pbc import gto
    >>> cell = gto.M(h=numpy.eye(3)*4, atom='He 1 1 1', basis='6-31g', gs=[10]*3)
    '''
    cell = Cell()
    cell.build(**kwargs)
    return cell


def format_pseudo(pseudo_tab):

from pyscf import gto, scf
from pyscf.tools import molden
from pyscf.lib import parameters as param
from pyscf.lib import logger
from pyscf.lib import linalg_helper
from pyscf.scf import _vhf
from pyscf import ao2mo

import numpy as np
import scipy
#import time

import ctypes
import _ctypes
from pyscf.lib import load_library
liblocalizer = load_library('liblocalizer')

class localizer:

    def __init__( self, mol, orbital_coeff, thetype, use_full_hessian=True ):
        r'''Initializer for the localization procedure

        Args:
            mol : A molecule which has been built
            orbital_coeff: Set of orthonormal orbitals, expressed in terms of the AO, which should be localized
            thetype: Which cost function to optimize: 'boys' or 'edmiston'
            use_full_hessian: Whether to do augmented Hessian Newton-Raphson (True) or just -gradient/diag(hessian) (False)
        '''

        assert( ( thetype == 'boys' ) or ( thetype == 'edmiston' ) )

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <[email protected]>
#

import ctypes
import tempfile
import numpy
import h5py
from pyscf import lib

libao2mo = lib.load_library('libao2mo')

class load(object):
    '''load 2e integrals from hdf5 file

    Usage:
        with load(erifile) as eri:
            print(eri.shape)
    '''
    def __init__(self, eri, dataname='eri_mo'):
        self.eri = eri
        self.dataname = dataname
        self.feri = None

    def __enter__(self):
        if isinstance(self.eri, str):

#!/usr/bin/env python

import os
import ctypes
import _ctypes
import unittest
from functools import reduce
import numpy
import h5py
from pyscf import lib
from pyscf import scf
from pyscf import gto
from pyscf import ao2mo
import pyscf.scf._vhf as _vhf

libcvhf = lib.load_library("libcvhf")
libao2mo1 = lib.load_library("libao2mo")

mol = gto.Mole()
mol.verbose = 0
mol.output = None  #'out_h2o'
mol.atom.extend([["O", (0.0, 0.0, 0.0)], [1, (0.0, -0.757, 0.587)], [1, (0.0, 0.757, 0.587)]])

mol.basis = "cc-pvdz"
mol.build()
nao = mol.nao_nr()
naopair = nao * (nao + 1) // 2
numpy.random.seed(15)
mo = numpy.random.random((nao, nao))
mo = mo.copy(order="F")