#!/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>
#

import sys
import ctypes
import math
import numpy
from pyscf import lib

libfci = lib.load_library('libfci')

def gen_strings4orblist(orb_list, nelec):
    '''Generate string from the given orbital list.

    Returns:
        list of int64.  One int64 element represents one string in binary format.
        The binary format takes the convention that the one bit stands for one
        orbital, bit-1 means occupied and bit-0 means unoccupied.  The lowest
        (right-most) bit corresponds to the lowest orbital in the orb_list.

    Exampels:

    >>> [bin(x) for x in gen_strings4orblist((0,1,2,3),2)]
    [0b11, 0b101, 0b110, 0b1001, 0b1010, 0b1100]
    >>> [bin(x) for x in gen_strings4orblist((3,1,0,2),2)]
    [0b1010, 0b1001, 0b11, 0b1100, 0b110, 0b101]
    '''
    orb_list = list(orb_list)
    if len(orb_list) > 63:
        return _gen_occslst(orb_list, nelec)

    assert(nelec >= 0)
    if nelec == 0:
        return [0]
    elif nelec > len(orb_list):
        return []
    def gen_str_iter(orb_list, nelec):
        if nelec == 1:
            res = [(1<<i) for i in orb_list]
        elif nelec >= len(orb_list):
            n = 0
            for i in orb_list:
                n = n | (1<<i)
            res = [n]
        else:
            restorb = orb_list[:-1]
            thisorb = 1 << orb_list[-1]
            res = gen_str_iter(restorb, nelec)
            for n in gen_str_iter(restorb, nelec-1):
                res.append(n | thisorb)
        return res
    strings = gen_str_iter(orb_list, nelec)
    assert(strings.__len__() == num_strings(len(orb_list),nelec))
    return numpy.asarray(strings, dtype=numpy.int64)

def _gen_occslst(orb_list, nelec):
    '''Generate occupied orbital list for each string.
    '''
    orb_list = list(orb_list)
    assert(nelec >= 0)
    if nelec == 0:
        return numpy.zeros((1,nelec), dtype=numpy.int32)
    elif nelec > len(orb_list):
        return numpy.zeros((0,nelec), dtype=numpy.int32)
    def gen_occs_iter(orb_list, nelec):
        if nelec == 1:
            res = [[i] for i in orb_list]
        elif nelec >= len(orb_list):
            res = [orb_list]
        else:
            restorb = orb_list[:-1]
            thisorb = orb_list[-1]
            res = gen_occs_iter(restorb, nelec)
            for n in gen_occs_iter(restorb, nelec-1):
                res.append(n + [thisorb])
        return res
    occslst = gen_occs_iter(orb_list, nelec)
    return numpy.asarray(occslst, dtype=numpy.int32).view(OIndexList)
def _strs2occslst(strs, norb):
    na = len(strs)
    one_particle_strs = numpy.asarray([1<<i for i in range(norb)])
    occ_masks = (strs.reshape(-1,1) & one_particle_strs) != 0
    occslst = numpy.where(occ_masks)[1].reshape(na,-1)
    return numpy.asarray(occslst, dtype=numpy.int32).view(OIndexList)
def _occslst2strs(occslst):
    na, nelec = occslst.shape
    strs = numpy.zeros(na, dtype=numpy.int64)
    for i in range(nelec):
        strs ^= 1 << occslst[:,i]
    return strs
class OIndexList(numpy.ndarray):
    pass

def num_strings(n, m):
    if m < 0 or m > n:
        return 0
    else:
        return math.factorial(n) // (math.factorial(n-m)*math.factorial(m))

def gen_linkstr_index_o0(orb_list, nelec, strs=None):
    if strs is None:
        strs = gen_strings4orblist(orb_list, nelec)
    strdic = dict(zip(strs,range(strs.__len__())))
    def propgate1e(str0):
        occ = []
        vir = []
        for i in orb_list:
            if str0 & (1<<i):
                occ.append(i)
            else:
                vir.append(i)
        linktab = []
        for i in occ:
            linktab.append((i, i, strdic[str0], 1))
        for i in occ:
            for a in vir:
                str1 = str0 ^ (1<<i) | (1<<a)
                # [cre, des, target_address, parity]
                linktab.append((a, i, strdic[str1], cre_des_sign(a, i, str0)))
        return linktab

    t = [propgate1e(s) for s in strs.astype(numpy.int64)]
    return numpy.array(t, dtype=numpy.int32)

