#!/usr/bin/env python
#
# H5 Input Output
#
# Last rev: 19/02/2021
from __future__ import print_function, division

import numpy as np, h5py

from ..cr              import cr
from ..partition_table import PartitionTable
from ..utils.common    import raiseError
from ..utils.parallel  import MPI_RANK, MPI_SIZE, MPI_COMM, mpi_reduce, is_rank_or_serial


PYQVARSI_H5_STRING  = 'pyQvarsi HDF5 exchange file'
PYQVARSI_H5_VERSION = (2,5)


def h5_write_attributes(file):
	'''
	Writes the attributes for an HDF5 using pyQvarsi
	'''
	file.attrs['Description'] = PYQVARSI_H5_STRING
	file.attrs['Version']     = PYQVARSI_H5_VERSION	

def h5_check_attributes(file):
	'''
	Check the attributes of an HDF5 to see if it matches
	with a pyQvarsi file
	'''
	# Check the file version
	version = tuple(file.attrs['Version'])
	if not version == PYQVARSI_H5_VERSION:
		raiseError('File version <%s> not matching the tool version <%s>!'%(str(file.attrs['Version']),str(PYQVARSI_H5_VERSION)))


def h5_write_metadata(file,metadata):
	'''
	Write metadata inside an HDF5
	'''
	group = file.create_group('METADATA') if not 'METADATA' in file.keys() else file['METADATA']
	# Create the datasets - all ranks
	mdict = {}
	for var in metadata.keys():
		mdict[var] = group[var] if var in group else group.create_dataset(var,(1,),dtype=metadata[var])
	# Store in the datasets - only rank 0 or serial
	if is_rank_or_serial(0):
		for var in mdict.keys(): mdict[var][:] = metadata[var]


def h5_write_partition_table(file,ptable,is_field=False,overwrite=True):
	'''
	Write partition table data for field and mesh HDF5
	'''
	if 'PARTITION' in file.keys() and not overwrite: return
	if 'PARTITION' in file.keys() and overwrite: del file['PARTITION']
	group = file.create_group('PARTITION') 
	pdset = {}
	nsubd = ptable.n_partitions
	if is_field:
		# Write datasets for field
		pdset['NSubD']  = group.create_dataset('NSubD', (1,),    dtype='i')
		pdset['Ids']    = group.create_dataset('Ids',   (nsubd,),dtype='i')
		pdset['Points'] = group.create_dataset('Points',(nsubd,),dtype='i')
	else:
		# Write datasets for mesh
		pdset['NSubD']      = group.create_dataset('NSubD',     (1,),    dtype='i')
		pdset['Ids']        = group.create_dataset('Ids',       (nsubd,),dtype='i')
		pdset['Elements']   = group.create_dataset('Elements',  (nsubd,),dtype='i')
		pdset['Points']     = group.create_dataset('Points',    (nsubd,),dtype='i')
		pdset['Boundaries'] = group.create_dataset('Boundaries',(nsubd,),dtype='i')
	pdset['has_master'] = group.create_dataset('has_master',(1,),dtype='i')
	# Store
	if is_rank_or_serial(0):
		if 'NSubD'      in pdset.keys(): pdset['NSubD'][:]      = nsubd
		if 'Ids'        in pdset.keys(): pdset['Ids'][:]        = ptable.Ids
		if 'Elements'   in pdset.keys(): pdset['Elements'][:]   = ptable.Elements
		if 'Points'     in pdset.keys(): pdset['Points'][:]     = ptable.Points
		if 'Boundaries' in pdset.keys(): pdset['Boundaries'][:] = ptable.Boundaries
		if 'has_master' in pdset.keys(): pdset['has_master'][:] = ptable.has_master

def h5_load_partition_table(file,is_field=False):
	'''
	Load partition table data for field and mesh HDF5
	'''
	group = file['PARTITION']
	# Load common parts
	nsubd  = int(group['NSubD'][0])
	ids    = np.array(group['Ids'][:])
	points = np.array(group['Points'][:])
	# Load parts that are only stored in mesh
	if not is_field:
		elements   = np.array(group['Elements'][:])
		boundaries = np.array(group['Boundaries'][:])
	else:
		elements   = np.zeros((nsubd,),np.int32)
		boundaries = np.zeros((nsubd,),np.int32)
	has_master = bool(group['has_master'][0])
	# Return a proper partition table
	return PartitionTable(nsubd,ids,elements,points,boundaries,has_master=has_master)

