#Various helper routines for magnetosphere quick look plots

import argparse
from argparse import RawTextHelpFormatter
import matplotlib as mpl
import numpy as np
import kaipy.kaiViz as kv
import kaipy.kaiTools as kt
import kaipy.gamera.magsphere as msph
import kaipy.remix.remix as remix
import os

dbMax = 25.0
dbCM = "RdGy_r"
bzCM = "bwr"
cLW = 0.25
bz0Col = "magenta"
mpiCol = "deepskyblue"

jMax = 10.0 #Max current for contours

eMax = 5.0  #Max current for contours

#Default pressure colorbar
vP = kv.genNorm(vMin=1.0e-2,vMax=10.0,doLog=True)
szStrs = ['small','std','big','bigger','fullD','dm']
szBds = {}
szBds["std"]      = [-40.0 ,20.0,2.0]
szBds["big"]      = [-100.0,20.0,2.0]
szBds["bigger"]   = [-200.0,25.0,2.0]
szBds["fullD"]    = [-300.0,30.0,3.0] # full domain for double res
szBds["small"]    = [-10.0 , 5.0,2.0]
szBds["dm"]       = [-30.0 ,10.0,40.0/15.0]

#Add different size options to argument
def AddSizeArgs(parser):
	"""
	Add size arguments to the parser.

	Args:
		parser (argparse.ArgumentParser): The argument parser object.

	Returns:
		None
	"""
	parser.add_argument('-size', type=str, default="std", choices=szStrs, help="Domain bounds options (default: %(default)s)")


#Return domain size from parsed arguments
def GetSizeBds(args):

	szStr = args.size
	szBd = szBds[szStr]

	xTail = szBd[0]
	xSun  = szBd[1]
	yMax = (xSun-xTail)/szBd[2]
	xyBds = [xTail,xSun,-yMax,yMax]

	return xyBds

#Plot absolute error in the requested, or given, equatorial field
def PlotEqErrAbs(gsphP,gsphO,nStp,xyBds,Ax,fieldNames,AxCB=None,doClear=True,doDeco=True,vMin=1e-9,vMax=1e-4,doLog=True,doVerb=True):
	#specify two gsph objects as gsphPredicted (gsphP) and gsphObserved (gsphO)
	#
	normAbs = kv.genNorm(vMin,vMax,doLog=doLog)
	cmapAbs = "PRGn"
	if doLog:
		cmapAbs = "magma"
	
	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()

		cbString = ""
		for fn in fieldNames:
			cbString = cbString + "'" + fn + "', "
		cbString = cbString + " Absolute Error"
		kv.genCB(AxCB,normAbs,cbString,cM=cmapAbs)
	#Now do main plotting
	if (doClear):
		Ax.clear()
	dataAbs = None
	for fn in fieldNames:
		dataP = gsphP.EggSlice(fn,nStp,doEq=True,doVerb=doVerb)
		dataO = gsphO.EggSlice(fn,nStp,doEq=True,doVerb=doVerb)
		if dataAbs is None:
			dataAbs = np.square(dataO - dataP)
		else:
			dataAbs = dataAbs + np.square(dataO - dataP)
	dataAbs = np.sqrt(dataAbs)
	Ax.pcolormesh(gsphP.xxi,gsphP.yyi,dataAbs,cmap=cmapAbs,norm=normAbs)
	
	kv.SetAx(xyBds,Ax)
	
	if (doDeco):
		kv.addEarth2D(ax=Ax)
		Ax.set_xlabel('SM-X [Re]')
		Ax.set_ylabel('SM-Y [Re]')
	return dataAbs

