# Path: mmd/slater.py
# def common_index(det1,det2,Nint):
#     common = []
#     ishift = -64
#     for l in range(Nint):
#         tmp = det1[l] & det2[l]
#         ishift += 64
#         for n in range(63):
#             mask = 1 << n
#             idx_string = tmp & mask
#             if idx_string:
#                 common.append(trailz(mask) + ishift)
# 
#     return common
# 
# def get_excitation(det1,det2,Nint):
#     exc = np.zeros((3,2),dtype=np.int)
#     degree = n_excitations(det1,det2,Nint)
#     if degree > 2:
#         phase = 0
#     elif degree == 2:
#         exc, phase = get_double_excitation(det1,det2,Nint)
#     elif degree == 1:
#         exc, phase = get_single_excitation(det1,det2,Nint)
#     elif degree == 0:
#         phase = 1
# 
#     return exc, degree, phase
# 
# Path: mmd/utils/davidson.py
# def davidson(A,roots,tol=1e-8):
# 
#     mat_dim = A.shape[0]
#     sub_dim = 4*roots
#     V = np.eye(mat_dim,sub_dim) 
# 
#     converged = False
#     while not converged:
#         # subspace
#         S = np.dot(V.T,np.dot(A,V))
#  
#         # diag subspace
#         E,C = scipy.linalg.eigh(S) # note already ordered if using eigh
#         E,C = E[:roots], C[:,:roots]
#          
#         # get current approx. eigenvectors
#         X = np.dot(V,C)   
# 
#         # form residual vectors 
#         R = np.zeros((mat_dim,roots))
#         Delta = np.zeros((mat_dim,roots))
#         unconverged = []
#         for j in range(roots):
#             R[:,j] = np.dot((A - E[j] * np.eye(mat_dim)),X[:,j])
#             if np.linalg.norm(R[:,j]) > tol:
#                 unconverged.append(j)
# 
#         # check convergence
#         if len(unconverged) < 1:
#             converged = True
# 
#         for j in unconverged:
#             preconditioner = np.zeros(mat_dim)
#             for i in range(mat_dim):
#                 if np.abs(E[j] - A[i,i]) < 1e-4:
#                     continue # don't let preconditioner blow up -- keep as zero
#                 else:
#                     preconditioner[i] = -1/(A[i,i] - E[j])
# 
#             # normalize correction vectors
#             Delta[:,j] = preconditioner * R[:,j]
#             Delta[:,j] /= np.linalg.norm(Delta[:,j])
# 
#             # project corrections onto orthogonal complement
#             q = np.dot((np.eye(mat_dim)  - np.dot(V,V.T)),Delta[:,j])
#             norm = np.linalg.norm(q)
# 
#             if (norm > 1e-4):
#                 if (sub_dim + 1 > min(500,mat_dim//4)):
#                     # subspace collapse 
#                     print("Subspace too big: collapsing")
#                     print("Eigs at current: ", E)
#                     sub_dim = roots # restart uses best guess of eigenvecs
#                     V = X
#                     V, _ = np.linalg.qr(V)
#                     break
#                     
#                 else:
#                     V_copy = np.copy(V)
#                     sub_dim += 1
#                     V = np.eye(mat_dim,sub_dim) 
#                     V[:,:(sub_dim-1)] = V_copy
#                     # add new vector to end of subspace
#                     V[:,-1] = q/norm
# 
#     if converged:
#         return E, X 


import numpy as np
import scipy
import sys
from itertools import product, combinations
from bitstring import BitArray
from mmd.slater import common_index, get_excitation
from mmd.utils.davidson import davidson
from scipy.special import comb
from scipy.linalg import sqrtm, lu
                EMP2 += numer/denom
            self.mol.emp2 = EMP2 + self.mol.energy   

        print('E(MP2) = ', self.mol.emp2.real) 

    @staticmethod
    def tuple2bitstring(bit_tuple):
        ''' From tuple of occupied orbitals, return bitstring representation '''
        string = ['0']*(max(bit_tuple) + 1)
        for i in bit_tuple:
            string[i] = '1'
        string = ''.join(string[::-1])
        return BitArray(bin=string)


    def hamiltonian_matrix_element(self,det1,det2,Nint):
        """ return general Hamiltonian matrix element <det1|H|det2> """

        exc, degree, phase = get_excitation(det1,det2,Nint)

        if degree > 2:
            return 0

        elif degree == 2:
            # sign * <hole1,hole2||particle1,particle2>
            return phase * self.mol.double_bar[exc[1,0], exc[2,0], exc[1,1], exc[2,1]]

        elif degree == 1:
            m = exc[1,0]
            p = exc[1,1]