def h5_edit_partition_table(fname,ptable):
	'''
	Edit the partition table of an existing HDF5 file.
	And by edit I mean change it for the input partition table
	'''
	file     = h5py.File(fname,'a')
	is_field = True if not 'MESH' in file.keys() else False
	h5_write_partition_table(file,ptable,is_field=is_field,overwrite=True)
	file.close()


@cr('h5IO.save_field')
def h5_save_field(fname,ptable,xyz,instant,time,varDict,*,mode='w',mpio=True,write_master=False,metadata={}):
	'''
	Save a field in HDF5
	'''
	if mpio and not MPI_SIZE == 1:
		h5_save_field_parallel(fname,mode,ptable,xyz,instant,time,varDict,write_master,metadata)
	else:
		h5_save_field_serial(fname,mode,ptable,xyz,instant,time,varDict,metadata)

def h5_save_variables_field(file,ptable,xyz,varDict,npoints,write_master,serial=False):
	'''
	Store variables into an HDF5 file
	'''
	group = file.create_group('VARIABLES')
	# Generate the datasets
	dset  = {}
	dset['xyz'] = file.create_dataset('xyz',(npoints,xyz.shape[1]),dtype=xyz.dtype)
	for var in varDict.keys():
		v     = varDict[var]
		shape = (npoints,) if len(v.shape) == 1 else (npoints,v.shape[1])
		dset[var] = group.create_dataset(var,shape,dtype=v.dtype)
	# Write into the datasets
	if serial or MPI_RANK != 0 or write_master: # To be changed in NOMASTER
		# Obtain start and end points for the array
		istart,iend = ptable.partition_bounds(MPI_RANK if not serial else 0,'Points')
		# Store the data
		dset['xyz'][istart:iend,:] = xyz
		for var in varDict.keys():
			v = varDict[var]
			if len(v.shape) == 1: # Scalar field
				dset[var][istart:iend]   = v
			else: # Vectorial or tensorial field
				dset[var][istart:iend,:] = v

def h5_save_field_serial(fname,mode,ptable,xyz,instant,time,varDict,metadata={}):
	'''
	Save a field in HDF5 in serial mode
	'''
	# Open file for writing
	file  = h5py.File(fname,mode)
	h5_write_attributes(file)
	# Metadata
	h5_write_metadata(file,metadata)
	# Partition table
	ptable = PartitionTable.new(1,nelems=0,npoints=xyz.shape[0],has_master=False)
	h5_write_partition_table(file,ptable,is_field=True,overwrite=False)
	# Store number of points, instant and time
	group = file.create_group('FIELD')
	group.create_dataset('npoints',(1,),dtype='i',data=xyz.shape[0])
	group.create_dataset('instant',(1,),dtype='i',data=instant)
	group.create_dataset('time'   ,(1,),dtype='f',data=time)
	# Store the variables
	h5_save_variables_field(group,ptable,xyz,varDict,xyz.shape[0],True,serial=True)
	file.close()

def h5_save_field_parallel(fname,mode,ptable,xyz,instant,time,varDict,write_master=False,metadata={},overwrite=True):
	'''
	Save a field in HDF5 in parallel mode
	'''
	# Compute the total number of points
	npoints  = 0 if MPI_RANK == 0 and not write_master else xyz.shape[0]
	npointsG = int(mpi_reduce(npoints,op='sum',all=True))
	# Open file
	file  = h5py.File(fname,mode,driver='mpio',comm=MPI_COMM)
	h5_write_attributes(file)
	# Metadata
	h5_write_metadata(file,metadata)
	# Partition table
	h5_write_partition_table(file,ptable,is_field=True,overwrite=False)
	# Create groups for number of points, instant, time and coordinates
	group = file.create_group('FIELD')
	dset = {}
	dset['npoints'] = group.create_dataset('npoints',(1,),dtype='i')
	dset['instant'] = group.create_dataset('instant',(1,),dtype='i')
	dset['time']    = group.create_dataset('time'   ,(1,),dtype='f')
	# Rank 0 stores npoints, instant and time
	if is_rank_or_serial(0):
		dset['npoints'][:] = npointsG
		dset['instant'][:] = instant
		dset['time'][:]    = time
	# Store the variables
	h5_save_variables_field(group,ptable,xyz,varDict,npointsG,write_master)
	file.close()


