#!/usr/bin/env python
#Make video of Gamera magnetosphere run

# Standard modules
import argparse
from argparse import RawTextHelpFormatter
import os
import errno

# Third-party modules
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt

import matplotlib.gridspec as gridspec
import numpy as np
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

# Kaipy modules
import kaipy.kaiViz as kv
import kaipy.gamera.msphViz as mviz
import kaipy.remix.remix as remix
import kaipy.gamera.magsphere as msph
import kaipy.gamera.gampp as gampp
import kaipy.gamera.rcmpp as rcmpp

cLW = 0.25
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.
	"""
	#Defaults
	fdir = os.getcwd()
	ftag = "msphere"
	oDir = "vid2D"
	ts = 0    #[min]
	te = 200  #[min]
	dt = 60.0 #[sec]
	noIon = False
	noRCM = False
	Nblk = 1 #Number of blocks
	nID = 1 #Block ID of this job
	doJy = False
	doBz = False
	doBigRCM = False

	MainS = """Creates simple multi-panel figure for Gamera magnetosphere run
	Left Panel - Residual vertical magnetic field
	Right Panel - Pressure (or density) and hemispherical insets
	"""

	parser = argparse.ArgumentParser(description=MainS, formatter_class=RawTextHelpFormatter)
	parser.add_argument('-d',type=str,metavar="directory",default=fdir,help="Directory to read from (default: %(default)s)")
	parser.add_argument('-id',type=str,metavar="runid",default=ftag,help="RunID of data (default: %(default)s)")
	parser.add_argument('-o',type=str,metavar="directory",default=oDir,help="Subdirectory to write to (default: %(default)s)")
	parser.add_argument('-ts' ,type=int,metavar="tStart",default=ts,help="Starting time [min] (default: %(default)s)")
	parser.add_argument('-te' ,type=int,metavar="tEnd"  ,default=te,help="Ending time   [min] (default: %(default)s)")
	parser.add_argument('-dt' ,type=int,metavar="dt"    ,default=dt,help="Cadence       [sec] (default: %(default)s)")
	parser.add_argument('-Nblk' ,type=int,metavar="Nblk",default=Nblk,help="Number of job blocks (default: %(default)s)")
	parser.add_argument('-nID' ,type=int,metavar="nID"  ,default=nID,help="Block ID of this job [1-Nblk] (default: %(default)s)")
	#parser.add_argument('-nompi', action='store_true', default=noMPI,help="Don't show MPI boundaries (default: %(default)s)")
	parser.add_argument('-bz'   , action='store_true', default=doBz ,help="Show Bz instead of dBz (default: %(default)s)")
	parser.add_argument('-jy'   , action='store_true', default=doJy ,help="Show Jy instead of pressure (default: %(default)s)")
	parser.add_argument('-bigrcm', action='store_true',default=doBigRCM,help="Show entire RCM domain (default: %(default)s)")
	parser.add_argument('-noion', action='store_true', default=noIon,help="Don't show ReMIX data (default: %(default)s)")
	parser.add_argument('-norcm', action='store_true', default=noRCM,help="Don't show RCM data (default: %(default)s)")

	mviz.AddSizeArgs(parser)

	return parser

def main():
	#Defaults
	doDen = False
	doMPI = False #[Add MPI tiling]
	doBig = False #[Use big window]
	noMPI = False
	# Set up the command-line parser.
	parser = create_command_line_parser()
	#Finalize parsing
	args = parser.parse_args()
	fdir = args.d
	ftag = args.id
	ts  = args.ts
	te  = args.te
	dt  = args.dt
	oSub = args.o
	Nblk = args.Nblk
	nID = args.nID
	doJy = args.jy
	doBz = args.bz
	doBigRCM = args.bigrcm
	
	#Setup timing info
	tOut = np.arange(ts*60.0,te*60.0,dt)
	Nt = len(tOut)
	vO = np.arange(0,Nt)

	print("Writing %d outputs between minutes %d and %d"%(Nt,ts,te))
	if (Nblk>1):
		#Figure out work bounds
		dI = (Nt//Nblk)
		i0 = (nID-1)*dI
		i1 = i0+dI
		if (nID == Nblk):
			i1 = Nt #Make sure we get last bit
		print("\tBlock #%d: %d to %d"%(nID,i0,i1))
	else:
		i0 = 0
		i1 = Nt

	#Setup output directory
	oDir = fdir + "/" + oSub
	print("Writing output to %s"%(oDir))

	#Check/create directory if necessary
	if (not os.path.exists(oDir)):
		try:
			print("Creating directory %s"%(oDir))
			os.makedirs(oDir)
		except OSError as exc:
			if exc.errno == errno.EEXIST and os.path.isdir(oDir):
				pass
			else:
				raise

	#Get domain size
	xyBds = mviz.GetSizeBds(args)
	
	#---------
	#Figure parameters
	figSz = (12,7.5)


	#======
	#Init data
	gsph = msph.GamsphPipe(fdir,ftag)

	#Check for remix
	rcmChk = fdir + "/%s.mhdrcm.h5"%(ftag)
	rmxChk = fdir + "/%s.mix.h5"%(ftag)
	doRCM = os.path.exists(rcmChk)
	doMIX = os.path.exists(rmxChk)

	if (doRCM and (not args.norcm)):
		print("Found RCM data")
		rcmdata = gampp.GameraPipe(fdir,ftag+".mhdrcm")
		mviz.vP = kv.genNorm(1.0e-2,100.0,doLog=True)
		rcmpp.doEll = not doBigRCM
	if (doMIX and (not args.noion)):
		print("Found ReMIX data")
		
		
	#======
	#Setup figure
	fig = plt.figure(figsize=figSz)
	gs = gridspec.GridSpec(3,6,height_ratios=[20,1,1],hspace=0.025)
	

	AxL = fig.add_subplot(gs[0,0:3])
	AxR = fig.add_subplot(gs[0,3:])

	AxC1 = fig.add_subplot(gs[-1,0:2])
	AxC2 = fig.add_subplot(gs[-1,2:4])
	AxC3 = fig.add_subplot(gs[-1,4:6])

	cbM = kv.genCB(AxC2,kv.genNorm(remix.facMax),"FAC",cM=remix.facCM,Ntk=4)
	AxC2.xaxis.set_ticks_position('top')

	#Loop over sub-range
	for i in range(i0,i1):
		#Convert time (in seconds) to Step #
		nStp = np.abs(gsph.T-tOut[i]).argmin()+gsph.s0
		print("Minute = %5.2f / Step = %d"%(tOut[i]/60.0,nStp))
		npl = vO[i]

		AxL.clear()
		AxR.clear()

		Bz = mviz.PlotEqB(gsph,nStp,xyBds,AxL,AxC1,doBz=doBz)

		if (doJy):
			mviz.PlotJyXZ(gsph,nStp,xyBds,AxR,AxC3)
		else:
			mviz.PlotMerid(gsph,nStp,xyBds,AxR,doDen,doRCM,AxC3)

		gsph.AddTime(nStp,AxL,xy=[0.025,0.89],fs="x-large")
		gsph.AddSW(nStp,AxL,xy=[0.625,0.025],fs="small")

			
		#Add inset RCM plot
		if (doRCM and (not args.norcm)):
			AxRCM = inset_axes(AxL,width="30%",height="30%",loc=3)
			rcmpp.RCMInset(AxRCM,rcmdata,nStp,mviz.vP)
			AxRCM.contour(kv.reWrap(gsph.xxc),kv.reWrap(gsph.yyc),kv.reWrap(Bz),[0.0],colors=mviz.bz0Col,linewidths=mviz.cLW)
			rcmpp.AddRCMBox(AxL)

		if (doMIX and (not args.noion)):
			ion = remix.remix(rmxChk,nStp)
			gsph.AddCPCP(nStp,AxR,xy=[0.610,0.925])
			mviz.AddIonBoxes(gs[0,3:],ion)		

		#Add MPI decomp
		if (doMPI):
			mviz.PlotMPI(gsph,AxL)
			mviz.PlotMPI(gsph,AxR)

		fOut = oDir+"/vid.%04d.png"%(npl)
		kv.savePic(fOut,bLenX=45)


if __name__ == "__main__":
	main()