#Plot relative error in the requested, or given, equatorial field
def PlotEqErrRel(gsphP,gsphO,nStp,xyBds,Ax,fieldNames,AxCB=None,doClear=True,doDeco=True,vMin=1e-16,vMax=1,doLog=True,doVerb=True):
	#specify two gsph objects as gsphPredicted (gsphP) and gsphObserved (gsphO)
	#
	normRel = kv.genNorm(vMin,vMax,doLog=doLog)
	cmapRel = "PRGn"
	if doLog:
		cmapRel = "magma"
	
	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()
		cbString = ""
		for fn in fieldNames:
			cbString = cbString + "'" + fn + "', "
		cbString = cbString + " Relative Error"
		kv.genCB(AxCB,normRel,cbString,cM=cmapRel)
	#Now do main plotting
	if (doClear):
		Ax.clear()
	dataAbs = None
	dataBase = None
	for fn in fieldNames:
		dataP = gsphP.EggSlice(fn,nStp,doEq=True,doVerb=doVerb)
		dataO = gsphO.EggSlice(fn,nStp,doEq=True,doVerb=doVerb)
		if dataAbs is None:
			dataBase = np.square(dataP)
			dataAbs = np.square(dataO - dataP)
		else:
			dataBase = dataBase + np.square(dataP)
			dataAbs = dataAbs + np.square(dataO - dataP)
	dataBase = np.sqrt(dataBase)
	dataAbs = np.sqrt(dataAbs)
	# math breaks when Base field is exactly 0
	dataBase[dataBase == 0] = np.finfo(np.float32).tiny
	dataRel = np.absolute(dataAbs/dataBase)
	Ax.pcolormesh(gsphP.xxi,gsphP.yyi,dataRel,cmap=cmapRel,norm=normRel)
	
	kv.SetAx(xyBds,Ax)
	
	if (doDeco):
		kv.addEarth2D(ax=Ax)
		Ax.set_xlabel('SM-X [Re]')
		Ax.set_ylabel('SM-Y [Re]')
	return dataRel

#Plot absolute error along the requested logical axis
def PlotLogicalErrAbs(gsphP,gsphO,nStp,Ax,fieldNames,meanAxis,AxCB=None,doClear=True,doDeco=True,vMin=1e-16,vMax=1,doLog=True,doVerb=True):
	#specify two gsph objects as gsphPredicted (gsphP) and gsphObserved (gsphO)
	#
	normAbs = kv.genNorm(vMin,vMax,doLog=doLog)
	cmapAbs = "PRGn"
	if doLog:
		cmapAbs = "magma"

	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()
		cbString = ""
		for fn in fieldNames:
			cbString = cbString + "'" + fn + "', "
		cbString = cbString + " Absolute Error along "
		if meanAxis == 0:
			cbString = cbString + "I axis"
		elif meanAxis == 1:
			cbString = cbString + "J axis"
		elif meanAxis == 2:
			cbString = cbString + "K axis"
		kv.genCB(AxCB,normAbs,cbString,cM=cmapAbs)
	#Now do main plotting
	if (doClear):
		Ax.clear()
	dataAbs = CalcTotalErrAbs(gsphP,gsphO,nStp,fieldNames,doVerb=doVerb,meanAxis=meanAxis)
	dataAbs = np.transpose(dataAbs) # transpose to put I/J on the horizontal axis
	Ax.pcolormesh(dataAbs,cmap=cmapAbs,norm=normAbs)

	if (doDeco):
		if meanAxis == 0:
			Ax.set_xlabel('J Indices')
			Ax.set_ylabel('K Indices')
		elif meanAxis == 1:
			Ax.set_xlabel('I Indices')
			Ax.set_ylabel('K Indices')
		elif meanAxis == 2:
			Ax.set_xlabel('I Indices')
			Ax.set_ylabel('J Indices')
	return dataAbs