@cr('h5IO.load_field')
def h5_load_field(fname,mpio=True,vars=[],inods=None):
	'''
	Load a field in HDF5
	'''
	if mpio and not MPI_SIZE == 1:
		return h5_load_field_parallel(fname,vars,inods)
	else:
		return h5_load_field_serial(fname,vars,inods)

def h5_load_variables_field(file,ptable,npoints,varList,inods):
	'''
	Load variables from an HDF5 file
	'''
	# Redo the partition table
	if not ptable.check_split():
		ptable = PartitionTable.new(MPI_SIZE,nelems=0,npoints=npoints)
	# Use the partition table to get the indices to read
	inods = ptable.get_partition_points(MPI_RANK) if inods is None else inods
	# Load node coordinates
	xyz     = np.array(file['xyz'][inods,:])
	varDict = {}
	# Load the variables in the varDict
	group = file['VARIABLES']
	if len(varList) == 0: varList = list(group.keys())
	for var in varList:
		if len(group[var].shape) == 1:
			# Scalar field
			varDict[var] = np.array(group[var][inods])
		else:
			# Vectorial field
			varDict[var] = np.array(group[var][inods,:])
	# Return
	return xyz, varDict, ptable

def h5_load_field_serial(fname,varList,inods):
	'''
	Load a field in HDF5 in serial
	'''
	# Open file for reading
	file    = h5py.File(fname,'r')
	h5_check_attributes(file)
	# Read partition table
	ptable  = h5_load_partition_table(file,is_field=True)
	# Read number of points, instant and time
	group   = file['FIELD']
	npoints = int(group['npoints'][0])
	instant = int(group['instant'][0])
	time    = float(group['time'][0])
	# Read variables
	xyz, varDict, ptable = h5_load_variables_field(group,ptable,npoints,varList,inods if not inods is None else np.s_[:])
	file.close()
	return ptable, xyz, instant, time, varDict

def h5_load_field_parallel(fname,varList,inods):
	'''
	Load a field in HDF5 in parallel
	'''
	# Open file for reading
	file    = h5py.File(fname,'r',driver='mpio',comm=MPI_COMM)
	h5_check_attributes(file)
	# Read partition table
	ptable  = h5_load_partition_table(file,is_field=True)
	# Read number of points, instant and time
	group   = file['FIELD']
	npoints = int(group['npoints'][0])
	instant = int(group['instant'][0])
	time    = float(group['time'][0])
	# Read variables
	xyz, varDict, ptable = h5_load_variables_field(group,ptable,npoints,varList,inods)
	file.close()
	return ptable, xyz, instant, time, varDict


@cr('h5IO.save_mesh')
def h5_save_mesh(fname,xyz,lnods,ltype,lninv,leinv,ptable,ngauss,*,mode='w',metadata={},mpio=True,write_master=False):
	'''
	Save a field in HDF5
	'''
	if mpio and not MPI_SIZE == 1:
		h5_save_mesh_parallel(fname,mode,xyz,lnods,ltype,lninv,leinv,ptable,ngauss,write_master,metadata)
	else:
		h5_save_mesh_serial(fname,mode,xyz,lnods,ltype,lninv,leinv,ptable,ngauss,metadata)

def h5_save_variables_mesh_serial(file,xyz,lnods,ltype,lninv,leinv,ngauss):
	'''
	Store the mesh variables into an HDF5 file
	'''
	# Compute the total number of points and elements
	nnodeG = lninv.shape[0]
	nelemG = leinv.shape[0]
	# Number of nodes, elements and gauss points
	file.create_dataset('nnodes' ,(1,),dtype='i',data=nnodeG)
	file.create_dataset('nelems' ,(1,),dtype='i',data=nelemG)
	file.create_dataset('ngauss' ,(1,),dtype='i',data=ngauss)
	# Arrays
	file.create_dataset('xyz'  ,(nnodeG,xyz.shape[1])  ,dtype=xyz.dtype  ,data=xyz)
	file.create_dataset('lnods',(nelemG,lnods.shape[1]),dtype=lnods.dtype,data=lnods)
	file.create_dataset('ltype',(nelemG,)              ,dtype=ltype.dtype,data=ltype)
	file.create_dataset('lninv',(nnodeG,)              ,dtype=lninv.dtype,data=lninv)
	file.create_dataset('leinv',(nelemG,)              ,dtype=leinv.dtype,data=leinv)