def gen_linkstr_index_o1(orb_list, nelec, strs=None, tril=False):
    if nelec == 0:
        return numpy.zeros((0,0,4), dtype=numpy.int32)

    if strs is None:
        strs = _gen_occslst(orb_list, nelec)
    occslst = strs

    orb_list = numpy.asarray(orb_list)
    norb = len(orb_list)
    assert(numpy.all(numpy.arange(norb) == orb_list))

    strdic = dict((tuple(s), i) for i,s in enumerate(occslst))
    nvir = norb - nelec
    def propgate1e(str0):
        addr0 = strdic[tuple(str0)]
        tab = numpy.empty((nelec,4), dtype=numpy.int32)
        tab[:,0] = tab[:,1] = str0
        tab[:,2] = addr0
        tab[:,3] = 1
        linktab = [tab]

        virmask = numpy.ones(norb, dtype=bool)
        virmask[str0] = False
        vir = orb_list[virmask]
        str0 = numpy.asarray(str0)
        # where to put vir-orb, ie how many occ-orb in the left
        where_vir = numpy.sum(str0.reshape(-1,1) < vir, axis=0)
        parity_occ_orb = 1  # parity for annihilating occupied orbital
        for n,i in enumerate(str0):  # loop over all occupied orbitals
            # o,v which index is bigger, to determine whether to annihilate occ-orb first 
            reorder_to_ov = vir > i
            str1s = numpy.empty((nvir,nelec), dtype=int)
            str1s[:] = str0
            str1s[:,n] = vir
            str1s.sort(axis=1)
            addr = [strdic[tuple(s)] for s in str1s]
            parity = (where_vir + reorder_to_ov + 1) % 2 #? +1 so that even parity has +1, odd parity = 0
            parity[parity == 0] = -1
            parity *= parity_occ_orb
            tab = numpy.empty((nvir,4), dtype=numpy.int32)
            tab[:,0] = vir
            tab[:,1] = i
            tab[:,2] = addr
            tab[:,3] = parity
            linktab.append(tab)
            parity_occ_orb *= -1
        return numpy.vstack(linktab)

    lidx = [propgate1e(s) for s in occslst]
    lidx = numpy.asarray(lidx, dtype=numpy.int32)
    if tril:
        lidx = reform_linkstr_index(lidx)
    return lidx

# return [cre, des, target_address, parity]
def gen_linkstr_index(orb_list, nocc, strs=None, tril=False):
    '''Look up table, for the strings relationship in terms of a
    creation-annihilating operator pair.

    For given string str0, index[str0] is (nocc+nocc*nvir) x 4 array.
    The first nocc rows [i(:occ),i(:occ),str0,sign] are occupied-occupied
    excitations, which do not change the string. The next nocc*nvir rows
    [a(:vir),i(:occ),str1,sign] are occupied-virtual exciations, starting from
    str0, annihilating i, creating a, to get str1.
    '''
    if strs is None:
        strs = gen_strings4orblist(orb_list, nocc)

    if isinstance(strs, OIndexList):
        return gen_linkstr_index_o1(orb_list, nocc, strs, tril)

    strs = numpy.array(strs, dtype=numpy.int64)
    assert(all(strs[:-1] < strs[1:]))
    norb = len(orb_list)
    nvir = norb - nocc
    na = strs.shape[0]
    link_index = numpy.empty((na,nocc*nvir+nocc,4), dtype=numpy.int32)
    libfci.FCIlinkstr_index(link_index.ctypes.data_as(ctypes.c_void_p),
                            ctypes.c_int(norb), ctypes.c_int(na),
                            ctypes.c_int(nocc),
                            strs.ctypes.data_as(ctypes.c_void_p),
                            ctypes.c_int(tril))
    return link_index

