import numpy as np
import pylab as pl
from gpaw.basis_data import Basis, BasisPlotter
from gpaw.atom.atompaw import AtomPAW

symbol = 'Au'

ordinary_basis = Basis(symbol, 'sharp.sz')
rcut = ordinary_basis.d * ordinary_basis.ng

#calc = AtomPAW(symbol, [[[2], [3]]], h=.02, rcut=rcut, setups='hgh',
#               basis='sz')

calc = AtomPAW(symbol, [[[1.], [0.], [10.]]], h=.05, rcut=2*rcut, 
               setups='hgh.sc',
               basis='sz')


#calc = AtomPAW(symbol, [[[1.]]], h=.05, rcut=2*rcut, 
#               setups='hgh',
#               basis='sz')


e = calc.get_potential_energy()

basis = calc.extract_basis_functions()
#for bf in basis.bf_j:
#    if bf.phit_g[1] < 0:
#        bf.phit_g *= -1.0

bp = BasisPlotter(premultiply=1)

#basis.write_xml()

plot = 1
if plot:
    bp.plot(basis)
    bp.plot(ordinary_basis)
    pl.show()



#epaw = []
#ehgh = []
#hvalues = np.arange(0.01, 0.3601, 0.01)
#print hvalues
#for h in hvalues:
#    calc = AtomPAW('N', [[[2], [3]]], h=h, rcut=rcut, setups='hgh',
#                   basis='sz')
#    e = calc.get_potential_energy()
#    ehgh.append(e)
#    calc = AtomPAW('N', [[[2], [3]]], h=h, rcut=rcut, setups='paw',
#                   basis='sz')
#    e = calc.get_potential_energy()
#    epaw.append(e)
#out = open('results.txt', 'w')
#print >> out, repr(epaw)
#print >> out, repr(ehgh)
#out.close()
#pl.plot(hvalues, epaw)
#pl.figure()
#pl.plot(hvalues, ehgh)
#pl.show()