#import psi4
#import numpy as np


molecule wat {
symmetry c1
no_reorient
O      -1.769142     -0.076181     -0.000000
H      -2.065745      0.837492      0.000000
H      -0.809034      0.001317     -0.000000
--
O       1.283899      0.088882      0.000000
H       1.580006     -0.425760      0.756361
H       1.580006     -0.425760     -0.756361
}

set basis 3-21G
#set print 3
method='scf'
set guess sad
set maxiter 0
set fail_on_maxiter false
set scf_type df
#set S_ORTHOGONALIZATION  canonical

wat.fix_com(True)
wat.fix_orientation(True)

print "== monomer A =="
molA=wat.extract_subsets(1)
A_e_scf, Awfn = energy(method, molecule=molA, return_wfn=True)
print A_e_scf

A_occ = np.array(Awfn.Ca_subset("SO", "OCC").to_array())
A_virt = np.array(Awfn.Ca_subset("SO", "VIR").to_array())
Anocc=A_occ.shape[1]
Amo=A_occ.shape[0]
Anvirt=A_virt.shape[1]
#print Anocc, Anvirt,Amo
Aeigen=np.array( Awfn.epsilon_a().to_array() )


print "== monomer B =="

molB=wat.extract_subsets(2)

B_e_scf, Bwfn = energy(method, molecule=molB, return_wfn=True)
print B_e_scf

B_occ = np.array(Bwfn.Ca_subset("SO", "OCC").to_array())
B_virt = np.array(Bwfn.Ca_subset("SO", "VIR").to_array())
Bnocc=B_occ.shape[1]
Bnvirt=B_virt.shape[1]
Bmo=B_occ.shape[0]
Beigen=np.array( Bwfn.epsilon_a().to_array() )

print "= promo wfn =="
clean()
ABnmo=Anocc+Bnocc+Anvirt+Bnvirt
AB=np.zeros((ABnmo,ABnmo))
ABeigen=np.zeros(ABnmo)
ABwfn = core.Wavefunction.build(wat, core.get_global_option('BASIS'))


# merge MOs and eigenvalues
cnt=0
for i in range(Anocc):
   ABeigen[cnt]=Aeigen[i]
   for j in range(Amo):
      AB[j][cnt]=A_occ[j][i]
   cnt+=1

for i in range(Bnocc):
   ABeigen[cnt]=Beigen[i]
   for j in range(Bmo):
      AB[j+Amo][cnt]=B_occ[j][i]
   cnt+=1

for i in range(Anvirt):
   ABeigen[cnt]=Aeigen[Anocc+i]
   for j in range(Amo):
      AB[j][cnt]=A_virt[j][i]
   cnt+=1

for i in range(Bnvirt):
   ABeigen[cnt]=Beigen[Bnocc+i]
   for j in range(Bmo):
      AB[j+Amo][cnt]=B_virt[j][i]
   cnt+=1



# build new wfn object and set MOs as guess
AB_Ca=core.Matrix.from_array(AB)

AB_Ca.print_out()

ref_wfn = scf_wavefunction_factory("scf", ABwfn, core.get_option('SCF', 'REFERENCE'))
core.set_legacy_wavefunction(ref_wfn)
ref_wfn.guess_Ca(AB_Ca) # this defines also the occupation :(
ABnocc=[Anocc+Bnocc]
ABoccpi=core.Dimension.from_list(ABnocc)
ref_wfn.force_doccpi(Aoccpi+Boccpi) # not found
ref_wfn.reset_occ(True) # seems to be overwritten later hf.initialize()
molden(ref_wfn,filename='api.molden')

print " == numpy SCF  =="
# MANUAL HF

mints = psi4.core.MintsHelper(ref_wfn.basisset())
S = np.asarray(mints.ao_overlap())

# Get nbf and ndocc for closed shell molecules
nbf = S.shape[0]
ndocc = ref_wfn.nalpha()

print('Number of occupied orbitals: %d' % ndocc)
print('Number of basis functions: %d' % nbf)

# Compute required quantities for SCF
V = np.asarray(mints.ao_potential())
T = np.asarray(mints.ao_kinetic())
I = np.asarray(mints.ao_eri())

Enuc = wat.nuclear_repulsion_energy()


# Build H_core
H = T + V

# Orthogonalizer A = S^(-1/2) using Psi4's matrix power.
#A = mints.ao_overlap()
#A.power(-0.5, 1.e-16)
#A = np.asarray(A)


# non-orthonormalized fragment MO density and its energy
C = AB
Cocc = C[:, :ndocc]
D = np.einsum('pi,qi->pq', Cocc, Cocc)

# Build fock matrix
J = np.einsum('pqrs,rs->pq', I, D)
K = np.einsum('prqs,rs->pq', I, D)
F = H + J * 2 - K

# SCF energy and update
SCF_E = np.einsum('pq,pq->', F + H, D) + Enuc
print SCF_E


#==============
print " == psi4 SCF  =="
AB_e_scf = ref_wfn.compute_energy()
print "ESCF(initial) =", AB_e_scf