from gpaw import parsize, parsize_bands
from gpaw.io.hdf5_highlevel import File, HyperslabSelection
from gpaw.mpi import world
import numpy as np

# TAU set-up
import pytau
pytau.setNode(world.rank)
pytau.setNode(world.rank)
tau_0 = pytau.profileTimer("Main")
tau_1 = pytau.profileTimer("OpenFile")
tau_2 = pytau.profileTimer("SetupArray")
tau_3 = pytau.profileTimer("Dataset")
tau_4 = pytau.profileTimer("Hyperslab")
tau_5 = pytau.profileTimer("Write")

pytau.start(tau_0)
# A = np.ones((25, 1024, 1024), dtype=float) # use about 200 MB

pytau.start(tau_1)
# parsize is the domain-decomposition, a tuple of three integers
# parsize_bands is the state-parallelization, an integer

file = File('psi.hdf5', mode='w', comm=world.get_c_object())
pytau.stop(tau_1)

# Global dimensions of the psi array. We make it 2D for simplicity:
# nbands x G
pytau.start(tau_2)
dtype = float
nbands = 600
Gx = 20
Gy = 20
Gz = 20

# Check to make sure parallelization options exist and 
# make sense. Additionally, paralleziation options must
# be integer divisible with the dimensions of the psi
#
# For a 64-node partion in VN mode, just hard code
# parsize = [4, 4, 4]
# parsize_bands = 4
# same as partition dimensions
#
assert parsize is not None
assert parsize_bands is not None
domain_size = parsize[0]*parsize[1]*parsize[2]
assert parsize_bands*domain_size == world.size 

assert nbands % parsize_bands == 0
mynbands = nbands / parsize_bands
assert Gx % parsize[0] == 0
assert Gy % parsize[1] == 0
assert Gz % parsize[2] == 0
myGx = Gx / parsize[0]
myGy = Gy / parsize[1]
myGz = Gz / parsize[2]

# Create a simple distributed psi array for testing:
# - even ranks in world communicator get 2.0
# - odd ranks in world communicator get 0.5 
global_shape = (nbands, Gx, Gy, Gz)
if world.rank % 2 == 0:
    factor = 2.0 # even ranks
else:
    factor = 0.5 # odd ranks
data = factor*np.ones((mynbands, myGx, myGy, myGz), dtype)
pytau.stop(tau_2)

pytau.start(tau_3)
dset = file.create_dataset('psi', global_shape, dtype)
pytau.stop(tau_3)

pytau.start(tau_4)
# Calculating indices is the hardest part. 
# Assume a simple 2D block distribution striding
# down the G-index first in the world communicator
band_rank = world.rank / domain_size
domain_rank = world.rank % domain_size
coord_x = domain_rank / (parsize[1]*parsize[2])
coord_y = (domain_rank - coord_x*parsize[1]*parsize[2]) / parsize[2]
coord_z = domain_rank - coord_x*parsize[1]*parsize[2] - coord_y*parsize[2]

# slice type is start, stop, step
step = 1
band_indices = slice(band_rank*mynbands, band_rank*mynbands + mynbands, step)
domain_indices_x = slice(coord_x*myGx, coord_x*myGx + myGx, step)
domain_indices_y = slice(coord_y*myGy, coord_y*myGy + myGy, step)
domain_indices_z = slice(coord_z*myGz, coord_z*myGz + myGz, step)

indices = (band_indices, domain_indices_x, domain_indices_y, domain_indices_z)
print indices
myslice = HyperslabSelection(indices, dset.shape)
pytau.stop(tau_4)

pytau.start(tau_5)
dset.write(data, selection=myslice, collective=True)
# world.barrier() # make sure everyone has finished writing
pytau.stop(tau_5)

dset.close()
file.close()
del data # memory should go back to near initial levels
pytau.stop(tau_0)