def h5_save_variables_mesh_parallel(file,ptable,xyz,lnods,ltype,lninv,leinv,ngauss,write_master):
	'''
	Store the mesh variables into an HDF5 file
	'''
	# Compute the total number of points and elements
	nnodeG = mpi_reduce(lninv.shape[0] if not np.all(np.isnan(xyz)) else 0.,op='sum',all=True)
	nelemG = mpi_reduce(leinv.shape[0] if not np.all(np.isnan(xyz)) else 0.,op='sum',all=True)
	dset   = {}
	# Number of nodes, elements and gauss points
	dset['nnodes']  = file.create_dataset('nnodes' ,(1,),dtype='i')
	dset['nelems']  = file.create_dataset('nelems' ,(1,),dtype='i')
	dset['ngauss']  = file.create_dataset('ngauss' ,(1,),dtype='i')
	# Arrays
	dset['xyz']     = file.create_dataset('xyz'  ,(nnodeG,xyz.shape[1])  ,dtype=xyz.dtype)
	dset['lnods']   = file.create_dataset('lnods',(nelemG,lnods.shape[1]),dtype=lnods.dtype)
	dset['ltype']   = file.create_dataset('ltype',(nelemG,)              ,dtype=ltype.dtype)
	dset['lninv']   = file.create_dataset('lninv',(nnodeG,)              ,dtype=lninv.dtype)
	dset['leinv']   = file.create_dataset('leinv',(nelemG,)              ,dtype=leinv.dtype)
	# Write the number of nodes, elements and gauss points
	if is_rank_or_serial(0):
		# Metadata
		dset['nnodes'][:]  = nnodeG
		dset['nelems'][:]  = nelemG
		dset['ngauss'][:]  = ngauss
	if MPI_RANK != 0 or write_master:
		# Partition for point arrays
		istart, iend = ptable.partition_bounds(MPI_RANK,'Points')
		# Store point arrays
		dset['xyz'][istart:iend,:]   = xyz
		dset['lninv'][istart:iend]   = lninv
		# Partition for element arrays
		istart, iend = ptable.partition_bounds(MPI_RANK,'Elements')
		# Store element arrays
		dset['lnods'][istart:iend,:] = lnods + istart # To global ids
		dset['ltype'][istart:iend]   = ltype
		dset['leinv'][istart:iend]   = leinv

def h5_save_mesh_serial(fname,mode,xyz,lnods,ltype,lninv,leinv,ptable,ngauss,metadata={}):
	'''
	Save a mesh in HDF5 in serial mode
	'''
	# Open file for writing
	file  = h5py.File(fname,mode)
	h5_write_attributes(file)
	# Metadata
	h5_write_metadata(file,metadata)
	# Partition table
	ptable = PartitionTable.new(1,nelems=lnods.shape[0],npoints=xyz.shape[0],has_master=False)
	h5_write_partition_table(file,ptable,is_field=False)
	# Mesh data
	group = file.create_group('MESH')
	h5_save_variables_mesh_serial(group,xyz,lnods,ltype,lninv,leinv,ngauss)
	file.close()

def h5_save_mesh_parallel(fname,mode,xyz,lnods,ltype,lninv,leinv,ptable,ngauss,write_master=False,metadata={}):
	'''
	Save a mesh in HDF5 in parallel mode
	'''
	# Open file
	file  = h5py.File(fname,mode,driver='mpio',comm=MPI_COMM)
	h5_write_attributes(file)
	# Metadata
	h5_write_metadata(file,metadata)
	# Partition table
	h5_write_partition_table(file,ptable,is_field=False)
	# Mesh data
	group = file.create_group('MESH')
	h5_save_variables_mesh_parallel(group,ptable,xyz,lnods,ltype,lninv,leinv,ngauss,write_master)
	file.close()


