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

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

	#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)')


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




# #!/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)


