#Various helper routines for heliosphere quick look plots

import argparse
from argparse import RawTextHelpFormatter
import matplotlib as mpl
import numpy as np

import kaipy.kaiViz as kv
import kaipy.gamhelio.heliosphere as hsph
import os

VMax = 800.
VMin = 300.
MagVCM = "inferno"

DMax = 150.
DMin = 2000.
DCM = "copper_r"

D0Max = 15.
D0Min = 1.
D0CM = "copper_r"

TMax = 1.
TMin = 0.01
TCM = "copper"

BMax = 150.
BMin = -150.
BCM = "coolwarm"


#Function to Add different size options to argument
#not used for helio right now
def AddSizeArgs(parser):
	parser.add_argument('-small' , action='store_true', default=False,help="Use smaller domain bounds (default: %(default)s)")
	parser.add_argument('-big'   , action='store_true', default=False,help="Use larger domain bounds (default: %(default)s)")
	parser.add_argument('-bigger', action='store_true', default=False,help="Use larger-er domain bounds (default: %(default)s)")
	parser.add_argument('-huge'  , action='store_true', default=False,help="Use huge domain bounds (default: %(default)s)")

#Return domain size from parsed arguments; see msphViz for options
def GetSizeBds(pic):
	if (pic == "pic1" or pic == "pic2"):
		xyBds = [-216.,216.,-216.,216.]
	elif (pic == "pic3"):
		xyBds = [0.,360.,-75.,75.]
	else:		
		print ("No pic type specified.")
	return xyBds

