#!/usr/bin/env python

import os, sys
import numpy as np

from ase import Atom, Atoms
from ase.units import Bohr
from gpaw import GPAW, restart
from gpaw.mpi import world
from gpaw.poisson import PoissonSolver
from gpaw.mixer import Mixer

def output(txt, comm=world):
    if comm.rank == 0:
        sys.stdout.write(txt)
        sys.stdout.flush()
    comm.barrier()

def prepare(comm=world, parallel=None):
    a = 6.0 # Ang
    d = 3.22 # Ang
    atoms = Atoms([Atom('Na', (a, d/2, a))],
                  pbc=(False, True, False),
                  cell=(2*a, d, 2*a))

    # Determine number of kpoints per axis
    pbc_c = atoms.get_pbc()
    cell_c = np.sum(atoms.get_cell()**2, axis=1)**0.5 / Bohr
    cellprod = 50 / Bohr
    kpts_c = 1 + (np.ceil(cellprod / cell_c).astype(int) - 1) * pbc_c

    # Determine number of grid points
    h = 0.2 / Bohr
    div = 16
    gpts_c = div * np.round(cell_c / (h*div))

    xc = 'LDA'
    if parallel is None:
        parallel = {}
    calc = GPAW(xc=xc,
                gpts=tuple(gpts_c),
                kpts=tuple(kpts_c),
                spinpol=False,
                nbands=4, #2,
                basis='dzp',
                poissonsolver=PoissonSolver(3, 'GS'),
                mixer=Mixer(0.05, 5, 50.0),
                convergence={'eigenstates': 1e-7, 'density': 1e-3},
                communicator=comm,
                parallel=parallel,
                )
    atoms.set_calculator(calc)
    return atoms, calc

if __name__ in ['__main__', '__builtin__']:

    """
    if world.rank == 0 and not os.path.isfile('Na.dzp.basis'):
        from gpaw.atom.basis import BasisMaker
        basis = BasisMaker('Na', 'dzp', xc=xc).generate(2, 1)
        basis.write_xml()
        del basis
    world.barrier()

    from gpaw import setup_paths
    setup_paths.insert(0, '.')
    ...
    del setup_paths[0]
    """

    filename = 'parallel_io_%02d.hdf5' % world.size
    atoms, calc = prepare()
    e0 = atoms.get_potential_energy()
    calc.write(filename, mode='all')
    del atoms, calc

    runs = {1: [(1,1)],
            2: [(1,1), (1,2), (2,1)],
            4: [(1,1), (1,2), (2,1), (1,4), (2,2), (4,1)],
            8: [(1,1), (1,2), (2,1), (1,4), (2,2), (4,1), (1,8), (2,4), (4,2), (8,1)],
            }

    # 16 bands (~25 electrons) and 8 spin/kpts (2,2,2)
    
    commsizes = 2**np.arange(np.floor(np.log2(world.size))+1, dtype=int)

    output('\n\nRestarting for commsizes: %s\n' % commsizes)

    for commsize in commsizes:
        for parsize, parsize_bands in runs[commsize]:
            #assert commsize % (parsize * parsize_bands) == 0
            comm = world.new_communicator(np.arange(commsize, dtype=int))
            if comm is not None:
                output('world: %d, ' % comm.size, comm)
                atoms, calc = restart(filename,
                                      parallel={'domain': parsize,
                                                'band': parsize_bands},
                                      communicator=comm,
                                      txt=None)
                calc.initialize_positions()
                calc.scf.reset()
                output('gd: %d, bd: %d, kd: %d, ' % (calc.wfs.gd.comm.size, calc.wfs.bd.comm.size, calc.wfs.kpt_comm.size), comm)
                e1 = atoms.get_potential_energy()
                niter = calc.get_number_of_iterations()
                output('niter: %d, e0: %f, e1: %f, diff: %f\n' % (niter, e0, e1, abs(e1-e0)), comm)
            world.barrier()