#Plot relative error along the requested logical axis
def PlotLogicalErrRel(gsphP,gsphO,nStp,Ax,fieldNames,meanAxis,AxCB=None,doClear=True,doDeco=True,vMin=1e-16,vMax=1,doLog=True,doVerb=True):
	#specify two gsph objects as gsphPredicted (gsphP) and gsphObserved (gsphO)
	#
	normRel = kv.genNorm(vMin,vMax,doLog=doLog)
	cmapRel = "PRGn"
	if doLog:
		cmapRel = "magma"

	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()
		cbString = ""
		for fn in fieldNames:
			cbString = cbString + "'" + fn + "', "
		cbString = cbString + " Relative Error along "
		if meanAxis == 0:
			cbString = cbString + "I axis"
		elif meanAxis == 1:
			cbString = cbString + "J axis"
		elif meanAxis == 2:
			cbString = cbString + "K axis"
		kv.genCB(AxCB,normRel,cbString,cM=cmapRel)
	#Now do main plotting
	if (doClear):
		Ax.clear()
	dataRel = CalcTotalErrRel(gsphP,gsphO,nStp,fieldNames,doVerb=doVerb,meanAxis=meanAxis)
	dataRel = np.transpose(dataRel) # transpose to put I/J on the horizontal axis
	Ax.pcolormesh(dataRel,cmap=cmapRel,norm=normRel)

	if (doDeco):
		if meanAxis == 0:
			Ax.set_xlabel('J Indices')
			Ax.set_ylabel('K Indices')
		elif meanAxis == 1:
			Ax.set_xlabel('I Indices')
			Ax.set_ylabel('K Indices')
		elif meanAxis == 2:
			Ax.set_xlabel('I Indices')
			Ax.set_ylabel('J Indices')
	return dataRel

#Calculate total cumulative absolute error between two cases
def CalcTotalErrAbs(gsphP,gsphO,nStp,fieldNames,doVerb=True,meanAxis=None):
	dataAbs = None
	for fn in fieldNames:
		dataP = gsphP.GetVar(fn,nStp,doVerb=doVerb)
		dataO = gsphO.GetVar(fn,nStp,doVerb=doVerb)
		if dataAbs is None:
			dataAbs = np.square(dataO - dataP)
		else:
			dataAbs = dataAbs + np.square(dataO - dataP)
	dataAbs = np.sqrt(dataAbs)
	return np.mean(dataAbs,axis=meanAxis)

#Calculate total cumulative relative error between two cases
def CalcTotalErrRel(gsphP,gsphO,nStp,fieldNames,doVerb=True,meanAxis=None):
	dataAbs = None
	dataBase = None
	for fn in fieldNames:
		dataP = gsphP.GetVar(fn,nStp,doVerb=doVerb)
		dataO = gsphO.GetVar(fn,nStp,doVerb=doVerb)
		if dataAbs is None:
			dataBase = np.square(dataP)
			dataAbs = np.square(dataO - dataP)
		else:
			dataBase = dataBase + np.square(dataP)
			dataAbs = dataAbs + np.square(dataO - dataP)
	dataBase = np.sqrt(dataBase)
	dataAbs = np.sqrt(dataAbs)
	# math breaks when Base field is exactly 0
	dataBase[dataBase == 0] = np.finfo(np.float32).tiny
	dataRel = np.absolute(dataAbs/dataBase)
	return np.mean(dataRel,axis=meanAxis)

#Plot equatorial field
def PlotEqB(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True,doBz=False):
	vBZ = kv.genNorm(dbMax)
	vDB = kv.genNorm(dbMax)

	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()
		if (doBz):
			kv.genCB(AxCB,vBZ,"Vertical Field [nT]",cM=bzCM,Ntk=7)
		else:
			kv.genCB(AxCB,vDB,"Residual Field [nT]",cM=dbCM,Ntk=7)
	#Now do main plotting
	if (doClear):
		Ax.clear()
	Bz = gsph.EggSlice("Bz",nStp,doEq=True)
	if (doBz):
		Ax.pcolormesh(gsph.xxi,gsph.yyi,Bz,cmap=bzCM,norm=vBZ)
	else:
		dbz = gsph.DelBz(nStp)
		Ax.pcolormesh(gsph.xxi,gsph.yyi,dbz,cmap=dbCM,norm=vDB)
	Ax.contour(kv.reWrap(gsph.xxc),kv.reWrap(gsph.yyc),kv.reWrap(Bz),[0.0],colors=bz0Col,linewidths=cLW)

	kv.SetAx(xyBds,Ax)

	if (doDeco):
		kv.addEarth2D(ax=Ax)
		Ax.set_xlabel('SM-X [Re]')
		Ax.set_ylabel('SM-Y [Re]')
	return Bz

