#! RHF cc-pVDZ energy for water, automatically scanning the symmetric stretch and bending coordinates
#! using Python's built-in loop mechanisms.  The geometry is apecified using a Z-matrix with variables
#! that are updated during the potential energy surface scan, and then the same procedure is performed
#! using polar coordinates, converted to Cartesian coordinates.
import math
import qcdb

refENuc = [ 9.785885838936569, 9.780670106434425, 8.807297255042920,     #TEST
            8.802603095790996, 8.006633868220828, 8.002366450719077 ]    #TEST

# Define the points on the potential energy surface using standard Python list functions
Rvals = [ 0.9, 1.0, 1.1 ]
Avals = range(102, 106, 2)


# Start with a potentital energy scan in Z-matrix coordinates

h2o = qcdb.Molecule("""
    O
    H 1 R
    H 1 R 2 A
""")

print "\n Testing Z-matrix coordinates\n"  #TEST

count = 0
for R in Rvals:
    h2o.set_variable('R', R)  # alternately, h2o.R = R
    for A in Avals:
        h2o.A = A  # alternately, h2o.set_variable('A', A)
        h2o.update_geometry()
        qcdb.compare_values(refENuc[count], h2o.nuclear_repulsion_energy(), 10, "Nuclear repulsion energy %d" % count)
        count = count + 1

# And now the same thing, using Python's trigonometry functions, and Cartesian input.  This time
# we want to reset the Cartesian positions every time the angles and bond lengths change, so we
# define the geometry inside the loops.  N.B. this requires the basis set to be re-specified after
# every change of geometry


print "\n Testing polar coordinates\n"  #TEST

count = 0
for R in Rvals:
    for A in Avals:
        h2o = qcdb.Molecule("""
            O   0.0    0.0    0.0
            H   0.0      R    0.0
            H   0.0  RCosA  RSinA
        """)
        # The non-numeric entries above just define placeholders with names.  They still need
        # to be set, which we do below.
        h2o.R = R
        h2o.set_variable('RCosA', R * math.cos(math.radians(A)))
        h2o.RSinA = R * math.sin(math.radians(A))
        h2o.update_geometry()

        qcdb.compare_values(refENuc[count], h2o.nuclear_repulsion_energy(),                 #TEST
                       10, "Nuclear repulsion energy %d" % count)                      #TEST
        count = count + 1