@cr('h5IO.save_meshAlya')
def h5_save_meshAlya(fname,xyz,lmast,codno,exnor,commu,ptable,consmas,massm,mode='a',mpio=True,write_master=False):
	'''
	Save a field in HDF5
	'''
	if mpio and not MPI_SIZE == 1:
		h5_save_meshAlya_parallel(fname,mode,xyz,lmast,codno,exnor,commu,ptable,consmas,massm,write_master)
	else:
		h5_save_meshAlya_serial(fname,mode,xyz,lmast,codno,exnor,commu,consmas,massm)

def h5_save_variables_meshAlya_serial(file,xyz,lmast,codno,exnor,commu,consmas,massm):
	'''
	Store the mesh variables into an HDF5 file
	'''
	# Compute the total number of points and elements
	nnodeG = xyz.shape[0]
	# Number of nodes, elements and gauss points
	file.create_dataset('consmas',(1,),dtype='i',data=consmas)
	# Optional arrays
	write_lmast = nnodeG == lmast.shape[0]
	write_codno = nnodeG == codno.shape[0]
	write_exnor = nnodeG == exnor.shape[0]
	write_commu = nnodeG == commu.shape[0]
	file.create_dataset('lmast',(nnodeG,)              ,dtype=lmast.dtype,data=lmast) if write_lmast else None
	file.create_dataset('codno',(nnodeG,codno.shape[1]),dtype=codno.dtype,data=codno) if write_codno else None
	file.create_dataset('exnor',(nnodeG,exnor.shape[1]),dtype=exnor.dtype,data=exnor) if write_exnor else None
	file.create_dataset('commu',(nnodeG,commu.shape[1]),dtype=commu.dtype,data=commu) if write_commu else None
	if not consmas:
		write_massm = nnodeG == massm.shape[0]
		file.create_dataset('massm',(nnodeG,),dtype=massm.dtype,data=massm) if write_massm else None

def h5_save_variables_meshAlya_parallel(file,ptable,xyz,lmast,codno,exnor,commu,consmas,massm,write_master):
	'''
	Store the mesh variables into an HDF5 file
	'''
	# Compute the total number of points and elements
	nnodeG = mpi_reduce(xyz.shape[0] if not np.all(np.isnan(xyz)) else 0.,op='sum',all=True)
	dset   = {}
	# Number of nodes, elements and gauss points
	dset['ngauss']  = file.create_dataset('ngauss' ,(1,),dtype='i')
	dset['consmas'] = file.create_dataset('consmas',(1,),dtype='i')
	# Optional arrays
	write_lmast     = nnodeG == mpi_reduce(lmast.shape[0] if not np.all(np.isnan(xyz)) else 0.,op='sum',all=True)
	write_codno     = nnodeG == mpi_reduce(codno.shape[0] if not np.all(np.isnan(xyz)) else 0.,op='sum',all=True)
	write_exnor     = nnodeG == mpi_reduce(exnor.shape[0] if not np.all(np.isnan(xyz)) else 0.,op='sum',all=True)
	write_commu     = nnodeG == mpi_reduce(commu.shape[0] if not np.all(np.isnan(xyz)) else 0.,op='sum',all=True)
	dset['lmast']   = file.create_dataset('lmast',(nnodeG,)              ,dtype=lmast.dtype) if write_lmast else None
	dset['codno']   = file.create_dataset('codno',(nnodeG,codno.shape[1]),dtype=codno.dtype) if write_codno else None
	dset['exnor']   = file.create_dataset('exnor',(nnodeG,exnor.shape[1]),dtype=exnor.dtype) if write_exnor else None
	dset['commu']   = file.create_dataset('commu',(nnodeG,commu.shape[1]),dtype=commu.dtype) if write_commu else None
	if not consmas:
		write_massm   = nnodeG == mpi_reduce(massm.shape[0] if not np.all(np.isnan(xyz)) else 0.,op='sum',all=True)
		dset['massm'] = file.create_dataset('massm',(nnodeG,),dtype=massm.dtype) if write_massm else None
	# Write the number of nodes, elements and gauss points
	if is_rank_or_serial(0):
		dset['consmas'][:] = consmas
	if MPI_RANK != 0 or write_master:
		# Partition for point arrays
		istart, iend = ptable.partition_bounds(MPI_RANK,'Points')
		# Optional arrays
		if dset['lmast'] is not None: dset['lmast'][istart:iend]   = lmast
		if dset['codno'] is not None: dset['codno'][istart:iend,:] = codno
		if dset['exnor'] is not None: dset['exnor'][istart:iend,:] = exnor
		if dset['commu'] is not None: dset['commu'][istart:iend,:] = commu
		if dset['massm'] is not None: dset['massm'][istart:iend]   = massm

