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='J', eps=1e-10, nn='M'),
                   convergence=dict(eigenstates=5e-8,
                                    density=5e-4),
                   xc='RPBE',
                   txt='gpaw.txt',
                   maxiter=240,
                   mixer=Mixer(0.03, 7, 60.0),
                   parallel=dict(sl_default=(8, 8, 32),
                                 buffer_size=2048,
                                 band=8,
                                 domain=(8,16,4)),
                   occupations=FermiDirac(0.01, maxiter=4000))
    kwargs1.update(**kwargs)
    return kwargs1
            
calc = GPAW(**getkwargs())

if world.rank == 0:
    import os
    if os.path.exists('TERMINATE'):
        os.remove('TERMINATE')
def checkconv():
    decision = 0
    if world.rank == 0:
        if os.path.exists('TERMINATE'):
            decision = 1
    decision = world.sum(decision)
    if decision:
        calc.scf.converged = True
calc.attach(checkconv, 1)

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

if world.rank == 0:
    fd = open('energy.txt', 'w')
    print >> fd, 'e =', e
    fd.close()
    timings = open('timings.txt', 'w')
    calc.timer.write(timings)
    timings.close()


import pickle
import numpy as np

def savedensity(calc):
    nt_G = calc.density.nt_sG[0] # XX
    gd = calc.wfs.gd
    rank = gd.rank

    if calc.wfs.world.rank == 0:
        f1 = open('parsize.dat', 'w')
        #f2 = open('globalshape.dat', 'w')
        print >> f1, '%d %d %d' % tuple(gd.parsize_c)
        #print >> f2, '%d %d %d' % tuple(gd.N_c)
        f1.close()
        #f2.close()

    if calc.wfs.bd.rank == 0:
        # XXX nspins
        densityfile = open('density.rank%d.spin%d.dat' % (gd.rank, 0), 'w')
        np.save(densityfile, nt_G)
        densityfile.close()


savedensity(calc)