import Numeric as na
from ASE.Utilities.BEE import GetEnsembleEnergies
from ASE import Atom, ListOfAtoms
from Dacapo import Dacapo

a = 4.0  # Size of unit cell (Angstrom)


# Hydrogen atom:
atom = ListOfAtoms([Atom('H', (0, 0, 0), magmom=1)],
                   cell=(a, a, a), periodic=True)

# Hydrogen molecule:
d = 0.74  # Experimental bond length
molecule = ListOfAtoms([Atom('H', (0, 0, 0)),
                        Atom('H', (d, 0, 0))],
                       cell=(a, a, a), periodic=True)

for xc in ['PBE','RPBE']:
    print xc + ':'
    calc = Dacapo(planewavecutoff=400, spinpol=True,nbands=1, xc=xc, out='H1.%s.out' % xc)
    calc.SetExecutable('/home/niflheim/lhansen/svol/dacapo/src/pglinux_serial/dacapo.run')
    atom.SetCalculator(calc)

    e1 = atom.GetPotentialEnergy()
    c1 = calc.GetEnsembleCoefficients()

    calc = Dacapo(planewavecutoff=400, nbands=1, xc=xc, out='H2.%s.out' % xc)
    calc.SetExecutable('/home/niflheim/lhansen/svol/dacapo/src/pglinux_serial/dacapo.run')
    molecule.SetCalculator(calc)
    
    e2 = molecule.GetPotentialEnergy()
    c2 = calc.GetEnsembleCoefficients()

    print 'hydrogen atom energy:     %5.2f eV' % e1
    print 'hydrogen molecule energy: %5.2f eV' % e2
    print 'atomization energy:       %5.2f eV' % (2 * e1 - e2)
    print c1
    print c2
    print c2 - 2 * c1
    
    e1i = GetEnsembleEnergies(c1)
    e2i = GetEnsembleEnergies(c2)
    eai = 2 * e1i - e2i

    n = len(eai)
    ea0 = na.sum(eai) / n
    sigma = (na.sum((eai - ea0)**2) / n)**0.5
    print 'Best fit:', ea0, '+-', sigma, 'eV'