# Path: mmd/slater.py
# def common_index(det1,det2,Nint):
#     common = []
#     ishift = -64
#     for l in range(Nint):
#         tmp = det1[l] & det2[l]
#         ishift += 64
#         for n in range(63):
#             mask = 1 << n
#             idx_string = tmp & mask
#             if idx_string:
#                 common.append(trailz(mask) + ishift)
# 
#     return common
# 
# def get_excitation(det1,det2,Nint):
#     exc = np.zeros((3,2),dtype=np.int)
#     degree = n_excitations(det1,det2,Nint)
#     if degree > 2:
#         phase = 0
#     elif degree == 2:
#         exc, phase = get_double_excitation(det1,det2,Nint)
#     elif degree == 1:
#         exc, phase = get_single_excitation(det1,det2,Nint)
#     elif degree == 0:
#         phase = 1
# 
#     return exc, degree, phase
# 
# Path: mmd/utils/davidson.py
# def davidson(A,roots,tol=1e-8):
# 
#     mat_dim = A.shape[0]
#     sub_dim = 4*roots
#     V = np.eye(mat_dim,sub_dim) 
# 
#     converged = False
#     while not converged:
#         # subspace
#         S = np.dot(V.T,np.dot(A,V))
#  
#         # diag subspace
#         E,C = scipy.linalg.eigh(S) # note already ordered if using eigh
#         E,C = E[:roots], C[:,:roots]
#          
#         # get current approx. eigenvectors
#         X = np.dot(V,C)   
# 
#         # form residual vectors 
#         R = np.zeros((mat_dim,roots))
#         Delta = np.zeros((mat_dim,roots))
#         unconverged = []
#         for j in range(roots):
#             R[:,j] = np.dot((A - E[j] * np.eye(mat_dim)),X[:,j])
#             if np.linalg.norm(R[:,j]) > tol:
#                 unconverged.append(j)
# 
#         # check convergence
#         if len(unconverged) < 1:
#             converged = True
# 
#         for j in unconverged:
#             preconditioner = np.zeros(mat_dim)
#             for i in range(mat_dim):
#                 if np.abs(E[j] - A[i,i]) < 1e-4:
#                     continue # don't let preconditioner blow up -- keep as zero
#                 else:
#                     preconditioner[i] = -1/(A[i,i] - E[j])
# 
#             # normalize correction vectors
#             Delta[:,j] = preconditioner * R[:,j]
#             Delta[:,j] /= np.linalg.norm(Delta[:,j])
# 
#             # project corrections onto orthogonal complement
#             q = np.dot((np.eye(mat_dim)  - np.dot(V,V.T)),Delta[:,j])
#             norm = np.linalg.norm(q)
# 
#             if (norm > 1e-4):
#                 if (sub_dim + 1 > min(500,mat_dim//4)):
#                     # subspace collapse 
#                     print("Subspace too big: collapsing")
#                     print("Eigs at current: ", E)
#                     sub_dim = roots # restart uses best guess of eigenvecs
#                     V = X
#                     V, _ = np.linalg.qr(V)
#                     break
#                     
#                 else:
#                     V_copy = np.copy(V)
#                     sub_dim += 1
#                     V = np.eye(mat_dim,sub_dim) 
#                     V[:,:(sub_dim-1)] = V_copy
#                     # add new vector to end of subspace
#                     V[:,-1] = q/norm
# 
#     if converged:
#         return E, X 


# Path: mmd/postscf.py
import numpy as np
import scipy
import sys
from itertools import product, combinations
from bitstring import BitArray
from mmd.slater import common_index, get_excitation
from mmd.utils.davidson import davidson
from scipy.special import comb
from scipy.linalg import sqrtm, lu

                EMP2 += numer/denom
            self.mol.emp2 = EMP2 + self.mol.energy   

        print('E(MP2) = ', self.mol.emp2.real) 

    @staticmethod
    def tuple2bitstring(bit_tuple):
        ''' From tuple of occupied orbitals, return bitstring representation '''
        string = ['0']*(max(bit_tuple) + 1)
        for i in bit_tuple:
            string[i] = '1'
        string = ''.join(string[::-1])
        return BitArray(bin=string)


    def hamiltonian_matrix_element(self,det1,det2,Nint):
        """ return general Hamiltonian matrix element <det1|H|det2> """

        exc, degree, phase = get_excitation(det1,det2,Nint)

        if degree > 2:
            return 0

        elif degree == 2:
            # sign * <hole1,hole2||particle1,particle2>
            return phase * self.mol.double_bar[exc[1,0], exc[2,0], exc[1,1], exc[2,1]]

        elif degree == 1:
            m = exc[1,0]
            p = exc[1,1]

            common = common_index(det1,det2,Nint)