import pickle
import numpy as np

def loaddensity():
    nx, ny, nz = map(int, open('parsize.dat').read().split())
    #globalshape = tuple(map(int, open('globalshape.dat').read().split()))
    rank = 0
    nt_xyz = []
    for x in range(nx):
        nt_yz = []
        for y in range(ny):
            nt_z = []
            for z in range(nz):
                fd = open('density.rank%d.spin%d.dat' % (rank, 0))
                nt_G = np.load(fd)
                nt_z.append(nt_G)
                print nt_G.shape
                rank += 1
            nt_z = np.concatenate(nt_z, -1)
            nt_yz.append(nt_z)
            #print nt_z.shape
        nt_yz = np.concatenate(nt_yz, -2)
        nt_xyz.append(nt_yz)
    nt_xyz = np.concatenate(nt_xyz, -3)
    return nt_xyz

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()

def main():
    from gpaw import GPAW
    from ase.data.molecules import molecule

    # note -- test runs on special number of cpus
    calc = GPAW(gpts=(32,40,36),
                nbands=8,
                basis='dzp',
                parallel=dict(domain=(1,1,2), band=2))
    system = molecule('H2O')
    system.center(vacuum=3.0)

    system.set_calculator(calc)
    system.get_potential_energy()

    savedensity(calc)

    from gpaw.mpi import world
    from ase.units import Bohr
    world.barrier()
    nt1_G = loaddensity()
    nt2_G = calc.get_pseudo_density(pad=False) * Bohr**3
    #nt_G = gd.collect(nt_G, broadcast)

    print 'SHAPE', nt1_G.shape, nt2_G.shape

    print 'val', nt1_G[5,6,7], nt2_G[5,6,7]

    err = np.abs(nt2_G - nt1_G).max()
    print 'ERR', err

    from ase.io import write
    write('nt1.cube', system, data=nt1_G)
    write('nt2.cube', system, data=nt2_G)

    #print nt1_G.shape, nt2_G.shape
    #assert nt1_G.shape == nt2_G.shape

if __name__ == '__main__':
    main()