#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# 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.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#         Alexander Sokolov <alexander.y.sokolov@gmail.com>
#

'''
Selected CI using Heat-Bath CI algorithm
(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')

    libhci.contract_h_c(h1.ctypes.data_as(ctypes.c_void_p), 
                        eri.ctypes.data_as(ctypes.c_void_p), 
                        ctypes.c_int(norb), 
                        ctypes.c_int(nelec[0]), 
                        ctypes.c_int(nelec[1]), 
                        strs.ctypes.data_as(ctypes.c_void_p), 
                        civec.ctypes.data_as(ctypes.c_void_p), 
                        hdiag.ctypes.data_as(ctypes.c_void_p), 
                        ctypes.c_ulonglong(ndet), 
                        ci1.ctypes.data_as(ctypes.c_void_p))

    return ci1

def contract_2e(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)

    eri = eri.reshape([norb]*4)

    for ip in range(ndet):
        for jp in range(ip):
            stria, strib = strs[ip].reshape(2,-1)
            strja, strjb = strs[jp].reshape(2,-1)
            desa, crea = str_diff(stria, strja)
            if len(desa) > 2:
                continue
            desb, creb = str_diff(strib, strjb)
            if len(desb) + len(desa) > 2:
                continue
            if len(desa) + len(desb) == 1:
# alpha->alpha
                if len(desb) == 0:
                    i,a = desa[0], crea[0]
                    occsa = str2orblst(stria, norb)[0]
                    occsb = str2orblst(strib, norb)[0]
                    fai = h1[a,i]
                    for k in occsa:
                        fai += eri[k,k,a,i] - eri[k,i,a,k]
                    for k in occsb:
                        fai += eri[k,k,a,i]
                    sign = cre_des_sign(a, i, stria)
                    ci1[jp] += sign * fai * civec[ip]
                    ci1[ip] += sign * fai * civec[jp]
# beta ->beta
                elif len(desa) == 0:
                    i,a = desb[0], creb[0]
                    occsa = str2orblst(stria, norb)[0]
                    occsb = str2orblst(strib, norb)[0]
                    fai = h1[a,i]
                    for k in occsb:
                        fai += eri[k,k,a,i] - eri[k,i,a,k]
                    for k in occsa:
                        fai += eri[k,k,a,i]
                    sign = cre_des_sign(a, i, strib)
                    ci1[jp] += sign * fai * civec[ip]
                    ci1[ip] += sign * fai * civec[jp]

            else:
# alpha,alpha->alpha,alpha
                if len(desb) == 0:
                    i,j = desa
                    a,b = crea
# 6 conditions for i,j,a,b
# --++, ++--, -+-+, +-+-, -++-, +--+ 
                    if a > j or i > b:
# condition --++, ++--
                        v = eri[a,j,b,i]-eri[a,i,b,j]
                        sign = cre_des_sign(b, i, stria)
                        sign*= cre_des_sign(a, j, stria)
                    else:
# condition -+-+, +-+-, -++-, +--+ 
                        v = eri[a,i,b,j]-eri[a,j,b,i]
                        sign = cre_des_sign(b, j, stria)
                        sign*= cre_des_sign(a, i, stria)
                    ci1[jp] += sign * v * civec[ip]
                    ci1[ip] += sign * v * civec[jp]
# beta ,beta ->beta ,beta
                elif len(desa) == 0:
                    i,j = desb
                    a,b = creb
                    if a > j or i > b:
                        v = eri[a,j,b,i]-eri[a,i,b,j]
                        sign = cre_des_sign(b, i, strib)
                        sign*= cre_des_sign(a, j, strib)
                    else:
                        v = eri[a,i,b,j]-eri[a,j,b,i]
                        sign = cre_des_sign(b, j, strib)
                        sign*= cre_des_sign(a, i, strib)
                    ci1[jp] += sign * v * civec[ip]
                    ci1[ip] += sign * v * civec[jp]
# alpha,beta ->alpha,beta
                else:
                    i,a = desa[0], crea[0]
                    j,b = desb[0], creb[0]
                    v = eri[a,i,b,j]
                    sign = cre_des_sign(a, i, stria)
                    sign*= cre_des_sign(b, j, strib)
                    ci1[jp] += sign * v * civec[ip]
                    ci1[ip] += sign * v * civec[jp]
        ci1[ip] += hdiag[ip] * civec[ip]

    return ci1

def spin_square(civec, norb, nelec):
    ss = numpy.dot(civec.T, contract_ss(civec, norb, nelec))
    s = numpy.sqrt(ss+.25) - .5
    multip = s*2+1
    return ss, multip

def contract_ss(civec, norb, nelec):
    strs = civec._strs
    ndet = len(strs)
    ci1 = numpy.zeros_like(civec)

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

    libhci.contract_ss_c(ctypes.c_int(norb), 
                        ctypes.c_int(nelec[0]), 
                        ctypes.c_int(nelec[1]), 
                        strs.ctypes.data_as(ctypes.c_void_p), 
                        civec.ctypes.data_as(ctypes.c_void_p), 
                        ctypes.c_ulonglong(ndet), 
                        ci1.ctypes.data_as(ctypes.c_void_p))

    return ci1

def make_hdiag(h1e, eri, strs, norb, nelec):
    eri = ao2mo.restore(1, eri, norb)
    diagj = numpy.einsum('iijj->ij',eri)
    diagk = numpy.einsum('ijji->ij',eri)

    ndet = len(strs)
    hdiag = numpy.zeros(ndet)
    for idet, (stra, strb) in enumerate(strs.reshape(ndet,2,-1)):
        aocc = str2orblst(stra, norb)[0]
        bocc = str2orblst(strb, norb)[0]
        e1 = h1e[aocc,aocc].sum() + h1e[bocc,bocc].sum()
        e2 = diagj[aocc][:,aocc].sum() + diagj[aocc][:,bocc].sum() \
           + diagj[bocc][:,aocc].sum() + diagj[bocc][:,bocc].sum() \
           - diagk[aocc][:,aocc].sum() - diagk[bocc][:,bocc].sum()
        hdiag[idet] = e1 + e2*.5
    return hdiag

def cre_des_sign(p, q, string):
    nset = len(string)
    pg, pb = p//64, p%64
    qg, qb = q//64, q%64

    if pg > qg:
        n1 = 0
        for i in range(nset-pg, nset-qg-1):
            n1 += bin(string[i]).count('1')
        n1 += bin(string[-1-pg] & numpy.uint64((1<<pb) - 1)).count('1')
        n1 += string[-1-qg] >> numpy.uint64(qb+1)
    elif pg < qg:
        n1 = 0
        for i in range(nset-qg, nset-pg-1):
            n1 += bin(string[i]).count('1')
        n1 += bin(string[-1-qg] & numpy.uint64((1<<qb) - 1)).count('1')
        n1 += string[-1-pg] >> numpy.uint64(pb+1)
    else:
        if p > q:
            mask = numpy.uint64((1 << pb) - (1 << (qb+1)))
        else:
            mask = numpy.uint64((1 << qb) - (1 << (pb+1)))
        n1 = bin(string[-1-pg]&mask).count('1')

    if n1 % 2:
        return -1
    else:
        return 1

def argunique(strs):
    def order(x, y):
        for i in range(y.size):
            if x[i] > y[i]:
                return 1
            elif y[i] > x[i]:
                return -1
        return 0
    def qsort_idx(idx):
        nstrs = len(idx)
        if nstrs <= 1:
            return idx
        else:
            ref = idx[-1]
            group_lt = []
            group_gt = []
            for i in idx[:-1]:
                c = order(strs[i], strs[ref])
                if c == -1:
                    group_lt.append(i)
                elif c == 1:
                    group_gt.append(i)
            return qsort_idx(group_lt) + [ref] + qsort_idx(group_gt)
    return qsort_idx(range(len(strs)))

def argunique_ctypes(strs):
    nstrs, nset = strs.shape

    sort_idx = numpy.empty(nstrs, dtype=numpy.uint64)

    strs = numpy.asarray(strs, order='C')
    sort_idx = numpy.asarray(sort_idx, order='C')

    nstrs_ = numpy.array([nstrs])

    libhci.argunique(strs.ctypes.data_as(ctypes.c_void_p), 
                     sort_idx.ctypes.data_as(ctypes.c_void_p), 
                     nstrs_.ctypes.data_as(ctypes.c_void_p), 
                     ctypes.c_int(nset))

    sort_idx = sort_idx[:nstrs_[0]]

    return sort_idx.tolist()

def str_diff(string0, string1):
    des_string0 = []
    cre_string0 = []
    nset = len(string0)
    off = 0
    for i in reversed(range(nset)):
        df = string0[i] ^ string1[i]
        des_string0.extend([x+off for x in find1(df & string0[i])])
        cre_string0.extend([x+off for x in find1(df & string1[i])])
        off += 64
    return des_string0, cre_string0

def excitation_level(string, nelec=None):
    nset = len(string)
    if nelec is None:
        nelec = 0
        for i in range(nset):
            nelec += bin(string[i]).count('1')

    g, b = nelec//64, nelec%64
    tn = nelec - bin(string[-1-g])[-b:].count('1')
    for s in string[nset-g:]:
        tn -= bin(s).count('1')
    return tn

def find1(s):
    return [i for i,x in enumerate(bin(s)[2:][::-1]) if x is '1']

def toggle_bit(s, place):
    nset = len(s)
    g, b = place//64, place%64
    s[-1-g] ^= numpy.uint64(1<<b)
    return s

def select_strs_ctypes(myci, civec, h1, eri, jk, eri_sorted, jk_sorted, norb, nelec):
    strs = civec._strs
    ndet = strs.shape[0]
    ndet, nset = strs.shape
    nset = nset // 2
    neleca, nelecb = nelec

    h1 = numpy.asarray(h1, order='C')  
    eri = numpy.asarray(eri, order='C')
    jk = numpy.asarray(jk, order='C')
    civec = numpy.asarray(civec, order='C')
    strs = numpy.asarray(strs, order='C')
    eri_sorted = numpy.asarray(eri_sorted, order='C')
    jk_sorted = numpy.asarray(jk_sorted, order='C')

    str_add = numpy.empty((0,strs.shape[1]), dtype=numpy.uint64)

    batch_size = max(1, 8 * 4 * neleca * nelecb * (norb-neleca) * (norb-nelecb))
    ndet_batch = int(myci.max_memory * 1024**2) // batch_size
    nbatches = ndet // ndet_batch + 1

    for i in range(nbatches):
        ndet_start = ndet_batch * i
        ndet_finish = min(ndet_batch * (i + 1), ndet)
        ndet_select_max = 4 * neleca * nelecb * (norb-neleca) * (norb-nelecb) * ndet_batch

        str_add_batch = numpy.empty((ndet_select_max, strs.shape[1]), dtype=numpy.uint64)
        n_str_add_batch = numpy.array([str_add_batch.shape[0]])

        str_add_batch = numpy.asarray(str_add_batch, order='C')

        libhci.select_strs(h1.ctypes.data_as(ctypes.c_void_p), 
                           eri.ctypes.data_as(ctypes.c_void_p), 
                           jk.ctypes.data_as(ctypes.c_void_p), 
                           eri_sorted.ctypes.data_as(ctypes.c_void_p), 
                           jk_sorted.ctypes.data_as(ctypes.c_void_p), 
                           ctypes.c_int(norb), 
                           ctypes.c_int(neleca), 
                           ctypes.c_int(nelecb), 
                           strs.ctypes.data_as(ctypes.c_void_p), 
                           civec.ctypes.data_as(ctypes.c_void_p), 
                           ctypes.c_ulonglong(ndet_start), 
                           ctypes.c_ulonglong(ndet_finish), 
                           ctypes.c_double(myci.select_cutoff),
                           str_add_batch.ctypes.data_as(ctypes.c_void_p),
                           n_str_add_batch.ctypes.data_as(ctypes.c_void_p))

        n_str_add_batch = n_str_add_batch[0]
        str_add_batch = str_add_batch[:n_str_add_batch]
        str_add = numpy.vstack((str_add, str_add_batch))

    str_add = numpy.asarray(str_add)
    return str_add

def enlarge_space(myci, civec, h1, eri, jk, eri_sorted, jk_sorted, norb, nelec):
    if not isinstance(civec, (tuple, list)):
        civec = [civec]

    strs = civec[0]._strs

    nroots = len(civec)

    cidx = abs(civec[0]) > myci.ci_coeff_cutoff
    for p in range(1,nroots):
        cidx += abs(civec[p]) > myci.ci_coeff_cutoff

    strs = strs[cidx]

    ci_coeff = [as_SCIvector(c[cidx], strs) for c in civec]
 
    strs_new = strs.copy()

    for p in range(nroots):
        str_add = select_strs_ctypes(myci, ci_coeff[p], h1, eri, jk, eri_sorted, jk_sorted, norb, nelec)
        strs_new = numpy.vstack((strs, str_add))

    # Add strings together and remove duplicate strings
    tmp = numpy.ascontiguousarray(strs_new).view(numpy.dtype((numpy.void, strs_new.dtype.itemsize * strs_new.shape[1])))
    _, tmpidx = numpy.unique(tmp, return_index=True)

    new_ci = []
    for p in range(nroots):
        c = numpy.zeros(strs_new.shape[0])
        c[:ci_coeff[p].shape[0]] = ci_coeff[p]
        new_ci.append(c[tmpidx])

    strs_new = strs_new[tmpidx]

    return [as_SCIvector(ci, strs_new) for ci in new_ci]

def str2orblst(string, norb):
    occ = []
    vir = []
    nset = len(string)
    off = 0
    for k in reversed(range(nset)):
        s = string[k]
        occ.extend([x+off for x in find1(s)])
        for i in range(0, min(64, norb-off)): 
            if not (s & numpy.uint64(1<<i)):
                vir.append(i+off)
        off += 64
    return occ, vir

def orblst2str(lst, norb):
    nset = (norb+63) // 64
    string = numpy.zeros(nset, dtype=numpy.uint64)
    for i in lst:
        toggle_bit(string, i)
    return string

def kernel_float_space(myci, h1e, eri, norb, nelec, ci0=None,
                       tol=None, lindep=None, max_cycle=None, max_space=None,
                       nroots=None, davidson_only=None, max_iter=None,
                       max_memory=None, verbose=None, ecore=0, return_integrals=False, 
                       eri_sorted=None, jk=None, jk_sorted=None, **kwargs):
    if verbose is None:
        log = logger.Logger(myci.stdout, myci.verbose)
    elif isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(myci.stdout, verbose)
    if tol is None: tol = myci.conv_tol
    if lindep is None: lindep = myci.lindep
    if max_cycle is None: max_cycle = myci.max_cycle
    if max_space is None: max_space = myci.max_space
    if max_memory is None: max_memory = myci.max_memory
    if nroots is None: nroots = myci.nroots
    if max_iter is None: max_iter = myci.max_iter
    if myci.verbose >= logger.WARN:
        myci.check_sanity()

    log.info('\nStarting heat-bath CI algorithm...\n')
    log.info('Selection threshold:                  %8.5e',    myci.select_cutoff)
    log.info('CI coefficient cutoff:                %8.5e',    myci.ci_coeff_cutoff)
    log.info('Energy convergence tolerance:         %8.5e',    tol)
    log.info('Number of determinants tolerance:     %8.5e',    myci.conv_ndet_tol)
    log.info('Number of electrons:                  %s',       nelec)
    log.info('Number of orbitals:                   %3d',      norb)
    log.info('Number of roots:                      %3d',    nroots)

    nelec = direct_spin1._unpack_nelec(nelec, myci.spin)
    eri = ao2mo.restore(1, eri, norb)

    # Avoid resorting the integrals by storing them in memory
    eri = eri.ravel()

    if eri_sorted is None and jk is None and jk_sorted is None:
        log.debug("\nSorting two-electron integrals...")
        t_start = time.time()
        eri_sorted = abs(eri).argsort()[::-1]
        jk = eri.reshape([norb]*4)
        jk = jk - jk.transpose(2,1,0,3)
        jk = jk.ravel()
        jk_sorted = abs(jk).argsort()[::-1]
        t_current = time.time() - t_start
        log.debug('Timing for sorting the integrals: %10.3f', t_current)

    # Initial guess
    if ci0 is None:
        hf_str = numpy.hstack([orblst2str(range(nelec[0]), norb), orblst2str(range(nelec[1]), norb)]).reshape(1,-1)
        ci0 = [as_SCIvector(numpy.ones(1), hf_str)]
    else:
        assert(nroots == len(ci0))

    ci0 = myci.enlarge_space(ci0, h1e, eri, jk, eri_sorted, jk_sorted, norb, nelec)

    def hop(c):
        hc = myci.contract_2e((h1e, eri), as_SCIvector(c, ci_strs), norb, nelec, hdiag)
        return hc.ravel()
    precond = lambda x, e, *args: x/(hdiag-e+myci.level_shift)

    e_last = 0
    float_tol = 3e-4
    conv = False
    for icycle in range(max_iter):
        ci_strs = ci0[0]._strs
        float_tol = max(float_tol*.3, tol*1e2)
        log.info('\nMacroiteration %d', icycle)
        log.info('Number of CI configurations: %d', ci_strs.shape[0])
        hdiag = myci.make_hdiag(h1e, eri, ci_strs, norb, nelec)
        t_start = time.time()
        e, ci0 = myci.eig(hop, ci0, precond, tol=float_tol, lindep=lindep,
                          max_cycle=max_cycle, max_space=max_space, nroots=nroots,
                          max_memory=max_memory, verbose=log, **kwargs)
        if not isinstance(ci0, (tuple, list)):
            ci0 = [ci0]
            e = [e]
        t_current = time.time() - t_start
        log.debug('Timing for solving the eigenvalue problem: %10.3f', t_current)
        ci0 = [as_SCIvector(c, ci_strs) for c in ci0]
        de, e_last = min(e)-e_last, min(e)
        log.info('Cycle %d  E = %s  dE = %.8g', icycle, numpy.array(e)+ecore, de)

        if abs(de) < tol*1e3:
            conv = True
            break

        last_ci0_size = float(len(ci_strs))
        t_start = time.time()
        ci0 = myci.enlarge_space(ci0, h1e, eri, jk, eri_sorted, jk_sorted, norb, nelec)
        t_current = time.time() - t_start
        log.debug('Timing for selecting configurations: %10.3f', t_current)
        if (((1 - myci.conv_ndet_tol) < len(ci0[0]._strs)/last_ci0_size < (1 + myci.conv_ndet_tol))):
            conv = True
            break

    ci_strs = ci0[0]._strs
    log.info('\nExtra CI in the final selected space')
    log.info('Number of CI configurations: %d', ci_strs.shape[0])
    hdiag = myci.make_hdiag(h1e, eri, ci_strs, norb, nelec)
    e, c = myci.eig(hop, ci0, precond, tol=tol, lindep=lindep,
                    max_cycle=max_cycle, max_space=max_space, nroots=nroots,
                    max_memory=max_memory, verbose=log, **kwargs)
    if not isinstance(c, (tuple, list)):
        c = [c]
        e = [e]
    log.info('\nSelected CI  E = %s', numpy.array(e)+ecore)

    if (return_integrals):
        return (numpy.array(e)+ecore), [as_SCIvector(ci, ci_strs) for ci in c], eri_sorted, jk, jk_sorted
    else:
        return (numpy.array(e)+ecore), [as_SCIvector(ci, ci_strs) for ci in c]

def fix_spin(myci, shift=.2, ss=None, **kwargs):
    r'''If Selected CI solver cannot stick on spin eigenfunction, modify the solver by
    adding a shift on spin square operator

    .. math::

        (H + shift*S^2) |\Psi\rangle = E |\Psi\rangle

    Args:
        myci : An instance of :class:`SelectedCI`

    Kwargs:
        shift : float
            Level shift for states which have different spin
        ss : number
            S^2 expection value == s*(s+1)

    Returns
            A modified Selected CI object based on myci.
    '''
    if 'ss_value' in kwargs:
        sys.stderr.write('fix_spin_: kwarg "ss_value" will be removed in future release. '
                         'It was replaced by "ss"\n')
        ss_value = kwargs['ss_value']
    else:
        ss_value = ss

    def contract_2e(h1_h2, civec, norb, nelec, hdiag=None, **kwargs):
        if isinstance(nelec, (int, numpy.number)):
            sz = (nelec % 2) * .5
        else:
            sz = abs(nelec[0]-nelec[1]) * .5
        if ss_value is None:
            ss = sz*(sz+1)
        else:
            ss = ss_value

        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)
        ci2 = 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')
        ci2 = numpy.asarray(ci2, order='C')

        libhci.contract_h_c_ss_c(h1.ctypes.data_as(ctypes.c_void_p), 
                                 eri.ctypes.data_as(ctypes.c_void_p), 
                                 ctypes.c_int(norb), 
                                 ctypes.c_int(nelec[0]), 
                                 ctypes.c_int(nelec[1]), 
                                 strs.ctypes.data_as(ctypes.c_void_p), 
                                 civec.ctypes.data_as(ctypes.c_void_p), 
                                 hdiag.ctypes.data_as(ctypes.c_void_p), 
                                 ctypes.c_ulonglong(ndet), 
                                 ci1.ctypes.data_as(ctypes.c_void_p),
                                 ci2.ctypes.data_as(ctypes.c_void_p))

        if ss < sz*(sz+1)+.1:
# (S^2-ss)|Psi> to shift state other than the lowest state
            ci2 -= ss * civec
        else:
# (S^2-ss)^2|Psi> to shift states except the given spin.
# It still relies on the quality of initial guess
            tmp = ci2.copy()
            tmp -= ss * civec
            ci2 = -ss * tmp
            ci2 += myci.contract_ss(as_SCIvector_if_not(tmp, strs), norb, nelec)
            tmp = None
        ci2 *= shift
        ci1 += ci2

        return as_SCIvector_if_not(ci1, strs)

    myci.contract_2e = contract_2e
    return myci

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 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 make_rdm12s(civec, norb, nelec):
    '''Spin orbital 1- and 2-particle reduced density matrices (aa, bb, aaaa, aabb, bbbb)
    '''
    strs = civec._strs
    ndet = len(strs)
    rdm1a = numpy.zeros(norb*norb)
    rdm1b = numpy.zeros(norb*norb)
    rdm2aa = numpy.zeros(norb*norb*norb*norb)
    rdm2ab = numpy.zeros(norb*norb*norb*norb)
    rdm2bb = numpy.zeros(norb*norb*norb*norb)

    civec = numpy.asarray(civec, order='C')
    strs = numpy.asarray(strs, order='C')
    rdm1a  = numpy.asarray(rdm1a, order='C')
    rdm1b  = numpy.asarray(rdm1b, order='C')
    rdm2aa = numpy.asarray(rdm2aa, order='C')
    rdm2ab = numpy.asarray(rdm2ab, order='C')
    rdm2bb = numpy.asarray(rdm2bb, order='C')

    # Compute 1- and 2-RDMs
    libhci.compute_rdm12s(ctypes.c_int(norb), 
                          ctypes.c_int(nelec[0]), 
                          ctypes.c_int(nelec[1]), 
                          strs.ctypes.data_as(ctypes.c_void_p), 
                          civec.ctypes.data_as(ctypes.c_void_p), 
                          ctypes.c_ulonglong(ndet), 
                          rdm1a.ctypes.data_as(ctypes.c_void_p),
                          rdm1b.ctypes.data_as(ctypes.c_void_p),
                          rdm2aa.ctypes.data_as(ctypes.c_void_p),
                          rdm2ab.ctypes.data_as(ctypes.c_void_p),
                          rdm2bb.ctypes.data_as(ctypes.c_void_p))

    rdm1a = rdm1a.reshape([norb]*2)
    rdm1b = rdm1b.reshape([norb]*2)
    rdm2aa = rdm2aa.reshape([norb]*4)
    rdm2ab = rdm2ab.reshape([norb]*4)
    rdm2bb = rdm2bb.reshape([norb]*4)

    # Sort 2-RDM into chemists' notation: <p_1 q_2|r_1 s_2> -> (p_1 r_1| q_2 s_2)
    rdm2aa = rdm2aa.transpose(0,2,1,3)
    rdm2ab = rdm2ab.transpose(0,2,1,3)
    rdm2bb = rdm2bb.transpose(0,2,1,3)

    return (rdm1a, rdm1b), (rdm2aa, rdm2ab, rdm2bb)

class SelectedCI(direct_spin1.FCISolver):
    def __init__(self, mol=None):
        direct_spin1.FCISolver.__init__(self, mol)
        self.ci_coeff_cutoff = .5e-3
        self.select_cutoff = .5e-3
        self.conv_tol = 1e-9
        self.conv_ndet_tol = 0.001
        self.nroots = 1
        self.max_iter = 10
        # Maximum memory in MB for storing lists of selected strings
        self.max_memory = 1000

##################################################
# don't modify the following attributes, they are not input options
        #self.converged = False
        #self.ci = None
        self._strs = None
        self._keys = set(self.__dict__.keys())

    def dump_flags(self, verbose=None):
        direct_spin1.FCISolver.dump_flags(self, verbose)
        logger.info(self, 'ci_coeff_cutoff %g', self.ci_coeff_cutoff)
        logger.info(self, 'select_cutoff   %g', self.select_cutoff)

    # define absorb_h1e for compatibility to other FCI solver
    def absorb_h1e(h1, eri, *args, **kwargs):
        return (h1, eri)

    def contract_2e(self, h1_h2, civec, norb, nelec, hdiag=None, **kwargs):
        if hasattr(civec, '_strs'):
            self._strs = civec._strs
        else:
            assert(civec.size == len(self._strs))
            civec = as_SCIvector(civec, self._strs)
        return contract_2e_ctypes(h1_h2, civec, norb, nelec, hdiag, **kwargs)
#        return contract_2e(h1_h2, civec, norb, nelec, hdiag, **kwargs)

    def contract_ss(self, civec, norb, nelec):
        if hasattr(civec, '_strs'):
            self._strs = civec._strs
        else:
            assert(civec.size == len(self._strs))
            civec = as_SCIvector(civec, self._strs)
        return contract_ss(civec, norb, nelec)

    def spin_square(self, civec, norb, nelec):
        if hasattr(civec, '_strs'):
            self._strs = civec._strs
        else:
            assert(civec.size == len(self._strs))
            civec = as_SCIvector(civec, self._strs)
        return spin_square(civec, norb, nelec)

    def make_hdiag(self, h1e, eri, strs, norb, nelec):
        return make_hdiag(h1e, eri, strs, norb, nelec)

    def to_fci(self, civec, norb, nelec):

        if hasattr(civec, '_strs'):
            self._strs = civec._strs
        else:
            assert(civec.size == len(self._strs))
            civec = as_SCIvector(civec, self._strs)

        return to_fci(civec, norb, nelec)

    def make_rdm12s(self, civec, norb, nelec):

        if hasattr(civec, '_strs'):
            self._strs = civec._strs
        else:
            assert(civec.size == len(self._strs))
            civec = as_SCIvector(civec, self._strs)

        return make_rdm12s(civec, norb, nelec)

    enlarge_space = enlarge_space
    kernel = kernel_float_space

SCI = SelectedCI


class _SCIvector(numpy.ndarray):
    def __array_finalize__(self, obj):
        self._strs = getattr(obj, '_strs', None)

def as_SCIvector(civec, ci_strs):
    civec = civec.view(_SCIvector)
    civec._strs = ci_strs
    return civec

def as_SCIvector_if_not(civec, ci_strs):
    if not hasattr(civec, '_strs'):
        civec = as_SCIvector(civec, ci_strs)
    return civec


if __name__ == '__main__':
    numpy.random.seed(3)
    strs = (numpy.random.random((14,3)) * 4).astype(numpy.uint64)
    print(strs)
    print(argunique(strs))

    norb = 6
    nelec = 3,3
    hf_str = numpy.hstack([orblst2str(range(nelec[0]), norb),
                           orblst2str(range(nelec[1]), norb)]).reshape(1,-1)
    numpy.random.seed(3)
    h1 = numpy.random.random([norb]*2)**4 * 1e-2
    h1 = h1 + h1.T
    eri = numpy.random.random([norb]*4)**4 * 1e-2
    eri = eri + eri.transpose(0,1,3,2)
    eri = eri + eri.transpose(1,0,2,3)
    eri = eri + eri.transpose(2,3,0,1)
    eri_sorted = abs(eri).argsort()[::-1]
    jk = eri.reshape([norb]*4)
    jk = jk - jk.transpose(2,1,0,3)
    jk = jk.ravel()
    jk_sorted = abs(jk).argsort()[::-1]
    ci1 = [as_SCIvector(numpy.ones(1), hf_str)]

    myci = SelectedCI()
    myci.select_cutoff = .001
    myci.ci_coeff_cutoff = .001

    ci2 = enlarge_space(myci, ci1, h1, eri, jk, eri_sorted, jk_sorted, norb, nelec)
    print(len(ci2[0]))

    ci2 = enlarge_space(myci, ci1, h1, eri, jk, eri_sorted, jk_sorted, norb, nelec)
    numpy.random.seed(1)
    ci3 = numpy.random.random(ci2[0].size)
    ci3 *= 1./numpy.linalg.norm(ci3)
    ci3 = [ci3]
    ci3 = enlarge_space(myci, ci2, h1, eri, jk, eri_sorted, jk_sorted, norb, nelec)

    efci = direct_spin1.kernel(h1, eri, norb, nelec, verbose=5)[0]

    ci4 = contract_2e_ctypes((h1, eri), ci3[0], norb, nelec)

    fci3 = to_fci(ci3, norb, nelec)
    h2e = direct_spin1.absorb_h1e(h1, eri, norb, nelec, .5)
    fci4 = direct_spin1.contract_2e(h2e, fci3, norb, nelec)
    fci4 = from_fci(fci4, ci3[0]._strs, norb, nelec)
    print(abs(ci4-fci4).sum())

    e = myci.kernel(h1, eri, norb, nelec, verbose=5)[0]
    print(e, efci)
