import numpy as np

#evaluate nuclear repulsion energy using the built-in function of PSI4
def rep(molecule):
	charge = molecule.molecular_charge()
	mult = molecule.multiplicity()
	Z = 0
	for A in range(molecule.natom()):
		Z += molecule.Z(A)
	ndocc = int(Z / 2) - (charge / 2) # number of doubly-occupied orbitals
	Enuc = molecule.nuclear_repulsion_energy()
	return Enuc

#atomic coordinate of water
molecule mol {
O
H 1 1.1
H 1 1.1 2 105
}

#The basis set I used (the STO-3G basis set)
set {
 basis = STO-3G
}

#Evaluation of the kinetic energy term, potential energy term, the two-electron integral and the overlap integral
mints = MintsHelper()
T = mints.ao_kinetic()
V = mints.ao_potential()
I = mints.ao_eri()
S = mints.ao_overlap()

nbf = S.rows(0)


T = np.array(T)
V = np.array(V)
S = np.array(S)
np.array(I)
H = np.zeros((nbf,nbf))

#the one-electron integral, H, is basically the sum of kinetic energy term and potential energy term

for i in range(nbf):
	for j in range (nbf):
		H[i][j] = T[i][j]+V[i][j]

I = np.array(I).reshape(nbf, nbf, nbf, nbf) 

#save all these values to corresponding files
np.savetxt('overlap', S)
np.savetxt('kinetic', T)
np.savetxt('potential', V)
np.savetxt('one-electron', H)

text_file1 = open('two-electron','w')
for i in range(nbf):
	for j in range(nbf):
		for k in range(nbf):
			for p in range(nbf):
				ii = str(I[i][j][k][p])
				text_file1.write(str(i))
				text_file1.write('\t')
				text_file1.write(str(j))
				text_file1.write('\t')
				text_file1.write(str(k))
				text_file1.write('\t')
				text_file1.write(str(p))
				text_file1.write('\t')
				text_file1.write(ii)
				text_file1.write('\n')
text_file1.close()

Enuc = rep(mol)
text_file1 = open('vnn','w')
text_file1.write(str(Enuc))
text_file1.close()