from ase import Atoms
from gpaw import GPAW, Mixer, MixerSum, ConvergenceError, PoissonSolver, restart
from gpaw.poisson import PoissonSolver
from gpaw.utilities.timing import Timer
from gpaw.eigensolvers.rmm_diis import RMM_DIIS
ps = PoissonSolver(nn='M', relax='GS', eps=1e-9)
es = RMM_DIIS(keep_htpsit=False)
# from gpaw.hs_operators import MatrixOperator
# MatrixOperator.nblocks = 2

class MyGPAW(GPAW):
    timer_class = Timer
    
# Makes Python traceback appear only on the root processor
import sys
from gpaw.mpi import world
from ase.utils import devnull
# if world.rank != 0:
#     sys.stderr = devnull
    
a = 4.08
bulk = Atoms('Au4', positions=((0, 0 ,0),
                               (0.5, 0.5, 0),
                               (0.5, 0, 0.5),
                               (0, 0.5, 0.5)), pbc=True)

bulk.set_cell((a,a,a), scale_atoms=True)
bulk = bulk.repeat((3, 3, 3))
calc = MyGPAW(gpts=(120, 120, 120),
            maxiter=5,
            mode = 'fd',
            poissonsolver = ps,
            nbands = 600,
            spinpol = False,
            xc = 'LDA',
            width = 0.01,
            mixer = Mixer(0.10, 3), 
            # mixer = MixerSum(0.1, 3), # spinpol = True
            eigensolver = es,
            parallel = {'buffer_size':2048},
            txt = 'Au_bulk.out')

bulk.set_calculator(calc)
filename = 'Au_bulk.hdf5'
# spinpol = True
# magmom = 2.0
# mms = [magmom for i in range(len(bulk))]
# bulk.set_initial_magnetic_memonts(mms)

# Run intial 5 SCF steps
try:
    bulk.get_potential_energy()
except ConvergenceError:
    pass

calc.write(filename, mode = 'all')
# calc.write(filename)
# For a Real Restart with reading wavefunctions
del bulk, calc
bulk, calc = restart(filename, maxiter=5, txt = 'Au_bulk2.out')
# calc.initialize_positions()
# calc.scf.reset()
# calc.set(convergence=dict(eigenstates=5e-8,
#                           density=5e-4)) # change criteria from default
# Re-start from previous wfs and run 5 more SCF steps
try:
    bulk.get_potential_energy()
except ConvergenceError:
    pass