def h5_save_meshAlya_serial(fname,mode,xyz,lmast,codno,exnor,commu,consmas,massm):
	'''
	Save a mesh in HDF5 in serial mode
	'''
	# Open file for writing
	file  = h5py.File(fname,mode)
	h5_write_attributes(file)
	# Mesh data
	group = file.create_group('MESH_ALYA')
	h5_save_variables_meshAlya_serial(group,xyz,lmast,codno,exnor,commu,consmas,massm)
	file.close()

def h5_save_meshAlya_parallel(fname,mode,xyz,lmast,codno,exnor,commu,ptable,consmas,massm,write_master=False):
	'''
	Save a mesh in HDF5 in parallel mode
	'''
	# Open file
	file  = h5py.File(fname,mode,driver='mpio',comm=MPI_COMM)
	h5_write_attributes(file)
	# Mesh data
	group = file.create_group('MESH_ALYA')
	h5_save_variables_meshAlya_parallel(group,ptable,xyz,lmast,codno,exnor,commu,consmas,massm,write_master)
	file.close()

@cr('h5IO.save_meshSOD2D')
def h5_save_meshSOD2D(fname,xyz,porder,commu,boucodes,bouconec,ndime,ptable,mode='a',mpio=True):
	'''
	Save a field in HDF5
	'''
	if mpio and not MPI_SIZE == 1:
		h5_save_meshSOD2D_parallel(fname,mode,xyz,porder,commu,boucodes,bouconec,ndime,ptable)
	else:
		h5_save_meshSOD2D_serial(fname,mode,xyz,porder,commu,boucodes,bouconec,ndime)

def h5_save_variables_meshSOD2D_serial(file,xyz,porder,commu,boucodes,bouconec,ndime):
	'''
	Store the mesh variables into an HDF5 file
	'''
	# Compute the total number of points and elements
	nnodeG = xyz.shape[0]
	nbounG = boucodes.shape[0]
	# Optional arrays
	file.create_dataset('ndime'   ,(1,)                      ,dtype='i'           ,data=ndime)
	file.create_dataset('porder'  ,(1,)                      ,dtype='i'           ,data=porder)
	file.create_dataset('boucodes',(nbounG,)                 ,dtype=boucodes.dtype,data=boucodes)
	file.create_dataset('bouconec',(nbounG,bouconec.shape[1]),dtype=bouconec.dtype,data=bouconec)
	file.create_dataset('commu'   ,(nnodeG,commu.shape[1])   ,dtype=commu.dtype   ,data=commu)

def h5_save_variables_meshSOD2D_parallel(file,xyz,porder,commu,boucodes,bouconec,ndime,ptable):
	'''
	Store the mesh variables into an HDF5 file
	'''
	# Compute the total number of points and elements
	nnodeG = mpi_reduce(xyz.shape[0] if not np.all(np.isnan(xyz)) else 0.,op='sum',all=True)
	nbounG = mpi_reduce(boucodes.shape[0] if not np.all(np.isnan(boucodes)) else 0.,op='sum',all=True)
	dset   = {}
	# Number of nodes, elements and gauss points
	dset['ndime']  = file.create_dataset('ndime' ,(1,),dtype='i')
	dset['porder'] = file.create_dataset('porder',(1,),dtype='i')
	# Optional arrays
	dset['boucodes'] = file.create_dataset('boucodes',(nbounG,)                 ,dtype=boucodes.dtype)
	dset['bouconec'] = file.create_dataset('bouconec',(nbounG,bouconec.shape[1]),dtype=bouconec.dtype)
	dset['commu']    = file.create_dataset('commu',   (nnodeG,commu.shape[1])   ,dtype=commu.dtype)
	if is_rank_or_serial(0):
		dset['ndime'][:]  = ndime
		dset['porder'][:] = porder
	# Partition for point arrays
	istart, iend = ptable.partition_bounds(MPI_RANK,'Boundaries')
	# Optional arrays
	dset['boucodes'][istart:iend] = boucodes
	dset['bouconec'][istart:iend,:] = bouconec
	istart, iend = ptable.partition_bounds(MPI_RANK,'Points')
	dset['commu'][istart:iend,:]    = commu