def reform_linkstr_index(link_index):
    '''Compress the (a, i) pair index in linkstr_index to a lower triangular
    index, to match the 4-fold symmetry of integrals.
    '''
    #for k, tab in enumerate(link_index):
    #    for j, (a, i, str1, sign) in enumerate(tab):
    #        if a > i:
    #            ai = a*(a+1)//2+i
    #        else:
    #            ai = i*(i+1)//2+a
    #        link_new[k,j] = (ai,0,str1,sign)
    link_new = link_index.copy()
    a = link_index[:,:,0]
    i = link_index[:,:,1]
    link_new[:,:,0] = numpy.maximum(a*(a+1)//2+i, i*(i+1)//2+a)
    link_new[:,:,1] = 0
    return link_new

def gen_linkstr_index_trilidx(orb_list, nocc, strs=None):
    r'''Generate linkstr_index with the assumption that :math:`p^+ q|0\rangle`
    where :math:`p > q`.
    So the resultant link_index has the structure ``[pq, *, str1, sign]``.
    It is identical to a call to ``reform_linkstr_index(gen_linkstr_index(...))``.
    '''
    return gen_linkstr_index(orb_list, nocc, strs, True)

# return [cre, des, target_address, parity]
def gen_cre_str_index_o0(orb_list, nelec):
    '''Slow version of gen_cre_str_index function'''
    cre_strs = gen_strings4orblist(orb_list, nelec+1)
    if isinstance(cre_strs, OIndexList):
        raise NotImplementedError('System with 64 orbitals or more')

    credic = dict(zip(cre_strs,range(cre_strs.__len__())))
    def progate1e(str0):
        linktab = []
        for i in orb_list:
            if not str0 & (1<<i):
                str1 = str0 | (1<<i)
                linktab.append((i, 0, credic[str1], cre_sign(i, str0)))
        return linktab

    strs = gen_strings4orblist(orb_list, nelec)
    t = [progate1e(s) for s in strs.astype(numpy.int64)]
    return numpy.array(t, dtype=numpy.int32)
def gen_cre_str_index_o1(orb_list, nelec):
    '''C implementation of gen_cre_str_index function'''
    norb = len(orb_list)
    assert(nelec < norb)
    strs = gen_strings4orblist(orb_list, nelec)
    if isinstance(strs, OIndexList):
        raise NotImplementedError('System with 64 orbitals or more')

    strs = numpy.array(strs, dtype=numpy.int64)
    na = strs.shape[0]
    link_index = numpy.empty((len(strs),norb-nelec,4), dtype=numpy.int32)
    libfci.FCIcre_str_index(link_index.ctypes.data_as(ctypes.c_void_p),
                            ctypes.c_int(norb), ctypes.c_int(na),
                            ctypes.c_int(nelec),
                            strs.ctypes.data_as(ctypes.c_void_p))
    return link_index
def gen_cre_str_index(orb_list, nelec):
    '''linkstr_index to map between N electron string to N+1 electron string.
    It maps the given string to the address of the string which is generated by
    the creation operator.

    For given string str0, index[str0] is nvir x 4 array.  Each entry
    [i(cre),--,str1,sign] means starting from str0, creating i, to get str1.
    '''
    return gen_cre_str_index_o1(orb_list, nelec)

# return [cre, des, target_address, parity]
def gen_des_str_index_o0(orb_list, nelec):
    '''Slow version of gen_des_str_index function'''
    des_strs = gen_strings4orblist(orb_list, nelec-1)
    if isinstance(des_strs, OIndexList):
        raise NotImplementedError('System with 64 orbitals or more')

    desdic = dict(zip(des_strs,range(des_strs.__len__())))
    def progate1e(str0):
        linktab = []
        for i in orb_list:
            if str0 & (1<<i):
                str1 = str0 ^ (1<<i)
                linktab.append((0, i, desdic[str1], des_sign(i, str0)))
        return linktab

    strs = gen_strings4orblist(orb_list, nelec)
    t = [progate1e(s) for s in strs.astype(numpy.int64)]
    return numpy.array(t, dtype=numpy.int32)
def gen_des_str_index_o1(orb_list, nelec):
    '''C implementation of gen_des_str_index function'''
    assert(nelec > 0)
    strs = gen_strings4orblist(orb_list, nelec)
    if isinstance(strs, OIndexList):
        raise NotImplementedError('System with 64 orbitals or more')

    strs = numpy.array(strs, dtype=numpy.int64)
    norb = len(orb_list)
    na = strs.shape[0]
    link_index = numpy.empty((len(strs),nelec,4), dtype=numpy.int32)
    libfci.FCIdes_str_index(link_index.ctypes.data_as(ctypes.c_void_p),
                            ctypes.c_int(norb), ctypes.c_int(na),
                            ctypes.c_int(nelec),
                            strs.ctypes.data_as(ctypes.c_void_p))
    return link_index
def gen_des_str_index(orb_list, nelec):
    '''linkstr_index to map between N electron string to N-1 electron string.
    It maps the given string to the address of the string which is generated by
    the annihilation operator.

    For given string str0, index[str0] is nvir x 4 array.  Each entry
    [--,i(des),str1,sign] means starting from str0, annihilating i, to get str1.
    '''
    return gen_des_str_index_o1(orb_list, nelec)



# Determine the sign of  p^+ q |string0>
def cre_des_sign(p, q, string0):
    if p == q:
        return 1
    else:
        if (string0 & (1<<p)) or (not (string0 & (1<<q))):
            return 0
        elif p > q:
            mask = (1 << p) - (1 << (q+1))
        else:
            mask = (1 << q) - (1 << (p+1))
        return (-1) ** bin(string0 & mask).count('1')

# Determine the sign of  p^+ |string0>
def cre_sign(p, string0):
    if (string0 & (1<<p)):
        return 0
    else:
        return (-1) ** bin(string0>>(p+1)).count('1')

# Determine the sign of  p |string0>
def des_sign(p, string0):
    if (not (string0 & (1<<p))):
        return 0
    else:
        return (-1) ** bin(string0>>(p+1)).count('1')

# Determine the sign of  string1 = p^+ q |string0>
def parity(string0, string1):
    sys.stderr.write('Function cistring.parity is deprecated\n')
    ss = string1 - string0
    def count_bit1(n):
        # see Hamming weight problem and K&R C program
        return bin(n).count('1')
    if ss > 0:
        # string1&ss gives the number of 1s between two strings
        return (-1) ** (count_bit1(string1&ss))
    elif ss == 0:
        return 1
    else:
        return (-1) ** (count_bit1(string0&(-ss)))

def addr2str_o0(norb, nelec, addr):
    assert(num_strings(norb, nelec) > addr)
    if addr == 0 or nelec == norb or nelec == 0:
        return (1<<nelec) - 1   # ..0011..11
    else:
        for i in reversed(range(norb)):
            addrcum = num_strings(i, nelec)
            if addrcum <= addr:
                return (1<<i) | addr2str_o0(i, nelec-1, addr-addrcum)
def addr2str_o1(norb, nelec, addr):
    assert(num_strings(norb, nelec) > addr)
    if addr == 0 or nelec == norb or nelec == 0:
        return (1<<nelec) - 1   # ..0011..11
    str1 = 0
    nelec_left = nelec
    for norb_left in reversed(range(norb)):
        addrcum = num_strings(norb_left, nelec_left)
        if nelec_left == 0:
            break
        elif addr == 0:
            str1 |= (1<<nelec_left) - 1
            break
        elif addrcum <= addr:
            str1 |= 1<<norb_left
            addr -= addrcum
            nelec_left -= 1
    return str1
def addr2str(norb, nelec, addr):
    '''Convert CI determinant address to string'''
    return addrs2str(norb, nelec, [addr])[0]
def addrs2str(norb, nelec, addrs):
    '''Convert a list of CI determinant address to string'''
    #return [addr2str_o1(norb, nelec, addr) for addr in addrs]
    addrs = numpy.asarray(addrs, dtype=numpy.int32)
    assert(all(num_strings(norb, nelec) > addrs))
    count = addrs.size
    strs = numpy.empty(count, dtype=numpy.int64)
    libfci.FCIaddrs2str(strs.ctypes.data_as(ctypes.c_void_p),
                        addrs.ctypes.data_as(ctypes.c_void_p),
                        ctypes.c_int(count),
                        ctypes.c_int(norb), ctypes.c_int(nelec))
    return strs

#def str2addr_o0(norb, nelec, string):
#    if norb <= nelec or nelec == 0:
#        return 0
#    elif (1<<(norb-1)) & string:  # remove the first bit
#        return num_strings(norb-1, nelec) \
#                + str2addr_o0(norb-1, nelec-1, string^(1<<(norb-1)))
#    else:
#        return str2addr_o0(norb-1, nelec, string)
#def str2addr_o1(norb, nelec, string):
#    #TODO: assert norb > first-bit-in-string, nelec == num-1-in-string
#    addr = 0
#    nelec_left = nelec
#    for norb_left in reversed(range(norb)):
#        if nelec_left == 0 or norb_left < nelec_left:
#            break
#        elif (1<<norb_left) & string:
#            addr += num_strings(norb_left, nelec_left)
#            nelec_left -= 1
#    return addr
def str2addr(norb, nelec, string):
    '''Convert string to CI determinant address'''
    if isinstance(string, str):
        assert(string.count('1') == nelec)
        string = int(string, 2)
    else:
        assert(bin(string).count('1') == nelec)
    libfci.FCIstr2addr.restype = ctypes.c_int
    return libfci.FCIstr2addr(ctypes.c_int(norb), ctypes.c_int(nelec),
                              ctypes.c_ulonglong(string))
def strs2addr(norb, nelec, strings):
    '''Convert a list of string to CI determinant address'''
    strings = numpy.asarray(strings, dtype=numpy.int64)
    count = strings.size
    addrs = numpy.empty(count, dtype=numpy.int32)
    libfci.FCIstrs2addr(addrs.ctypes.data_as(ctypes.c_void_p),
                        strings.ctypes.data_as(ctypes.c_void_p),
                        ctypes.c_int(count),
                        ctypes.c_int(norb), ctypes.c_int(nelec))
    return addrs

def sub_addrs(norb, nelec, orbital_indices, sub_nelec=0):
    '''The addresses of the determinants which include the specified orbital
    indices. The size of the returned addresses is equal to the number of
    determinants of (norb, nelec) system.
    '''
    assert(norb < 63)
    if sub_nelec == 0:
        strs = gen_strings4orblist(orbital_indices, nelec)
        return strs2addr(norb, nelec, strs)
    else:
        strs = gen_strings4orblist(range(norb), nelec)
        counts = numpy.zeros(len(strs), dtype=int)
        for i in orbital_indices:
            counts += (strs & (1<<i)) != 0
        sub_strs = strs[counts == sub_nelec]
        return strs2addr(norb, nelec, sub_strs)

def tn_strs(norb, nelec, n):
    '''Generate strings for Tn amplitudes.  Eg n=1 (T1) has nvir*nocc strings,
    n=2 (T2) has nvir*(nvir-1)/2 * nocc*(nocc-1)/2 strings.
    '''
    if nelec < n or norb-nelec < n:
        return numpy.zeros(0, dtype=int)
    occs_allow = numpy.asarray(gen_strings4orblist(range(nelec), n)[::-1])
    virs_allow = numpy.asarray(gen_strings4orblist(range(nelec,norb), n))
    hf_str = int('1'*nelec, 2)
    tns = (hf_str | virs_allow.reshape(-1,1)) ^ occs_allow
    return tns.ravel()

if __name__ == '__main__':
    #print([bin(i) for i in gen_strings4orblist(range(2,5), 2)])
    #print(gen_strings4orblist(range(4), 2))
    #print(gen_linkstr_index(range(6), 3))
#    index = gen_linkstr_index(range(8), 4)
#    idx16 = index[:16]
#    print(idx16[:,:,2])
    strs = gen_strings4orblist(range(8), 4)
    occlst = _gen_occslst(range(8), 4)
    print(abs(occlst - _strs2occslst(strs, 8)).sum())
    print(abs(strs - _occslst2strs(occlst)).sum())
    tab1 = gen_linkstr_index_o0(range(8), 4)
    tab2 = gen_linkstr_index(range(8), 4)
    print(abs(tab1 - tab2).sum())
    tab3 = gen_linkstr_index_o1(range(8), 4)
    print(abs(tab1 - tab3).sum())

    print(addr2str_o0(6, 3, 7) - addr2str(6, 3, 7))
    print(addr2str_o0(6, 3, 8) - addr2str(6, 3, 8))
    print(addr2str_o0(7, 4, 9) - addr2str(7, 4, 9))

    print(str2addr(6, 3, addr2str(6, 3, 7)) - 7)
    print(str2addr(6, 3, addr2str(6, 3, 8)) - 8)
    print(str2addr(7, 4, addr2str(7, 4, 9)) - 9)
    print(numpy.allclose(numpy.arange(20),
                         strs2addr(6, 3, addrs2str(6, 3, range(20)))))

    tab1 = gen_cre_str_index_o0(range(8), 4)
    tab2 = gen_cre_str_index_o1(range(8), 4)
    print(abs(tab1 - tab2).sum())
    tab1 = gen_des_str_index_o0(range(8), 4)
    tab2 = gen_des_str_index_o1(range(8), 4)
    print(abs(tab1 - tab2).sum())
