#!/usr/bin/env python
#Takes one MPI-decomposed restart and spits out an upscaled MPI restart

# Standard modules
import argparse
import os
from argparse import RawTextHelpFormatter

# Third-party modules
import numpy as np
import h5py

# Kaipy modules
import kaipy.kaiH5 as kh5
import kaipy.gamera.magsphereRescale as upscl

def create_command_line_parser():
	"""Create the command-line argument parser.

	Create the parser for command-line arguments.

	Returns:
		argparse.ArgumentParser: Command-line argument parser for this script.
	"""
	dIn = os.getcwd()

	#Input tiling
	iRi = 4
	iRj = 4
	iRk = 1
	#Output tiling
	oRi = 8
	oRj = 8
	oRk = 1

	inid  = "msphere"
	outid = "msphere"

	nRes = "0"
	grid = "lfmQ.h5"

	MainS = """Upscales and retiles a Gamera MPI resart

	
	(iRi,iRj,iRk) : Input MPI decomposition
	(oRi,oRj,oRk) : Output 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 input 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('-iRi',type=int,metavar="iRi",default=iRi,help="Input i-Ranks (default: %(default)s)")
	parser.add_argument('-iRj',type=int,metavar="iRj",default=iRj,help="Input j-Ranks (default: %(default)s)")
	parser.add_argument('-iRk',type=int,metavar="iRk",default=iRk,help="Input 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('-oRi',type=int,metavar="oRi",default=oRi,help="Input i-Ranks (default: %(default)s)")
	parser.add_argument('-oRj',type=int,metavar="oRj",default=oRj,help="Input j-Ranks (default: %(default)s)")
	parser.add_argument('-oRk',type=int,metavar="oRk",default=oRk,help="Input k-Ranks (default: %(default)s)")
	
	parser.add_argument('-grid',type=str,metavar="grid",default=grid,help="inGrid file to read from (default: %(default)s)")
	parser.add_argument('--keep',action='store_true',default=False,help='Keep intermediate files (default: %(default)s)')
	parser.add_argument('--norescale',action='store_true',default=False,help='Do not rescale (up or down) (default: %(default)s)')
	parser.add_argument('--down',action='store_true',default=False,help='Downscale instead of upscale (default: %(default)s)')

	return parser

def main():
	parser = create_command_line_parser()

	#Finalize parsing
	args = parser.parse_args()
	bStr = args.i
	nRes = args.n
	outid = args.o
	iRi = args.iRi
	iRj = args.iRj
	iRk = args.iRk

	oRi = args.oRi
	oRj = args.oRj
	oRk = args.oRk

	grid = args.grid
	doKeep = args.keep
	doUp = not args.down
	doRescale = not args.norescale

#Pull tiled restart, write to temp file
	#Stupidly writing temp restart to reuse old code
	fTmp = "tempRes.31337.h5" # temporarily written at the run directory.
	oH5 = h5py.File(fTmp,'w') # fTmp needs to include all necessary variables. Use oH5.create_dataset below.

	G,M,G0,oG,oM = upscl.PullRestartMPI(bStr,nRes,iRi,iRj,iRk,dIn=None,oH5=oH5)

	#Write main data
	print("Writing plasma and field data to temp file...")
	oH5.create_dataset("Gas",data=G)
	if (G0 is not None):
		doGas0 = True
		print("Writing Gas0")
		oH5.create_dataset("Gas0",data=G0)
	else:
		doGas0 = False

	oH5.create_dataset("magFlux",data=M)
	oH5.create_dataset("oGas",data=oG)
	oH5.create_dataset("omagFlux",data=oM)
	gVals = ['X','Y','Z']

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

#Upscale from temp file
# taking an mpi-decomposed restart, lazily stitching it together and writing it to one serial restart
	fTmp2X = "tempRes.31337.2x.h5"
	
	#Open input and output
	oH5 = h5py.File(fTmp2X,'w')
	iH5 = h5py.File(fTmp,'r') # iH5 is now the object for fTmp.

	Ns,Nv,Nk,Nj,Ni = iH5['Gas'].shape
	G = np.zeros((Ns,Nv,Nk,Nj,Ni))
	M = np.zeros((3,Nk+1,Nj+1,Ni+1))
	G[:,:,:,:,:] = iH5['Gas'][:]
	M[  :,:,:,:] = iH5['magFlux'][:]
	oG = np.zeros((Ns,Nv,Nk,Nj,Ni))
	oM = np.zeros((3,Nk+1,Nj+1,Ni+1))
	oG[:,:,:,:,:] = iH5['oGas'][:]
	oM[  :,:,:,:] = iH5['omagFlux'][:]
	if (doGas0):
		G0[:,:,:,:,:] = iH5['Gas0'][:]

	X = iH5['X'][:]
	Y = iH5['Y'][:]
	Z = iH5['Z'][:]

	#Transfer attributes to output
	for k in iH5.attrs.keys():
		aStr = str(k)
		oH5.attrs.create(k,iH5.attrs[aStr])
	#Close input
	iH5.close()

	if (doRescale):
		if (doUp):
			print("Upscaling data ...")
			#Do upscaling
			Xr,Yr,Zr = upscl.upGrid(X,Y,Z)
			Gr = upscl.upGas(X,Y,Z,G,Xr.T,Yr.T,Zr.T)
			FluxR = upscl.upFlux(X,Y,Z,M,Xr,Yr,Zr)
			oGr = upscl.upGas(X,Y,Z,oG,Xr.T,Yr.T,Zr.T)
			oFluxR = upscl.upFlux(X,Y,Z,oM,Xr,Yr,Zr)
			if (doGas0):
				G0r = upscl.upGas(X,Y,Z,G0,Xr.T,Yr.T,Zr.T)
		else:
			print("Downscaling data ...")
			Xr,Yr,Zr = upscl.downGrid(X,Y,Z)
			Gr = upscl.downGas(X,Y,Z,G,Xr.T,Yr.T,Zr.T)
			FluxR = upscl.downFlux(X,Y,Z,M,Xr,Yr,Zr)
			oGr = upscl.downGas(X,Y,Z,oG,Xr.T,Yr.T,Zr.T)
			oFluxR = upscl.downFlux(X,Y,Z,oM,Xr,Yr,Zr)
			if (doGas0):
				G0r = upscl.downGas(X,Y,Z,G0,Xr.T,Yr.T,Zr.T)
	else:
		#No rescale, just set variables
		Xr = X.T #Adding transpose to be consistent w/ rescaling code
		Yr = Y.T
		Zr = Z.T
		Gr = G
		FluxR = M
		oGr = oG
		oFluxR = oM
		if (doGas0):
			G0r = G0
			
	#Write out grid to restart
	oH5.create_dataset("X",data=Xr.T)
	oH5.create_dataset("Y",data=Yr.T)
	oH5.create_dataset("Z",data=Zr.T)

	#Write out gas/flux variables
	oH5.create_dataset("Gas",data=Gr)
	oH5.create_dataset("magFlux",data=FluxR)
	oH5.create_dataset("oGas",data=oGr)
	oH5.create_dataset("omagFlux",data=oFluxR)
	if (doGas0):
		oH5.create_dataset("Gas0",data=G0r)

	#Close output
	oH5.close()

#Split up upscaled file
	if (doGas0):
		upscl.PushRestartMPI(outid,nRes,oRi,oRj,oRk,Xr.T,Yr.T,Zr.T,Gr,FluxR,oGr,oFluxR,fTmp2X,G0r)
	else:
		upscl.PushRestartMPI(outid,nRes,oRi,oRj,oRk,Xr.T,Yr.T,Zr.T,Gr,FluxR,oGr,oFluxR,fTmp2X)

#Delete temp files
	if (not doKeep):
		os.remove(fTmp)
		os.remove(fTmp2X)

if __name__ == "__main__":
	main()


# #!/usr/bin/env python
# #Takes one MPI-decomposed restart and spits out an upscaled MPI restart

# import argparse
# import kaipy.embiggenUtils as upscl
# from argparse import RawTextHelpFormatter
# import numpy as np
# import h5py
# import os
# import kaipy.kaiH5 as kh5

# if __name__ == "__main__":
# 	#Input tiling
# 	iRi = 4
# 	iRj = 4
# 	iRk = 1
# 	#Output tiling
# 	oRi = 8
# 	oRj = 8
# 	oRk = 1

# 	doFast = True
# 	doTest = False
# 	inid  = "msphere"
# 	outid = "msphere"

# 	nRes = "0"

# 	MainS = """Upscales and retiles a Gamera MPI resart

