#!/usr/bin/env python
#Joins MPI restart into serial one

import argparse
import kaipy.gamera.magsphereRescale as upscl
from argparse import RawTextHelpFormatter
import numpy as np
import h5py
import os
import kaipy.kaiH5 as kh5

if __name__ == "__main__":
	dIn = os.getcwd()

	Ri = 3
	Rj = 6
	Rk = 1
	inid = "mpimsphere"
	
	nres = "0"

	outid = "msphere"
	grid = "lfmQ.h5"
	MainS = """Joins MPI-decomposed Gamera restart file into serial file

	
	(Ri,Rj,Rk) : Input MPI decomposition
	inid/nres : Run ID string and restart number, i.e. input file = inid.MPISTUFF.Res.#nres.h5
	outid : Output Run ID
	grid : Filename of grid corners file (with ghosts) generated by genLFM/genGrid
	"""

	parser = argparse.ArgumentParser(description=MainS, formatter_class=RawTextHelpFormatter)
	parser.add_argument('-i',metavar='inid',default=inid,help="Input Run ID string (default: %(default)s)")
	parser.add_argument('-n',type=int,metavar="nres",default=0,help="Restart number (default: %(default)s)")
	parser.add_argument('-Ri',type=int,metavar="Ri",default=Ri,help="i-Ranks (default: %(default)s)")
	parser.add_argument('-Rj',type=int,metavar="Rj",default=Rj,help="j-Ranks (default: %(default)s)")
	parser.add_argument('-Rk',type=int,metavar="Rk",default=Rk,help="k-Ranks (default: %(default)s)")
	parser.add_argument('-o',type=str,metavar="outid",default=outid,help="Output run ID (default: %(default)s)")
	parser.add_argument('-grid',type=str,metavar="grid",default=grid,help="Grid file to read from (default: %(default)s)")
	parser.add_argument('-d',type=str,metavar="directory",default=dIn,help="Directory to read from (default: %(default)s)")

	#Finalize parsing
	args = parser.parse_args()
	bStr = args.i
	nRes = args.n
	outid = args.o
	Ri = args.Ri
	Rj = args.Rj
	Rk = args.Rk
	grid = args.grid
	dIn = args.d

	#Open output file
	fOut = outid + ".Res.%05d.h5"%(nRes)
	oH5 = h5py.File(fOut,'w')

	G,M,G0 = upscl.PullRestartMPI(bStr,nRes,Ri,Rj,Rk,dIn,oH5)

	#Write main data
	print("Writing plasma and field data ...")
	oH5.create_dataset("Gas",data=G)
	if (G0 is not None):
		print("Writing Gas0")
		oH5.create_dataset("Gas0",data=G0)
	oH5.create_dataset("magFlux",data=M)
	gVals = ['X','Y','Z']

	fGrid = dIn  + "/" + grid
	print("Reading grid from %s ..."%(fGrid))
	
	iH5 = h5py.File(fGrid,'r')
	for g in gVals:
		oH5.create_dataset(g,data=iH5[g])

	iH5.close()
	oH5.close()		
