from ATK.KohnSham import *

def createOxygenMolecule(atom_separation):
    '''
    Simple function for creating a custom
    MoleculeConfiguration
    @param atom_separation - bond-length between
           oxygen atoms
    '''
    atoms = 2*[Oxygen]
    coordinates =  [
        ( 0.0, 0.0, atom_separation/2) ,
        ( 0.0, 0.0,-atom_separation/2)
        ]*Ang
    oxygen_molecule = MoleculeConfiguration(
        elements = atoms,
        cartesian_coordinates = coordinates
        )
    return oxygen_molecule

def createKohnShamMethod():
    '''
    @return Predefined KohnShamMethod
    '''
    dft_method=KohnShamMethod( 
        basis_set_parameters=basisSetParameters(),
        exchange_correlation_type=GGA.PBE,
        )
    return dft_method

# Create a dictionary for storing bond lengths and
# associated bond lengths
from numpy import arange
energy_dict ={}

print 'Separation (Angstrom)\tTotal Energy (eV)' 
# Iterate over a list of increasing bond lengths
for distance in  arange(1.10,1.30,0.02):
    dft_method = createKohnShamMethod()
    scf_oxygen = dft_method.apply(
        createOxygenMolecule(distance)
        )
    total_energy = calculateTotalEnergy(scf_oxygen).inUnitsOf(eV)
    print "%4.2f %s %9.4f"%(distance,'\t\t\t',calculateTotalEnergy(scf_oxygen).inUnitsOf(eV))
    energy_dict[total_energy] = distance

# Find the bond length with the minimum energy
energy_list = energy_dict.keys()
energy_list.sort()
print "Minimum total energy: %9.4f eV" % (energy_list[0])
print 'O-O equilibrium distance:',energy_dict[energy_list[0]],'A'