# 	(iRi,iRj,iRk) : Input MPI decomposition
# 	(oRi,oRj,oRk) : Output MPI decomposition
# 	inid/nres : Run ID string and restart number, i.e. input file = inid.MPISTUFF.Res.#nres.h5
# 	outid : Output Run ID
# 	"""	
# 	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('-iRi',type=int,metavar="iRi",default=iRi,help="Input i-Ranks (default: %(default)s)")
# 	parser.add_argument('-iRj',type=int,metavar="iRj",default=iRj,help="Input j-Ranks (default: %(default)s)")
# 	parser.add_argument('-iRk',type=int,metavar="iRk",default=iRk,help="Input 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('-oRi',type=int,metavar="oRi",default=oRi,help="Input i-Ranks (default: %(default)s)")
# 	parser.add_argument('-oRj',type=int,metavar="oRj",default=oRj,help="Input j-Ranks (default: %(default)s)")
# 	parser.add_argument('-oRk',type=int,metavar="oRk",default=oRk,help="Input k-Ranks (default: %(default)s)")
# 	parser.add_argument('--down',action='store_true',default=False,help='Downscale instead of upscale (default: %(default)s)')

# #Finalize parsing
# 	args = parser.parse_args()
# 	bStr = args.i
# 	nRes = args.n
# 	outid = args.o
# 	iRi = args.iRi
# 	iRj = args.iRj
# 	iRk = args.iRk

