import n2v
import psi4

H2O = psi4.geometry(
"""
0 1
O   0.000000   0.000000   0.000000
H   0.757459   0.586790   0.000000
H  -0.757459   0.586790   0.000000
noreorient
nocom
units bohr
symmetry c1
""" )


#n2v is driven by psi4's reference option. Make sure you set it accordingly.
psi4.set_options({"reference" : "rhf"})

#Perform a calculation for a target density.
#Remember that for post scf calculations, Psi4 does not update the density.
#Thus make sure you obtain something like a dipole in order to do so.
e, wfn = psi4.properties("ccsd/cc-pvtz", return_wfn=True, properties=["dipole"], molecule=H2O)

#Define inverter objects for each molcule. Simply use the wnf object from psi4 as an argument.
ine = n2v.Inverter()
ine.set_system(H2O, "cc-pvtz",wfn=wfn)
ine.from_wfn(wfn)


# how to change the increase lambda
start = 1
stop = 1000
step = 2
lam_list = []

lam = start
for i in range(int(start), int(stop), int(step)):
    lam_list.append(i)


# do zmp 
ine.invert("zmp", guide_components='fermi_amaldi', opt_max_iter=2000, opt_tol=1e-7, zmp_mixing=0, print_scf=False,
        lambda_list=lam_list)

print(ine)