def PlotMerid(gsph,nStp,xyBds,Ax,doDen=False,doRCM=False,AxCB=None,doClear=True,doDeco=True,doSrc=False):
	CMx = "viridis"
	if (doDen):
		
		if (doRCM):
			vN = kv.genNorm(vMin=1.0,vMax=1.0e+3,doLog=True)
		else:
			vN = kv.genNorm(0,25)
		if (doSrc):
			vID = "SrcD"
			cbStr = "Source Density [#/cc]"
			
		else:
			vID = "D"
			cbStr = "Density [#/cc]"
		Q = gsph.EggSlice(vID,nStp,doEq=False)
	else:
		vN = vP
		if (doSrc):
			vID = "SrcP"
			cbStr = "Source Pressure [nPa]"
		else:
			vID = "P"
			cbStr = "Pressure [nPa]"
		Q = gsph.EggSlice(vID,nStp,doEq=False)
	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()
		kv.genCB(AxCB,vN,cbStr,cM=CMx)
	Ax.pcolormesh(gsph.xxi,gsph.yyi,Q,cmap=CMx,norm=vN)

	kv.SetAx(xyBds,Ax)
	if (doDeco):
		kv.addEarth2D(ax=Ax)
		Ax.set_xlabel('SM-X [Re]')
		Ax.set_ylabel('SM-Z [Re]')
		Ax.yaxis.tick_right()
		Ax.yaxis.set_label_position('right')

def PlotJyXZ(gsph,nStp,xyBds,Ax,AxCB=None,jScl=None,doDeco=True):
	if (jScl is None):
		#Just assuming current scaling is nA/m2
		jScl = 1.0
	vJ = kv.genNorm(jMax)
	jCMap = "PRGn"
	Nc = 15
	cVals = np.linspace(-jMax,jMax,Nc)

	if (AxCB is not None):
		AxCB.clear()
		kv.genCB(AxCB,vJ,"Jy [nA/m2]",cM=jCMap)
	Q = jScl*gsph.EggSlice("Jy",nStp,doEq=False)
	#Zero out first shell b/c bad derivative
	print(Q.shape)
	Q[0:2,:] = 0.0
	#Ax.contour(kv.reWrap(gsph.xxc),kv.reWrap(gsph.yyc),kv.reWrap(Q),cVals,norm=vJ,cmap=jCMap,linewidths=cLW)
	Ax.pcolormesh(gsph.xxi,gsph.yyi,Q,norm=vJ,cmap=jCMap)
	kv.SetAx(xyBds,Ax)
	if (doDeco):
		kv.addEarth2D(ax=Ax)
		Ax.set_xlabel('SM-X [Re]')
		Ax.set_ylabel('SM-Z [Re]')
		Ax.yaxis.tick_right()
		Ax.yaxis.set_label_position('right')

