import psi4

output1 = "water_manual.out"

psi4.core.set_output_file(output1, False)
psi4.set_memory('500 MB')

h2o = psi4.geometry("""
  0 1
  O
  H 1 0.96
  H 1 1.10 2 104.5
  """)
dft_method = 'b3lyp'
geom_iter = 10
psi4.set_options({'reference':'rhf','basis':'dz'})
n = 1
while n <= geom_iter:
    psi4.core.print_out('-----Geometry optimization cycle: ')
    psi4.core.print_out(str(n) + '-----\n')
    h2o.update_geometry()
    psi4.energy(dft_method)
    psi4.core.print_out ("------------energy end-----------------")
    G, wfn = psi4.driver.gradient(dft_method, return_wfn=True, molecule=h2o)
    psi4.core.print_out ("------------gradiend end---------------")
    psi4.core.set_legacy_molecule(h2o)
    optimization_status = psi4.core.optking()
    h2o = psi4.core.get_legacy_molecule()
    h2o.update_geometry()
    print ("optimization status: ", optimization_status)
    if optimization_status is True:
        psi4.core.print_out ('optimization done')
        psi4.core.print_out('Exiting')
        print ('\n all done, exiting. \n')
        break
    n += 1