import numpy as np
import time

#from phase_mod import adadaa

m1 = 0x55555555
m2 = 0x33333333
m4 = 0x0f0f0f0f
m8 = 0x00ff00ff
m16 = 0x0000ffff

def bit_count(x):
    x = (x & m1) + ((x >> 1) & m1)
    x = (x & m2) + ((x >> 2) & m2)
    x = (x & m4) + ((x >> 4) & m4)
    x = (x & m8) + ((x >> 8) & m8)
    x = (x & m16) + ((x >> 16) & m16)
    return x
    

ORBITAL_DIR = "space/"
HAMILTONIAN_DIR = "hamiltonian/"
BASIS_DIR = "basis/"
INTERACTION_DIR = "interaction/"

#def hamiltonian(basisfile, onebodyfile, twobodyfile, orbitals):

#print()
#print("hamiltonian.py: Calculating the Hamiltonian matrix with the following input:")
#print("\torbitals = ", ORBITAL_DIR + orbitals)
#print("\tbasis = ", BASIS_DIR + basisfile)
#print("\tonebody = ", INTERACTION_DIR + oneparticle)
#print("\ttwobody = ", INTERACTION_DIR + twoparticle)
#print()

basisfile = "20o_basis_bit.npy"
onebodyfile = "sd_sp_np.int"
twobodyfile = "usdb_m.int"

# Basis of Slater determinants
basis = np.load(BASIS_DIR + basisfile)
#basis2 = np.loadtxt(BASIS_DIR + "basis.txt")
n_basis_states = np.shape(basis)[0]
n_particles = bit_count(basis[0])
obme = np.loadtxt(INTERACTION_DIR + onebodyfile)
n_obme = np.shape(obme)[0]
tbme = np.loadtxt(INTERACTION_DIR + twobodyfile)
n_tbme = np.shape(tbme)[0]

anni_ob = np.array(2**(obme[:,0] - 1), dtype=np.int32)
anni_tb = np.array(2**(tbme[:,2] - 1) + 2**(tbme[:,3] - 1), dtype=np.int32)
crea_tb = np.array(2**(tbme[:,0] - 1) + 2**(tbme[:,1] - 1), dtype=np.int32)
anni1_tb = np.array(2**(tbme[:,3] - 1), dtype=np.int32)
anni2_tb = np.array(2**(tbme[:,2] - 1), dtype=np.int32)

crea1_tb = np.array(2**(tbme[:,1] - 1), dtype=np.int32)
crea2_tb = np.array(2**(tbme[:,0] - 1), dtype=np.int32)

START = time.time()

mass_factor = (18./(16.+n_particles))**(0.3)

H = np.zeros((n_basis_states, n_basis_states))

###############################################################################
# Calculate the contribution of two-body matrix elements to the Hamiltonian
###############################################################################

# Loop over basis states
for i in range(n_basis_states):
#    if i == 5:
#        break
    bb = basis[i]
    # Loop over available two-body matrix elements
    for m in np.arange(n_tbme):
        b = bb
        phase = -1
        # Check if the action of the two annihilation operators in the matrix
        # element is valid
        if (bin(b & anni1_tb[m]).count('1') == 1) and (bin(b & anni2_tb[m]).count('1') == 1): # This one is faster than the version above!
#            print(bin(anni1_tb[m]), "\t", bin(anni2_tb[m]))
            # Act with the annihilation operators on the basis state
            bb = b
            b -= anni1_tb[m]
            phase = phase * (-1)**(bin(bb & ((1 << anni1_tb[m]) -1) ).count('1'))
            bb = b
            b -= anni2_tb[m]
            phase = phase * (-1)**(bin(bb & ((1 << anni2_tb[m]) -1) ).count('1'))

            # Check if the action of the two creation operators in the matrix
            # element is valid
            if (not b & crea1_tb[m]) and (not b & crea2_tb[m]):
                # Act with the creation operators on the basis state
#                print("crea")
                bb = b
                b += crea1_tb[m]
                phase = phase * (-1)**(bin(bb & ((1 << crea1_tb[m]) -1) ).count('1'))
                bb = b
                b += crea2_tb[m]
                phase = phase * (-1)**(bin(bb & ((1 << crea2_tb[m]) -1) ).count('1'))
                # Find the basis state that was generated by the action of the
                # creation and annihilation operators, and add the TBME to the
                # Hamiltonian
                for j in range(n_basis_states):
                    if b == basis[j]:
                        H[i][j] += phase*tbme[m][4]*mass_factor
#                        print(bin(anni_tb[m]), "\t", bin(crea_tb[m]), "\t", bin(anni_tb[m] - crea_tb[m]))
#                        print(phase*tbme[m][4]*mass_factor, "[", i, "][", j, "], phi = ", phase, ", ", bin(b), " ", bin(bb))
                        break
                    
# Loop over basis states
#for i in range(n_basis_states):
#    if i == 20:
#        break
#    bb = basis[i]
#    # Loop over available two-body matrix elements
#    for m in np.arange(n_tbme):
#        b = bb
#        # Check if the action of the two annihilation operators in the matrix
#        # element is valid
##        if bit_count(b & anni_tb[m]) == 2: # Using self-implemented bit counting. Slower than the code below
#        if bin(b & anni_tb[m]).count('1') == 2: # This one is faster than the version above!
#            # Act with the annihilation operators on the basis state
#            b -= anni_tb[m]
#            # Check if the action of the two creation operators in the matrix
#            # element is valid
#            if not b & crea_tb[m]:
#                # Act with the creation operators on the basis state
#                b += crea_tb[m]
#                # Find the basis state that was generated by the action of the
#                # creation and annihilation operators, and add the TBME to the
#                # Hamiltonian
#                for j in range(n_basis_states):
#                    if b == basis[j]:
#                        phase = (-1)**(bin(b & bb).count('1'))
##                        phase = (-1)**(bin((b & bb)).count('1'))
#                        H[i][j] += phase*tbme[m][4]*mass_factor
##                        print(bin(bb), "\t", bin(b), "\t", bin(bb & b))
#                        print(tbme[m])
#                        print(phase*tbme[m][4]*mass_factor, "[", i, "][", j, "], phi = ", phase, ", ", bin(b), " ", bin(bb))
#                        print()
#                        break

###############################################################################
# Calculate the contribution of one-body matrix elements to the Hamiltonian
# In the harmonic oscillator basis, the one-body term is orthogonal, so no
# action of annihilation and creation operators is required
###############################################################################
                    
for i in range(n_basis_states):
    for j in range(n_obme):
#        print(bin(2**j),"\t", bin(basis[i]))
        if 2**j & basis[i]:
            H[i][i] += obme[j][4]

STOP = time.time()

np.savetxt(HAMILTONIAN_DIR + "Hamiltonian_bit.txt", H, delimiter = " ")

print("Execution took", STOP - START, "seconds")
print(H)

print()
print("hamiltonian_bit.py: Saved Hamiltonian matrix to '", HAMILTONIAN_DIR, "Hamiltonian_bit.txt'")
print()