# 	oRi = args.oRi
# 	oRj = args.oRj
# 	oRk = args.oRk

# 	doUp = not args.down

# #Start by pulling tiled restart into one brick w/ halos
# 	X,Y,Z,nG,nM,nB,oG,oM,oB,G0,fIn = upscl.PullRestartMPI(bStr,nRes,iRi,iRj,iRk)

# 	if (doTest):
# 		rX = X
# 		rY = Y
# 		rZ = Z
# 		nrG = nG
# 		nrM = nM
# 		nrB = nB
# 		orG = oG
# 		orM = oM
# 		orB = oB

# 		rG0 = G0
# 		upscl.PushRestartMPI(outid,nRes,oRi,oRj,oRk,rX,rY,rZ,nrG,nrM,nrB,orG,orM,orB,rG0,fIn,dtScl=1.0)
# 	#Toy check
# 		upscl.CompRestarts(bStr,outid,nRes,iRi,iRj,iRk)

# 	else:

# 	#Do upscaling on all variables
# 		#Chop out last 2 cells on each side, then upscale

# 		#Start w/ grid
# 		rX,rY,rZ = upscl.upGrid(X,Y,Z)

# 		dVr = upscl.Volume(rX,rY,rZ)
# 		dV0 = upscl.Volume( X, Y, Z)

# 		#Face-centered fluxes
# 		nrM = upscl.upFlux(nM)

# 		#Now ready to do cell-centered variables
# 		nrG = upscl.upGas(nG,dV0,dVr,"Gas")
# 		nrB = upscl.upCCMag(nB,dV0,dVr,"Bxyz")
# 		if (G0 is not None):
# 			rG0 = upscl.upGas(G0,dV0,dVr,"Gas0")

# 		if (doFast):
# 			#Just replicate for oState
# 			orM = nrM
# 			orB = nrB
# 			orG = nrG
# 		else:
# 			orM = upscl.upFlux (oM)
# 			orB = upscl.upCCMag(oB,dV0,dVr,"Bxyz")
# 			orG = upscl.upGas  (oG,dV0,dVr,"Gas")

# 	#Push back out to arbitrary decomposition
# 		#Update dt0 by x1/2
# 		upscl.PushRestartMPI(outid,nRes,oRi,oRj,oRk,rX,rY,rZ,nrG,nrM,nrB,orG,orM,orB,rG0,fIn,dtScl=0.5)