#Plot speed in equatorial plane
def PlotEqMagV(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	vMagV = kv.genNorm(VMin, VMax, doLog=False, midP=None)
	
	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()
		kv.genCB(AxCB,vMagV,"Speed [km/s]",cM=MagVCM,Ntk=7)

	#Now do main plotting
	if (doClear):
		Ax.clear()

	MagV = gsph.eqMagV(nStp)
	Ax.pcolormesh(gsph.xxi,gsph.yyi,MagV,cmap=MagVCM,norm=vMagV)

	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('X [R_S]')
		Ax.set_ylabel('Y [R_S]')
	return MagV

#Plot speed in meridional plane Y=0
def PlotMerMagV(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
        vMagV = kv.genNorm(VMin, VMax, doLog=False, midP=None)

        if (AxCB is not None):
                #Add the colorbar to AxCB
                AxCB.clear()
                kv.genCB(AxCB,vMagV,"Speed [km/s]",cM=MagVCM,Ntk=7)

        if (doClear):
                Ax.clear()
	#r is for +X plane and l is for -X plane
        Vr, Vl = gsph.MerMagV(nStp)
	#cell corners
        xr, zr, xl, zl, r = gsph.MeridGridHalfs()
        Ax.pcolormesh(xr,zr,Vr,cmap=MagVCM,norm=vMagV)
        Ax.pcolormesh(xl,zl,Vl,cmap=MagVCM,norm=vMagV)

        kv.SetAx(xyBds,Ax)

        if (doDeco):
                Ax.set_xlabel('X [R_S]')
                Ax.set_ylabel('Z [R_S]')
        return Vr, Vl

#Plot normalized density n(r/r0)^2 in meridional plane Y=0
def PlotMerDNorm(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	vD = kv.genNorm(DMin, DMax, doLog=False, midP=None)

	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()
		kv.genCB(AxCB,vD,r"Density n$(r/r_0)^2$ [cm$^{-3}$]",cM=DCM,Ntk=7)

	if (doClear):
		Ax.clear()

	Dr, Dl = gsph.MerDNrm(nStp)
	xr, zr, xl, zl, r = gsph.MeridGridHalfs()
	Ax.pcolormesh(xr,zr,Dr,cmap=DCM,norm=vD, shading='auto')
	Ax.pcolormesh(xl,zl,Dl,cmap=DCM,norm=vD, shading='auto')

	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('X [R_S]')
		Ax.set_ylabel('Z [R_S]')
		Ax.yaxis.tick_right()
		Ax.yaxis.set_label_position('right')
	return Dr, Dl

#Plot normalized Br Br(r/r0)^2 in meridional plane
def PlotMerBrNorm(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	vB = kv.genNorm(BMin, BMax, doLog=False, midP=None)

	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()
		kv.genCB(AxCB,vB,r'Radial MF B$_r$(r/r$_0)^2$ [nT]',cM=BCM,Ntk=7)
	if (doClear):
		Ax.clear()

	Br_r, Br_l = gsph.MerBrNrm(nStp)
	xr, zr, xl, zl, r = gsph.MeridGridHalfs()
	Ax.pcolormesh(xr,zr,Br_r,cmap=BCM,norm=vB,shading='auto')
	Ax.pcolormesh(xl,zl,Br_l,cmap=BCM,norm=vB,shading='auto')

	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('X [R_S]')
		Ax.set_ylabel('Z [R_S]')
		Ax.yaxis.tick_right()
		Ax.yaxis.set_label_position('right')
	return Br_r, Br_l

#Plot normalized temperature T(r/r0) in meridional plane
def PlotMerTemp(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	vT = kv.genNorm(TMin, TMax, doLog=False, midP=None)

	if (AxCB is not None):
		AxCB.clear()
		kv.genCB(AxCB,vT, r'Temperature T(r/r$_0$) [MK]',cM=TCM,Ntk=7)
	if (doClear):
		Ax.clear()

	Tempr, Templ = gsph.MerTemp(nStp)
	xr, zr, xl, zl, r = gsph.MeridGridHalfs()
	Ax.pcolormesh(xr,zr,Tempr,cmap=TCM,norm=vT)
	Ax.pcolormesh(xl,zl,Templ,cmap=TCM,norm=vT)

	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('X [R_S]')
		Ax.set_ylabel('Z [R_S]')
	return Tempr, Templ

#Plot normalized density in equatorial plane n(r/r0)^2
def PlotEqD(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	vD = kv.genNorm(DMin, DMax, doLog=False, midP=None)
	
	if (AxCB is not None):
		#Add the colorbar to AxCB
		AxCB.clear()
		kv.genCB(AxCB,vD,r"Density n(r/r$_0)^2$ [cm$^{-3}$]",cM=DCM,Ntk=7)

	#Now do main plotting
	if (doClear):
		Ax.clear()

	NormD = gsph.eqNormD(nStp)
	Ax.pcolormesh(gsph.xxi,gsph.yyi,NormD,cmap=DCM,norm=vD)

	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('X [R_S]')
		Ax.set_ylabel('Y [R_S]')
		Ax.yaxis.tick_right()
		Ax.yaxis.set_label_position('right')
	return NormD

#Plot normalized Temperature in equatorial plane T(r/r0)
def PlotEqTemp(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	vT = kv.genNorm(TMin, TMax, doLog=False, midP=None)

	if (AxCB is not None):
		AxCB.clear()
		kv.genCB(AxCB,vT,r"Temperature T(r/r$_0$) [MK]",cM=TCM,Ntk=7)
	if (doClear):
		Ax.clear()

	Temp = gsph.eqTemp(nStp)
	Ax.pcolormesh(gsph.xxi,gsph.yyi,Temp,cmap=TCM,norm=vT)
	
	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('X [R_S]')
		Ax.set_ylabel('Y [R_S]')
	return Temp

#Plor Br in equatorial plane
def PlotEqBr(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	vB = kv.genNorm(BMin, BMax, doLog=False, midP=None)

	if (AxCB is not None):
		AxCB.clear()
		kv.genCB(AxCB,vB,r'Radial MF B$_r$(r/r$_0)^2$ [nT]',cM=BCM,Ntk=7)
	if (doClear):
		Ax.clear()

	Br = gsph.eqNormBr(nStp)
	Ax.pcolormesh(gsph.xxi,gsph.yyi,Br,cmap=BCM,norm=vB)

	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('X [R_S]')
		Ax.set_ylabel('Y [R_S]')
		Ax.yaxis.tick_right()
		Ax.yaxis.set_label_position('right')
	return Br


#Plot Speed at 1 AU
def PlotiSlMagV(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	vMagV = kv.genNorm(VMin, VMax, doLog=False, midP=None)

	if (AxCB is not None):
                #Add the colorbar to AxCB
                AxCB.clear()
                kv.genCB(AxCB,vMagV,"Speed [km/s]",cM=MagVCM,Ntk=7)

	#Now do main plotting
	if (doClear):
		Ax.clear()

	V = gsph.iSliceMagV(nStp)
	lat, lon = gsph.iSliceGrid()
	Ax.pcolormesh(lon,lat,V,cmap=MagVCM,norm=vMagV)

	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('Longitude')
		Ax.set_ylabel('Latitude')
	return V

#Plot Density at 1 AU
def PlotiSlD(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	vD = kv.genNorm(D0Min, D0Max, doLog=False, midP=None)
	if (AxCB is not None):
		AxCB.clear()
		kv.genCB(AxCB,vD,"Number density [cm-3]",cM=D0CM,Ntk=7)

	if (doClear):
		Ax.clear()

	D = gsph.iSliceD(nStp)
	lat, lon = gsph.iSliceGrid()
	Ax.pcolormesh(lon,lat,D,cmap=D0CM,norm=vD)
	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('Longitude')
		Ax.set_ylabel('Latitude')
		Ax.yaxis.tick_right()
		Ax.yaxis.set_label_position('right')
	return D

#Plot Br and current sheet (Br=0) at 1 AU
def PlotiSlBr(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	BMin = -5.
	BMax = 5.
	vB = kv.genNorm(BMin, BMax, doLog=False, midP=None)
	if (AxCB is not None):
		AxCB.clear()
		kv.genCB(AxCB,vB,"Radial magnetic field [nT]",cM=BCM,Ntk=7)
	if (doClear):
		Ax.clear()

	Br = gsph.iSliceBr(nStp)
	lat, lon = gsph.iSliceGrid()
	#for contour cell-centered lon lat coordinates
	lon_c = 0.25*( lon[:-1,:-1]+lon[:-1,1:]+lon[1:,:-1]+lon[1:,1:] )
	lat_c = 0.25*( lat[:-1,:-1]+lat[:-1,1:]+lat[1:,:-1]+lat[1:,1:] )

	Ax.pcolormesh(lon,lat,Br,cmap=BCM,norm=vB)
	Ax.contour(lon_c, lat_c,Br,[0.],colors='black')
	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('Longitude')
		Ax.set_ylabel('Latitude')
		Ax.yaxis.tick_right()
		Ax.yaxis.set_label_position('right')
	return Br

#Plot Temperature at 1 AU
def PlotiSlTemp(gsph,nStp,xyBds,Ax,AxCB=None,doClear=True,doDeco=True):
	#colorbar limits
	TMin = 0.02
	TMax = 0.12
	vT = kv.genNorm(TMin, TMax, doLog=False, midP=None)

	if (AxCB is not None):
		AxCB.clear()
		kv.genCB(AxCB,vT,"Temperature [MK]",cM=TCM,Ntk=7)
	if (doClear):
		Ax.clear()

	Temp = gsph.iSliceT(nStp)	
	lat, lon = gsph.iSliceGrid()
	Ax.pcolormesh(lon,lat,Temp,cmap=TCM,norm=vT)

	kv.SetAx(xyBds,Ax)

	if (doDeco):
		Ax.set_xlabel('Longitude')
		Ax.set_ylabel('Latitude')
	return Temp


#Adds MPI contours
#this function is from magnetosphere Viz script. PlotMPI is not used for helio as of now 
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)
			
