# Perform a geomtry optimization
from ase import Atoms, PickleTrajectory, read
from ase.optimize.lbfgs import LBFGS, LBFGSLineSearch
from gpaw import GPAW, Mixer, MixerSum, ConvergenceError, PoissonSolver
from gpaw.poisson import PoissonSolver
from gpaw.eigensolvers.rmm_diis import RMM_DIIS
ps = PoissonSolver(nn='M', relax='GS', eps=1e-9)
es = RMM_DIIS(keep_htpsit=False)

# 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

epsilon = 0.01 # small perturbation

a = 4.08
bulk = Atoms('Au4', positions=((0, 0 ,0),
                               (0.5, 0.5, 0),
                               (0.5, epsilon, 0.5),
                               (0, 0.5, 0.5)), pbc=True)

bulk.set_cell((a,a,a), scale_atoms=True)
bulk = bulk.repeat((1, 1, 1))
calc = GPAW(h=0.16,
            maxiter=200,
            mode = 'fd',
            poissonsolver = ps,
            nbands = 30,
            spinpol = False,
            xc = 'LDA',
            width = 0.01,
            mixer = Mixer(0.10, 5, 100.0), 
            # mixer = MixerSum(0.10, 5, 100.0), # spinpol = True
            eigensolver = es,
            parallel = {'buffer_size':2048},
            txt = 'Au_bulk.out')

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



mode = 'w'
traj = PickleTrajectory('qn.traj', mode, bulk)
relax = LBFGS(bulk, logfile = 'qn.log', trajectory = traj, restart= 'qn.pckl')

# Check if there is enough time left, if not shutdown the calculation
# Need to do this for the SCF and structural optimization seperatel
wallclocktime = 2 # wallclock time limit in minutes
from time import time
class MaxTime:
    def __init__(self, nminutes):
        self.starttime = time()
        self.maxtime = nminutes * 60
        
    def checkSCF(self):
        decision = 0
        if world.rank == 0:
            now = time()
            if now - self.starttime > self.maxtime:
                decision = 1
        decision = world.sum(decision)
        if decision:
            calc.scf.converged = True
            
    def checkOpt(self):
        decision = 0
        if world.rank == 0:
            now = time()
            if now - self.starttime > self.maxtime:
                decision = 1
        decision = world.sum(decision)
        if decision:
           relax.fmax = float('inf')

StopCalculation = MaxTime(wallclocktime)
calc.attach(StopCalculation.checkSCF)
relax.attach(StopCalculation.checkOpt)
relax.attach(lambda: calc.write('Au_bulk.%d.hdf5' % relax.nsteps))

relax.run(fmax=0.05, steps=5) # This should stop in the middle of the 4 relaxation step