#PBS -l nodes=1:ppn=8
#PBS -l walltime=0:25:0
#PBS -o emt.stdout.log
#PBS -e emt.stdout.err
#PBS -N EMT

import sys
import os
import gzip

from ase.io import read
from ase.lattice.cubic import FaceCenteredCubic
from ase import Atoms
from ase.optimize import QuasiNewton
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from ase.md.verlet import VelocityVerlet
from ase.units import kB, fs
from ase.io import write
from ase.io.trajectory import PickleTrajectory
from ase.md.langevin import Langevin
from ase.parallel import rank, barrier

from gpaw.mpi import world
from gpaw.utilities import devnull, h2gpts
from gpaw import GPAW, FermiDirac, Mixer, PoissonSolver

import numpy as np

fileid = os.path.split(__file__)[1].rsplit('.', 1)[0]

system = read('cluster.anneal.Au.traj')
system.pbc = 0
spos = system.get_scaled_positions()
maxspos = spos.max()
minspos = spos.min()
if maxspos > 0.9 or minspos < 0.1: # yuck, cluster divided across cell boundary
    for direction in range(3):
        if (spos[:, direction] > 0.9).any():
            spos[:, direction] = (spos[:, direction] + 0.5) % 1.0
    system.set_scaled_positions(spos)
    spos = system.get_scaled_positions()
    assert (0.1 < spos).all() and (spos < 0.9).all()

system.center(vacuum=7.5)
v = system.get_volume()
a0 = v**(1.0 / 3.0)
system.set_cell((a0, a0, a0))
system.center()

#if world.rank == 0:
#    txt = gzip.open('gpaw.%s.txt.gz' % fileid, 'w')
#else:
#    from gpaw.utilities import devnull
#    txt = devnull

txt = None
mixersetting = 0.1
if len(system) == 6 or len(system) % 15 == 0:
    txt = 'gpaw.%s.txt' % fileid
if len(system) < 10:
    mixersetting = 0.05

kwargs = dict(txt=txt,
              xc='RPBE',
              convergence=dict(density=1e-3),
              gpts=h2gpts(0.24, system.get_cell(), idiv=8),
              poissonsolver=PoissonSolver(relax='GS', eps=1e-7),
              occupations=FermiDirac(0.01),
              mixer=Mixer(mixersetting, 5, 50.0),
              mode='lcao',
              parallel=dict(sl_default=(world.size // 2, 2, 64)),
              basis='dzp')
if len(system) < 64:
    kwargs.pop('parallel')

calc = GPAW(**kwargs)
system.set_calculator(calc)
#system.get_forces()

class CheapPickleTrajectory(PickleTrajectory):
    write_energy = True
    write_forces = False
    write_stress = False
    write_magmoms = False
    write_momenta = False
traj = CheapPickleTrajectory('md.%s.traj' % fileid, 'w', atoms=system)

Tmax = 750
Tmin = 300
dyn = Langevin(system, 24.0 * fs, Tmax * kB, 0.06,
               trajectory=traj, logfile='md.%s.log' % fileid)
def recenter():
    system.center()
dyn.attach(recenter, 1)
MaxwellBoltzmannDistribution(system, Tmax * kB)

niter = 5 + int(len(system) * 0.5)

if world.rank == 0:
    fd = open('mdtemp.%s.log' % fileid, 'w')
for i, T in enumerate(range(Tmax, Tmin, -1)):
    if world.rank == 0:
        print >> fd, '%d: %f [%d]' % (i, T, 16*i)
    dyn.set_temperature(T * kB)
    dyn.run(niter)

opt = QuasiNewton(system, trajectory='opt.%s.traj' % fileid)
opt.run(fmax=0.05)
write('cluster.%s.traj' % fileid, system)