#Plot equatorial azimuthal electric field
def PlotEqEphi(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	vE = kv.genNorm(eMax)
	vEMap = "PRGn"
	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()
		kv.genCB(AxCB,vE,r"E$_{phi}$ [mV/m]",cM=vEMap)

	#Now do main plotting
	if (doClear):
		Ax.clear()
	Bx = gsph.EggSlice("Bx",nStp,doEq=True)
	By = gsph.EggSlice("By",nStp,doEq=True)
	Bz = gsph.EggSlice("Bz",nStp,doEq=True)
	Vx = gsph.EggSlice("Vx",nStp,doEq=True)
	Vy = gsph.EggSlice("Vy",nStp,doEq=True)
	Vz = gsph.EggSlice("Vz",nStp,doEq=True)

	# calculating some variables to to plot
	#E=-VxB
	Ex = -(Vy*Bz-Vz*By)*0.001 # [mV/m]
	Ey =  (Vx*Bz-Vz*Bx)*0.001
	Ez = -(Vx*By-Vy*Bx)*0.001

	# coordinate transform
	ppc = np.arctan2(gsph.yyc,gsph.xxc)
	theta = np.pi #eq plane
	Er,Et,Ep = kt.xyz2rtp(ppc,theta,Ex,Ey,Ez)

	Ax.pcolormesh(gsph.xxi,gsph.yyi,Ep,cmap=vEMap,norm=vE)

	kv.SetAx(xyBds,Ax)

	if (doDeco):
		kv.addEarth2D(ax=Ax)
		Ax.set_xlabel('SM-X [Re]')
		Ax.set_ylabel('SM-Y [Re]')
		Ax.yaxis.tick_right()
		Ax.yaxis.set_label_position('right')

#Add MPI contours
def PlotMPI(gsph,Ax,ashd=0.5):
	gCol = mpiCol
	for i in range(gsph.Ri):
		i0 = i*gsph.dNi
		Ax.plot(gsph.xxi[i0,:],gsph.yyi[i0,:],mpiCol,linewidth=cLW,alpha=ashd)

	if (gsph.Rj>1):
		for j in range(1,gsph.Rj):
			j0 = j*gsph.dNj
			Ax.plot(gsph.xxi[:,j0], gsph.yyi[:,j0],gCol,linewidth=cLW,alpha=ashd)
			Ax.plot(gsph.xxi[:,j0],-gsph.yyi[:,j0],gCol,linewidth=cLW,alpha=ashd)
		#X-axis (+)
		Ax.plot(gsph.xxi[:,0], gsph.yyi[:,0],gCol,linewidth=cLW,alpha=ashd)
		#X-axis (-)
		j0 = (gsph.Rj)*gsph.dNj
		Ax.plot(gsph.xxi[:,j0], gsph.yyi[:,j0],gCol,linewidth=cLW,alpha=ashd)

def AddIonBoxes(gs,ion):
	gsRM = gs.subgridspec(20,20)

	wXY = 6
	dX = 1
	dY = 1

	#Northern
	ion.init_vars('NORTH')
	ax = ion.plot('current'  ,gs=gsRM[dY:dY+wXY,dX:dX+wXY],doInset=True)
	#ax.set_title("North",fontsize="small")

	#Southern
	ion.init_vars('SOUTH')
	ax = ion.plot('current'  ,gs=gsRM[-dY-wXY:-dY,dX:dX+wXY],doInset=True)

def plotPlane(gsph,data,xyBds,Ax,AxCB,var='D',vMin=None,vMax=None,doDeco=True,cmap='viridis',doLog=False,midp=None):
	if (AxCB is not None):
		AxCB.clear()
	if (not midp):
		if (vMin is None):
			vMin = np.min(data)
		if (vMax is None):
			vMax = np.max(data)

	else:
		if ((vMin is None) and (vMax is None)):
			vMax = np.max(np.abs([np.min(data),np.max(data)]))
			vMin = -1.0*vMax

	vNorm = kv.genNorm(vMin,vMax=vMax,doLog=doLog,midP=midp)
	kv.genCB(AxCB,vNorm,cbT=var,cM=cmap,Ntk=7)
	Ax.pcolormesh(gsph.xxi,gsph.yyi,data,cmap=cmap,norm=vNorm)
	kv.SetAx(xyBds,Ax)

	return

def plotXY(gsph,nStp,xyBds,Ax,AxCB,var='D',vMin=None,vMax=None,doDeco=True,cmap='viridis',doLog=False,midp=None):
	data = gsph.EggSlice(var,nStp,doEq=True)
	plotPlane(gsph,data,xyBds,Ax,AxCB,var,vMin=vMin,vMax=vMax,doDeco=doDeco,cmap=cmap,doLog=doLog,midp=midp)
	if (doDeco):
		kv.addEarth2D(ax=Ax)
		Ax.set_xlabel('SM_X [Re]')
		Ax.set_ylabel('SM-Y [Re]')

	return data

def plotXZ(gsph,nStp,xzBds,Ax,AxCB,var='D',vMin=None,vMax=None,doDeco=True,cmap='viridis',doLog=False,midp=None):
	data = gsph.EggSlice(var,nStp,doEq=False)
	plotPlane(gsph,data,xzBds,Ax,AxCB,var,vMin=vMin,vMax=vMax,doDeco=doDeco,cmap=cmap,doLog=doLog,midp=midp)
	if (doDeco):
		kv.addEarth2D(ax=Ax)
	Ax.set_xlabel('SM_X [Re]')
	Ax.set_ylabel('SM-Z [Re]')
	return data
