from ATK.KohnSham import *

def createOxygenMolecule( atom_separation = 1.23):

    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( spin = 0.0):

    total_electron_density = electronDensityParameters(
        fixed_spin = spin*hbar
        )

    iteration_control_parameters = iterationControlParameters(tolerance = 1e-6)
    
    dft_method=KohnShamMethod(
        basis_set_parameters=basisSetParameters(),
        exchange_correlation_type=GGA.PBE,
        electron_density_parameters = total_electron_density,
	iteration_control_parameters = iteration_control_parameters
        )

    return dft_method

def determineBondLength(oxygen):
    import math
    coord = oxygen.cartesianCoordinates()
    oo_bond_length = coord[1][2] - coord[0][2]
    return math.fabs(oo_bond_length.inUnitsOf(Angstrom))

from numpy import arange
energy_dict ={}
geometry_opt_param = geometricOptimizationParameters()

print 'Separation (Angstrom)\tTotal Energy (eV)\tSpin (hbar)' 
for polarization in  arange(0.0,2.1,0.1):
#for polarization in  arange(0.0,0.1,0.1):
    dft_method = createKohnShamMethod(polarization)

    opt_oxygen = calculateOptimizedAtomicGeometry(
        atomic_configuration=createOxygenMolecule(),
        method=dft_method,
        optimization_parameters=geometry_opt_param
        )
    
    distance = determineBondLength(opt_oxygen)
    scf_oxygen = dft_method.apply(createOxygenMolecule(distance))
    total_energy = calculateTotalEnergy(scf_oxygen).inUnitsOf(eV)
    print "%8.4f %s %8.4f %s %2.1f" %(distance,'\t\t',
                                      calculateTotalEnergy(scf_oxygen).inUnitsOf(eV),\
                                      '\t\t',polarization)
    energy_dict[total_energy] = (distance,polarization)

energy_list = energy_dict.keys()
energy_list.sort()
print "Minimum total energy: %8.4f eV" % (energy_list[0])
print "O-O equilibrium distance: %8.4f A" % (energy_dict[energy_list[0]][0])