from ase.io import read
from gpaw import GPAW, FermiDirac
from gpaw.mixer import Mixer
from glob import glob
from gpaw.utilities import h2gpts
from gpaw.poisson import PoissonSolver
from gpaw.hs_operators import MatrixOperator
from gpaw.mpi import world

MatrixOperator.nblocks = 2

system = read('geom.traj')


def adjust_to_h(system, h, idiv=8):
    if isinstance(h, float):
        h = (h, h, h)
    cell = system.get_cell().diagonal()
    gpts = (cell / h).astype(int) / idiv * idiv + idiv
    cell2 = gpts * h
    system.set_cell(cell2)
    system.center()
    return gpts
gpts = adjust_to_h(system, 0.14)


def getkwargs(**kwargs):
    kwargs1 = dict(nbands=((len(system) * 6) // 8) * 8,
                   gpts=gpts,
                   #gpts=h2gpts(0.14, system.get_cell(), idiv=8),
                   poissonsolver=PoissonSolver(relax='GS', eps=1e-10, nn='M'),
                   convergence=dict(eigenstates=5e-8),
                   xc='RPBE',
                   txt='gpaw.txt',
                   maxiter=240,
                   mixer=Mixer(0.04, 7, 60.0),
                   parallel=dict(sl_default=(8, 8, 32),
                                 band=8,
                                 domain=(8,16,8)),
                   occupations=FermiDirac(0.01, maxiter=4000))
    kwargs1.update(**kwargs)
    return kwargs1
            
calc = GPAW(**getkwargs())

system.set_calculator(calc)
e = system.get_potential_energy()

if world.rank == 0:
    fd = open('energy.txt', 'w')
    print >> fd, 'e =', e