def h5_save_meshSOD2D_serial(fname,mode,xyz,porder,commu,boucodes,bouconec,ndime):
	'''
	Save a mesh in HDF5 in serial mode
	'''
	# Open file for writing
	file  = h5py.File(fname,mode)
	h5_write_attributes(file)
	# Mesh data
	group = file.create_group('MESH_SOD2D')
	h5_save_variables_meshSOD2D_serial(group,xyz,porder,commu,boucodes,bouconec,ndime)
	file.close()

def h5_save_meshSOD2D_parallel(fname,mode,xyz,porder,commu,boucodes,bouconec,ndime,ptable):
	'''
	Save a mesh in HDF5 in parallel mode
	'''
	# Open file
	file  = h5py.File(fname,mode,driver='mpio',comm=MPI_COMM)
	h5_write_attributes(file)
	# Mesh data
	group = file.create_group('MESH_SOD2D')
	h5_save_variables_meshSOD2D_parallel(group,xyz,porder,commu,boucodes,bouconec,ndime,ptable)
	file.close()


@cr('h5IO.load_mesh')
def h5_load_mesh(fname,mpio=True):
	'''
	Load a mesh in HDF5
	'''
	if mpio and not MPI_SIZE == 1:
		return h5_load_mesh_parallel(fname)
	else:
		return h5_load_mesh_serial(fname)

def h5_load_variables_mesh(file,ptable):
	'''
	Load variables from an HDF5 file
	'''
	repart  = False
	# Read the number of points, elements, Gauss points and consmas
	npoints = int(file['nnodes'][0])
	nelems  = int(file['nelems'][0])
	ngauss  = int(file['ngauss'][0])
	# If necessary, redo the partition table
	if not ptable.check_split():
		repart = True
		print(MPI_SIZE,nelems,npoints,ptable.has_master)
		ptable = PartitionTable.new(MPI_SIZE,nelems=nelems,npoints=npoints,has_master=ptable.has_master)
	# If we repartitioned we have a correct element partition
	# however the node partition will not be correct
	# Read element arrays
	istart, iend = ptable.partition_bounds(MPI_RANK, 'Elements')
	lnods = np.array(file['lnods'][istart:iend,:]) - istart
	ltype = np.array(file['ltype'][istart:iend])
	leinv = np.array(file['leinv'][istart:iend])
	# If necessary, update partition points
	if repart:
		inods = ptable.update_partition_points(MPI_RANK,npoints,conec=lnods+istart)
		ptable.update_points(inods.shape[0])
	# Read point arrays
	inods = ptable.get_partition_points(MPI_RANK)
	xyz   = np.array(file['xyz'][inods,:])
	lninv = np.array(file['lninv'][inods])
	# Return
	return xyz,lnods,ltype,lninv,leinv,ngauss,ptable

def h5_load_mesh_serial(fname):
	'''
	Load a mesh in HDF5 in serial mode
	'''
	# Open file for reading
	file  = h5py.File(fname,'r')
	h5_check_attributes(file)
	# Read partition table
	ptable = h5_load_partition_table(file,is_field=False)
	# Read mesh
	group = file['MESH']
	xyz,lnods,ltype,lninv,leinv,ngauss,ptable = h5_load_variables_mesh(group,ptable)
	file.close()
	return xyz,lnods,ltype,lninv,leinv,ptable,ngauss

def h5_load_mesh_parallel(fname):
	'''
	Load a mesh in HDF5 in parallel using a pre computed partition
	'''
	# Open file for reading
	file = h5py.File(fname,'r',driver='mpio',comm=MPI_COMM)
	h5_check_attributes(file)
	# Read partition table
	ptable = h5_load_partition_table(file,is_field=False)
	# Read mesh
	group = file['MESH']
	xyz,lnods,ltype,lninv,leinv,ngauss,ptable = h5_load_variables_mesh(group,ptable)
	file.close()
	return xyz,lnods,ltype,lninv,leinv,ptable,ngauss

@cr('h5IO.load_meshAlya')
def h5_load_meshAlya(fname,ptable,mpio=True):
	'''
	Load a mesh in HDF5
	'''
	if mpio and not MPI_SIZE == 1:
		return h5_load_meshAlya_parallel(fname,ptable)
	else:
		return h5_load_meshAlya_serial(fname,ptable)

