from MMTK import *
from MMTK.ForceFields import Amber99ForceField
from MMTK.Minimization import ConjugateGradientMinimizer
from MMTK.Trajectory import StandardLogOutput
from MMTK.NormalModes import VibrationalModes

universe = InfiniteUniverse(Amber99ForceField())
molecule = Molecule('water')
universe.addObject(molecule)

minimizer = ConjugateGradientMinimizer(universe,
                                       actions=[StandardLogOutput(1)])
minimizer(convergence = 1.e-4, steps = 100)

modes = VibrationalModes(universe)

for mode in modes:
    print mode

d_O_H1 = universe.distance(molecule.O, molecule.H1)
d_O_H2 = universe.distance(molecule.O, molecule.H2)
a_H1_O_H2 = universe.angle(molecule.H1, molecule.O, molecule.H2)

minimum_configuration = copy(universe.configuration())
displaced_configuration = minimum_configuration + modes[8]

universe.setConfiguration(displaced_configuration)
displaced_d_O_H1 = universe.distance(molecule.O, molecule.H1)
displaced_d_O_H2 = universe.distance(molecule.O, molecule.H2)
displaced_a_H1_O_H2 = universe.angle(molecule.H1, molecule.O,
                                     molecule.H2)
universe.setConfiguration(minimum_configuration)

print
print "Change in distance O-H1:", displaced_d_O_H1 - d_O_H1
print "Change in distance O-H2:", displaced_d_O_H2 - d_O_H2
print "Change in angle H1-O-H2:", displaced_a_H1_O_H2 - a_H1_O_H2