# Import the KohnSham module from ATK
from ATK.KohnSham import *

# Function for handling molecule setup   
def waterConfiguration(angle,bondLength):
    from math import sin,cos

    theta = angle.inUnitsOf(radians)
    pos = [(0.0, 0.0, 0.0)*Angstrom,
           (1.0, 0.0, 0.0)*bondLength,
           (cos(theta), sin(theta), 0.0)*bondLength]
    elm = [Oxygen] + [Hydrogen]*2

    return MoleculeConfiguration(elm,pos)


# Function for calculating H-O-H bond angle
def waterBendingAngle(water):
    from math import acos

    coords = water.cartesianCoordinates()
    oh1 = coords[1] - coords[0]
    oh2 = coords[2] - coords[0]
    prod = dot(oh1,oh2)/( length(oh1)*length(oh2) )

    return acos(prod)*radians


# Function for calculating O-H bond length
def waterBondLength(water):
    coords = water.cartesianCoordinates()
    oh1 = coords[1] - coords[0]

    return length(oh1)


# Dot product between two vectors
def dot(u,v):
    return (u*v).sum()
    

# Length of a vector containing elements of type PhysicalQuantity
def length(u):
    from math import sqrt

    return sqrt(((u*u).sum()).inUnitsOf(Angstrom**2))*Angstrom


# Setup water molecule
h2o = waterConfiguration(100.0*degrees, 1.1*Angstrom)

# Set DFT method with a low tolerance
low_tolerance = iterationControlParameters( tolerance = 1.0e-6 )
dft_method = KohnShamMethod(iteration_control_parameters = low_tolerance)

# Calculate the optimized geometry for H2O molecule.
opt_params = geometricOptimizationParameters(
    force_tolerance = 0.005*eV/Ang
    )
h2o_opt = calculateOptimizedAtomicGeometry(
    atomic_configuration = h2o,
    optimization_parameters=opt_params,
    method = dft_method
    )

# Calculate optimized bond length and bending angle
bondlength = waterBondLength(h2o_opt)
bendingangle = waterBendingAngle(h2o_opt)

# Print the results
print 'H-O-H angle    = ', bendingangle.inUnitsOf(degrees)
print 'O-H bondlength = ', bondlength