from ATK.KohnSham import *

def createOxygenMolecule(atom_separation):
    '''
    @param atom_seperation - bond-length for
           a oxygen molecule
    @return MoleculeConfiguration object for an Oxygen
            molecule.
    '''
    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):
    '''
    @param spin - Total spin for the electron density.
           Supplied with no unit.
    '''
    fixed_electron_density = electronDensityParameters(
        fixed_spin = spin*hbar
        )
    
    dft_method=KohnShamMethod(
        basis_set_parameters=basisSetParameters(),
        exchange_correlation_type=GGA.PBE,
        electron_density_parameters = fixed_electron_density
        )
    
    return dft_method

# Create a list of corresponding bond-lengths and
# values of total spin that we want to compare.
spin_types = [
    ('singlet',0.0,1.22),
    ('triplet',1.0,1.22),
    ('quintet',2.0,1.93)] 

vnl_file = VNLFile('oxygen.vnl')
for (spinType,spinValue,ooDistance) in spin_types:
    # Create spin-polarized KohnShamMethod
    dft_oxygen = createKohnShamMethod(spinValue)
    # Create oxygen molecule
    oxygen_molecule = createOxygenMolecule(ooDistance) 
    # Perform actual calculation
    scf_oxygen = dft_oxygen.apply(oxygen_molecule)
    print 'Spin type:',spinType,
    print "Total energy: %8.4f eV" %(calculateTotalEnergy(scf_oxygen).inUnitsOf(Units.eV)) 
    # Calculate eigenstates and export to VNL file.
    # ID reveals the value of the fixed spin.
    homoLumoEigenstates = calculateEigenstates(
        self_consistent_calculation = scf_oxygen, 
        quantum_numbers=([6,7], Spin.Up)
        )
    vnl_file.addToSample(oxygen_molecule, 'oxygen'+spinType)
    for eigenstate in homoLumoEigenstates:
        vnl_file.addToSample(eigenstate, 'oxygen'+spinType)