def h5_load_variables_meshAlya(file,ptable):
	'''
	Load variables from an HDF5 file
	'''
	consmas = bool(file['consmas'][0])
	# Partition table should be OK at this point
	# Read point arrays
	inods = ptable.get_partition_points(MPI_RANK)
	lmast = np.array(file['lmast'][inods])   if 'lmast' in file.keys() else np.array([],np.int32)
	codno = np.array(file['codno'][inods,:]) if 'codno' in file.keys() else np.array([[]],np.int32)
	exnor = np.array(file['exnor'][inods,:]) if 'exnor' in file.keys() else np.array([[]],np.double)
	commu = np.array(file['commu'][inods,:]) if 'commu' in file.keys() else np.array([[]],np.int32)
	massm = np.array(file['massm'][inods])   if 'massm' in file.keys() else np.array([],np.double)
	# Return
	return lmast,codno,exnor,commu,consmas,massm

def h5_load_meshAlya_serial(fname,ptable):
	'''
	Load a mesh in HDF5 in serial mode
	'''
	# Open file for reading
	file    = h5py.File(fname,'r')
	h5_check_attributes(file)
	# Read Mesh
	group = file['MESH_ALYA']
	lmast,codno,exnor,commu,consmas,massm = h5_load_variables_meshAlya(group,ptable)
	file.close()
	return lmast,codno,exnor,commu,consmas,massm

def h5_load_meshAlya_parallel(fname,ptable):
	'''
	Load a mesh in HDF5 in parallel using a pre computed partition
	'''
	# Open file for reading
	file = h5py.File(fname,'r',driver='mpio',comm=MPI_COMM)
	h5_check_attributes(file)
	# Read Mesh
	group = file['MESH_ALYA']
	lmast,codno,exnor,commu,consmas,massm = h5_load_variables_meshAlya(group,ptable)
	file.close()
	return lmast,codno,exnor,commu,consmas,massm

@cr('h5IO.load_meshSOD2D')
def h5_load_meshSOD2D(fname,ptable,mpio=True):
	'''
	Load a mesh in HDF5
	'''
	if mpio and not MPI_SIZE == 1:
		return h5_load_meshSOD2D_parallel(fname,ptable)
	else:
		return h5_load_meshSOD2D_serial(fname,ptable)

def h5_load_variables_meshSOD2D(file,ptable):
	'''
	Load variables from an HDF5 file
	'''
	# Partition table should be OK at this point
	# Read point arrays
	ndime    = np.array(file['ndime'])[0]  if 'ndime' in file.keys() else np.array([[]],np.int32)
	porder   = np.array(file['porder'])[0] if 'porder' in file.keys() else np.array([[]],np.int32)

	istart, iend = ptable.partition_bounds(MPI_RANK,'Boundaries')
	boucodes = np.array(file['boucodes'][istart:iend])   if 'boucodes' in file.keys() else np.array([[]],np.int32)
	bouconec = np.array(file['bouconec'][istart:iend,:]) if 'bouconec' in file.keys() else np.array([[]],np.double)
	inods = ptable.get_partition_points(MPI_RANK)
	commu    = np.array(file['commu'][inods,:])          if 'commu' in file.keys() else np.array([[]],np.int32)
	# Return
	return ndime,porder,boucodes,bouconec,commu

def h5_load_meshSOD2D_serial(fname,ptable):
	'''
	Load a mesh in HDF5 in serial mode
	'''
	# Open file for reading
	file    = h5py.File(fname,'r')
	h5_check_attributes(file)
	# Read Mesh
	group = file['MESH_SOD2D']
	ndime,porder,boucodes,bouconec,commu = h5_load_variables_meshSOD2D(group,ptable)
	file.close()
	return ndime,porder,boucodes,bouconec,commu

def h5_load_meshSOD2D_parallel(fname,ptable):
	'''
	Load a mesh in HDF5 in parallel using a pre computed partition
	'''
	# Open file for reading
	file = h5py.File(fname,'r',driver='mpio',comm=MPI_COMM)
	h5_check_attributes(file)
	# Read Mesh
	group = file['MESH_SOD2D']
	ndime,porder,boucodes,bouconec,commu = h5_load_variables_meshSOD2D(group,ptable)
	file.close()
	return ndime,porder,boucodes,bouconec,commu
