# Standard modules
import os
import datetime

# Third-party modules
import h5py as h5
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.dates as mdates
from astropy.time import Time

# Kaipy modules
import kaipy.kaiH5 as kh5
import kaipy.kaiViz as kv
import kaipy.kaijson as kj
import kaipy.kaiTools as kT
import kaipy.gamera.gampp as gampp
import kaipy.gamera.rcmpp as rcmpp
import kaipy.satcomp.scutils as scutils


#Optionals
doProgressBar = True
if doProgressBar:
	import progressbar

#Constants
massi = 1.67e-27 # mass of ions in kg
masse = 9.11e-31 # mass of lectrons in kg
ev = 1.607e-19 # 1ev in J
nt = 1.0e-9 # nt
re = 6.380e6 # radius of earth in m

pressure_factor   = 2./3.*ev/re*nt
specFlux_factor_i = 1/np.pi/np.sqrt(8)*np.sqrt(ev/massi)*nt/re/1.0e1  # [units/cm^2/keV/str]
specFlux_factor_e = 1/np.pi/np.sqrt(8)*np.sqrt(ev/masse)*nt/re/1.0e1  # [units/cm^2/keV/str]
TINY = 1.0e-8

to_center = lambda A: 0.25*(A[:-1,:-1]+A[1:,:-1]+A[:-1,1:]+A[1:,1:])

#Settings that should maybe be elsewhere
tkl_lInner = 2  # Exclude points within 2 Re

#Hard-coded filenames
RCM_TIME_JFNAME = 'rcm_times.json'
RCM_TKL_JFNAME = 'rcm_tkl.json'
RCM_WEDGE_JFNAME = 'rcm_wedgetkl.json'
RCM_EQLATLON_JFNAME = 'rcm_eqlatlon.json'
RCM_CUMULFRAC_JFNAME = 'rcm_cumulFrac.json'
MHDRCM_TIME_JFNAME = 'mhdrcm_times.json'

#Spacecraft strings for cdaweb retrieval
#Probably shouldn't be hard-coded
supportedSats = ["RBSPA", "RBSPB"]
supportedDsets = ["Hydrogen_omniflux_RBSPICE", 'Hydrogen_omniflux_RBSPICE_TOFXPHHHELT', "Electron_omniflux_RBSPICE", "Hydrogen_PAFlux_HOPE"]

#======
#Helpers
#======

#Generate json filename for given spacecraft dataset
def genSCD_jfname(jdir, scName, dSetName):
	"""
	Generates the JSON file name for a given satellite name and data set name.

	Args:
		jdir (str): The directory where the JSON file will be located.
		scName (str): The name of the satellite.
		dSetName (str): The name of the data set.

	Returns:
		str: The generated JSON file name.

	"""
	jfname = scName + "_" + dSetName + ".json"
	return os.path.join(jdir, jfname)


def genRCMTrack_jfname(jdir, scName):
	"""
	Generate the JSON file name for RCM track.

	Args:
		jdir (str): The directory where the JSON file will be located.
		scName (str): The name of the spacecraft.

	Returns:
		str: The full path of the JSON file name for RCM track.
	"""
	jfname = "rcmTrack_" + scName + ".json"
	return os.path.join(jdir, jfname)

#Return a 3D cube of a given variable to use in interpolation
def getVarCube(rcm5, varName, stepLow, stepHigh, ilon, ilat, k=None):
	def getVarCube(rcm5, varName, stepLow, stepHigh, ilon, ilat, k=None):
		"""
		Returns a variable cube from the given RCM5 dataset.

		Args:
			rcm5 (numpy.ndarray): The RCM5 dataset.
			varName (str): The name of the variable.
			stepLow (int): The lower time step.
			stepHigh (int): The higher time step.
			ilon (int): The longitude index.
			ilat (int): The latitude index.
			k (int, optional): The vertical level index. Defaults to None.

		Returns:
			numpy.ndarray: The variable cube.

		"""
	if k is None:
		v1 = np.expand_dims(rcm5[stepLow][varName][ilon:ilon+2, ilat:ilat+2], axis=2)
		v2 = np.expand_dims(rcm5[stepHigh][varName][ilon:ilon+2, ilat:ilat+2], axis=2)
	else:
		v1 = np.expand_dims(rcm5[stepLow][varName][k, ilon:ilon+2, ilat:ilat+2], axis=2)
		v2 = np.expand_dims(rcm5[stepHigh][varName][k, ilon:ilon+2, ilat:ilat+2], axis=2)
	return np.append(v1, v2, axis=2)

#electrons: kStart = kStart, kEnd = kIon
#ions: kStart = kIon, kEnd = len(rcmS0['alamc'])
def getSpecieslambdata(rcmS0, species='ions'):
	"""
	Calculate lambdata for a given species in the RCM model.

	Args:
		rcmS0 (dict): Dictionary containing RCM data.
		species (str): Species for which lambdata is calculated. Can be 'ions' or 'electrons'. Default is 'ions'.

	Returns:
		dict: Dictionary containing the calculated lambdata.

	"""
	# Determine what our bounds are
	kStart = 1 if rcmS0['alamc'][0] == 0 else 0  # Check channel 0 for plasmasphere
	kIon = (rcmS0['alamc'][kStart:] > 0).argmax() + kStart
	if species == 'electrons':
		kEnd = kIon
	elif species == 'ions':
		kStart = kIon
		kEnd = len(rcmS0['alamc'])

	ilamc = rcmS0['alamc'][kStart:kEnd]  # Cell centers
	Nk = len(ilamc)
	ilami = np.zeros(Nk + 1)  # Cell interfaces
	for n in range(0, Nk - 1):
		ilami[n + 1] = 0.5 * (ilamc[n] + ilamc[n + 1])
	ilami[Nk] = ilamc[-1] + 0.5 * (ilamc[-1] - ilamc[-2])

	ilamc = np.abs(ilamc)
	ilami = np.abs(ilami)
	lamscl = np.diff(ilami) * np.sqrt(ilamc)

	result = {
		'kStart': kStart,
		'kEnd': kEnd,
		'ilamc': ilamc,
		'ilami': ilami,
		'lamscl': lamscl
	}
	return result

# TODO Ask Anthony why not just use the one in kaipy.kaiViz
def genVarNorm(var, doLog=False):
	"""
	Generate normalized colorbar .

	Parameters:
	- var: numpy array
		The variable for which to generate normalized values.
	- doLog: bool, optional
		Flag indicating whether to use logarithmic normalization. Default is False.

	Returns:
	- norm: normalized colorbar object.

	"""
	vMin = np.min(var[var>0])
	vMax = np.max(var)
	norm = kv.genNorm(vMin, vMax, doLog=doLog)

	return norm

# TODO Ask Anthony why not just use the one in kaipy.kaiViz
def get_aspect(ax):
	"""
	Calculate the aspect ratio of an axis.

	Args:
		ax: The axis object for which to calculate the aspect ratio.

	Returns:
		The aspect ratio of the axis.

	"""
	from operator import sub

	# Total figure size
	figW, figH = ax.get_figure().get_size_inches()

	# Axis size on figure
	_, _, w, h = ax.get_position().bounds

	# Ratio of display units
	disp_ratio = (figH * h) / (figW * w)

	# Ratio of data units
	# Negative over negative because of the order of subtraction
	data_ratio = sub(*ax.get_ylim()) / sub(*ax.get_xlim())

	return disp_ratio #/ data_ratio

#======
#Main work
#======

#Given sc and dataset name (according to above dict), grab specifically omnidirecitonal differential flux
def getSCOmniDiffFlux(scName, dSetName, t0, t1, jdir=None, forceCalc=False):
	"""
	Retrieves spacecraft data and calculates omnidirectional differential flux.

	Args:
		scName (str): Name of the spacecraft.
		dSetName (str): Name of the dataset.
		t0 (float): Start time of the data retrieval.
		t1 (float): End time of the data retrieval.
		jdir (str, optional): Directory to store the JSON file. Defaults to None.
		forceCalc (bool, optional): Flag to force calculation even if JSON file exists. Defaults to False.

	Returns:
		ephdata (dict): Ephemeris data.
		dataset (dict): Dataset containing the calculated omnidirectional differential flux.

	Raises:
		None

	Examples:
		# Retrieve spacecraft data and calculate omnidirectional differential flux
		ephdata, dataset = getSCOmniDiffFlux("Satellite1", "Dataset1", 0, 10)

	"""
	if jdir is not None:
		dojson = True
		jfname = genSCD_jfname(jdir, scName, dSetName)
	else:
		dojson = False

	if dojson and not forceCalc:
		if os.path.exists(jfname):
			print("Grabbing spacecraft data from " + jfname)
			data = kj.load(jfname)
			ephdata = data['ephdata']
			dataset = data['dataset']
			return ephdata, dataset

	print("Pulling spacecraft data from cdaweb")
	# TODO: Add all desired datasets

	scStrs = scutils.getScIds()
	satStrs = scStrs[scName]
	ephemStrs = satStrs['Ephem']
	dsetStrs = satStrs[dSetName]
	epochStr = "Epoch" if "EpochStr" not in dsetStrs.keys() else dsetStrs['EpochStr']
	energyStr = dsetStrs['EnergyStr']

	# First get ephem data
	s, ephdata = scutils.pullVar(ephemStrs['Id'], ephemStrs['Data'], t0, t1)
	s, data = scutils.pullVar(dsetStrs['Id'], dsetStrs['Data'], t0, t1, doVerbose=True)

	dataset = {}
	dataset['name'] = scName
	species = 'electrons' if 'E' == dSetName[0] else 'ions'
	dataset['species'] = species
	dataset['epoch'] = data[epochStr]
	
	# Turn each dataset's data into omniflux
	if dSetName in ['Hydrogen_omniflux_RBSPICE', 'Hydrogen_omniflux_RBSPICE_TOFXPHHHELT', 'Electron_omniflux_RBSPICE']:
		# Already got omni flux, no problem
		ofStr = dsetStrs['OmnifluxStr']
		if dSetName == 'Hydrogen_omniflux_RBSPICE':
			dataset['OmniDiffFlux'] = data[ofStr]*1E-3  # Diferential flux [1/(MeV-cm^2-s-sr]*[1/MeV -> 1/keV]
			dataset['energies'] = data[energyStr]*1E3  # [MeV] -> [keV]
		else:
			dataset['OmniDiffFlux'] = data[ofStr]*1E-3  # Diferential flux [1/(MeV-cm^2-s-sr]*[1/MeV -> 1/keV]
			dataset['energies'] = data[energyStr]*1E3  # [MeV] -> [keV]
		
	elif dSetName == "Hydrogen_PAFlux_HOPE":
		# !!TODO: Properly calc OmniDiffFlux from FPDU. Currently just using pitch angle
		fluxStr = dsetStrs['PAFluxStr']
		iPA = np.abs(data['PITCH_ANGLE'] - 90).argmin()
		dataset['OmniDiffFlux'] = data[fluxStr][:,iPA,:]  # (epoch, energy)
		dataset['energies'] = data[energyStr]*1E-3  # [eV -> keV]

	# Pause to save to json
	if dojson:
		print("Saving to file")
		jdata = {'ephdata' : ephdata, 'dataset' : dataset}
		kj.dump(jfname, jdata)

	return ephdata, dataset

def getRCMtimes(rcmf5, mhdrcmf5, jdir=None, forceCalc=False):
	"""
	Grab RCM times, sIDs, and MJDs

	Args:
		rcmf5 (str): Path to the RCM hdf5 file.
		mhdrcmf5 (str): Path to the MHDRCM hdf5 file.
		jdir (str, optional): Directory to search for JSON files. Defaults to None.
		forceCalc (bool, optional): Flag to force calculation even if JSON files exist. Defaults to False.

	Returns:
		dict: Dictionary containing RCM time data.

	Notes:
		- If jdir is given, the function will try to find the files there. If not found, it will pull the data from the hdf5 files.
		- If JSON files exist and forceCalc is False, the function will load the data from the JSON files.
		- If JSON files do not exist or forceCalc is True, the function will calculate the RCM times based on the MHDRCM file.

	"""
	if jdir is not None:
		dojson = True
		rcmjfname = os.path.join(jdir, RCM_TIME_JFNAME)
		mhdrcmjfname = os.path.join(jdir, MHDRCM_TIME_JFNAME)
	else:
		dojson = False

	if dojson and not forceCalc:
		if os.path.exists(rcmjfname):
			print("Grabbing RCM time data from " + rcmjfname)
			# That's all that's needed, done
			return kj.load(rcmjfname)
		else:
			# Must create RCM MJD's based on MHDRCM file, so get MHDRCM data
			if os.path.exists(mhdrcmjfname):
				print("Grabbing MHDRCM time data from " + mhdrcmjfname)
				mhdrcmTimes = kj.load(mhdrcmjfname)
				mhdrcmT = mhdrcmTimes['T']
				mhdrcmMJDs = mhdrcmTimes['MJD']
			else:
				print("No usable files in " + jdir + ", grabbing all from hdf5's")

	if 'mhdrcmTimes' not in locals():
		print("Grabbing MHDRCM time data from " + mhdrcmf5)
		Nt, sIDs = kh5.cntSteps(mhdrcmf5)
		sIDs = np.sort(sIDs)
		mhdrcmT = kh5.getTs(mhdrcmf5, sIDs, aID='time')
		mhdrcmMJDs = kh5.getTs(mhdrcmf5, sIDs, aID='MJD')

		mhdrcmTimes = {'Nt': Nt,
					   'sIDs': sIDs,
					   'T': mhdrcmT,
					   'MJD': mhdrcmMJDs}
		# Pause to save to file
		if dojson:
			kj.dump(mhdrcmjfname, mhdrcmTimes)

	print("Grabbing RCM time data from " + rcmf5)
	Nt, sIDs = kh5.cntSteps(rcmf5)
	sIDs = np.sort(sIDs)
	sIDstrs = np.array(['Step#' + str(s) for s in sIDs])
	rcmT = kh5.getTs(rcmf5, sIDs, aID="time")

	if (mhdrcmT[:] == rcmT[:]).all():
		rcmMJDs = mhdrcmMJDs
	else:
		# !! This only works if all rcm steps are also in mhdrcm
		#    In the future: As long as a single timestep is in both, that MJD can be used to get all MJDs for rcm steps
		rcmMJDs = np.zeros((len(rcmT)))
		for i in range(len(rcmT)):
			idx = np.where(mhdrcmT == rcmT[i])[0]
			rcmMJDs[i] = mhdrcmMJDs[idx]

	rcmTimes = {'rcmf5': rcmf5,
				'mhdrcmf5': mhdrcmf5,
				'Nt': Nt,
				'sIDs': sIDs,
				'sIDstrs': sIDstrs,
				'T': rcmT,
				'MJD': rcmMJDs}
	# Pause to save rcm jfile
	if dojson:
		kj.dump(rcmjfname, rcmTimes)

	return rcmTimes

def getRCM_scTrack(trackf5, rcmf5, rcmTimes, jdir=None, forceCalc=False, scName=""):
	"""
	Pull RCM data along a given spacecraft track.

	Args:
		trackf5 (str): Path to the big spacecraft trajectory hdf5 file, generated from sctrack.x.
		rcmf5 (str): Path to the <tag>.rcm.h5 file.
		rcmTimes (dict): Dictionary containing RCM times information, including 'sIDstrs' and 'MJD'.
		jdir (str, optional): If included, will try to do json saving and loading to save time. Defaults to None.
		forceCalc (bool, optional): Whether to force calculation even if json file exists. Defaults to False.
		scName (str, optional): Name of the spacecraft. Defaults to "".

	Returns:
		dict: Dictionary containing RCM data along the spacecraft track, including time, MJD, MLAT, MLON, vm, energy, and eetas.

	"""
	if jdir is not None:
		dojson = True
		jfname = genRCMTrack_jfname(jdir, scName)
	else:
		dojson = False

	if dojson and not forceCalc:
		if os.path.exists(jfname):
			print("Grabbing RCM track data from " + jfname)
			return kj.load(jfname)

	print("Extracting RCM track data from " + rcmf5)
	kh5.CheckOrDie(trackf5)
	scMLATs = kh5.PullVar(trackf5, 'MLAT')
	scMLONs = kh5.PullVar(trackf5, 'MLON')
	scTs = kh5.PullVar(trackf5, 'T')
	scMJDs = kh5.PullVar(trackf5, 'MJDs')
	Nsc = len(scTs)

	#Get information for mirror ratio
	Bx = kh5.PullVar(trackf5,"Bx")
	By = kh5.PullVar(trackf5,"By")
	Bz = kh5.PullVar(trackf5,"Bz")
	Bmag = np.sqrt(Bx**2.0 + By**2.0 + Bz**2.0)
	Beq = kh5.PullVar(trackf5,"Beq")

	J0 = scutils.getJScl(Bmag,Beq)

	#Unpack what we need from rcmTimes
	sIDstrs = rcmTimes['sIDstrs']
	rcmMJDs = rcmTimes['MJD']

	#Init rcm h5 info
	rcm5 = h5.File(rcmf5,'r')
	rcmS0 = rcm5[sIDstrs[0]]
	Ni, Nj = rcmS0['aloct'].shape
	rcmMLAT = 90.0-rcmS0['colat'][0,:]*180/np.pi
	rcmMLON = rcmS0['aloct'][:,0]*180/np.pi
	rcmMLAT_min = np.min(rcmMLAT)
	rcmMLAT_max = np.max(rcmMLAT)

	#Init electron and ion dicts
	sdata = {}
	sdata['electrons'] = getSpecieslambdata(rcmS0, 'electrons')
	sdata['ions'     ] = getSpecieslambdata(rcmS0, 'ions')
	kStart_e = sdata['electrons']['kStart']
	kStart_i = sdata['ions']['kStart']
	Nk_e = len(sdata['electrons']['ilamc'])
	Nk_i = len(sdata['ions']['ilamc'])
	
	#Collect data along spacecraft track
	vms = np.zeros((Nsc))
	xmin = np.zeros((Nsc))
	ymin = np.zeros((Nsc))
	zmin = np.zeros((Nsc))
	energies_e = np.zeros((Nsc, Nk_e))
	energies_i = np.zeros((Nsc, Nk_i))
	eeta_e = np.zeros((Nsc, Nk_e))
	eeta_i = np.zeros((Nsc, Nk_i))
	diffFlux_e = np.zeros((Nsc, Nk_e))
	diffFlux_i = np.zeros((Nsc, Nk_i))

	nearest_i = np.zeros(Nsc)
	nearest_j = np.zeros(Nsc)
	
	if doProgressBar: bar = progressbar.ProgressBar(max_value=Nsc)

	for n in range(Nsc):
		if doProgressBar: bar.update(n)

		mjd_sc = scMJDs[n]
		mlat_sc = scMLATs[n]
		mlon_sc = scMLONs[n]
		#Make sure track and rcm domain overlap
		if mjd_sc < rcmMJDs[0] or mjd_sc > rcmMJDs[-1] or \
			mlat_sc < rcmMLAT_min or mlat_sc > rcmMLAT_max:
			continue

		# Get bounds in rcm space
		ilat = len(rcmMLAT)-1  # mlat_rcm goes from high to low
		while rcmMLAT[ilat] < mlat_sc: ilat -= 1
		ilon = 2
		while ilon < len(rcmMLON)-2 and rcmMLON[ilon+1] < mlon_sc: ilon += 1
		#while rcmMLON[ilon+1] < mlon_sc: ilon += 1
		imjd = 0
		while rcmMJDs[imjd+1] < mjd_sc: imjd += 1

		#For other things to use for less rigorous mapping
		nearest_i[n] = ilat if abs(mlat_sc - rcmMLAT[ilat]) < abs(mlat_sc - rcmMLAT[ilat+1]) else ilat+1
		nearest_j[n] = ilon if abs(mlon_sc - rcmMLON[ilon]) < abs(mlon_sc - rcmMLON[ilon+1]) else ilon+1

		latbnd = [rcmMLAT[ilat], rcmMLAT[ilat+1]]
		lonbnd = [rcmMLON[ilon], rcmMLON[ilon+1]]
		mjdbnd = [rcmMJDs[imjd], rcmMJDs[imjd+1]]		
		stepLow = sIDstrs[imjd]
		stepHigh = sIDstrs[imjd+1]

		vmcube = getVarCube(rcm5, 'rcmvm', stepLow, stepHigh, ilon, ilat)
		vms[n] = scutils.trilinterp(lonbnd, latbnd, mjdbnd, vmcube, mlon_sc, mlat_sc, mjd_sc)
		#Do the same for xeq, yeq, zeq
		xmincube = getVarCube(rcm5, 'rcmxmin', stepLow, stepHigh, ilon, ilat)
		ymincube = getVarCube(rcm5, 'rcmymin', stepLow, stepHigh, ilon, ilat)
		zmincube = getVarCube(rcm5, 'rcmzmin', stepLow, stepHigh, ilon, ilat)
		xmin[n] = scutils.trilinterp(lonbnd, latbnd, mjdbnd, xmincube, mlon_sc, mlat_sc, mjd_sc)
		ymin[n] = scutils.trilinterp(lonbnd, latbnd, mjdbnd, ymincube, mlon_sc, mlat_sc, mjd_sc)
		zmin[n] = scutils.trilinterp(lonbnd, latbnd, mjdbnd, zmincube, mlon_sc, mlat_sc, mjd_sc)

		def getSpecEetas(kOffset, Nk):
			eetas = np.zeros((Nk))
			for k in range(Nk):
				kr = k + kOffset
				eetacube = getVarCube(rcm5, 'rcmeeta', stepLow, stepHigh, ilon, ilat, kr)
				eetas[k] = scutils.trilinterp(lonbnd, latbnd, mjdbnd, eetacube, mlon_sc, mlat_sc, mjd_sc)
			return eetas

		eeta_e[n,:] = getSpecEetas(kStart_e, Nk_e)
		eeta_i[n,:] = getSpecEetas(kStart_i, Nk_i)
		energies_e[n,:] = vms[n]*sdata['electrons']['ilamc']
		energies_i[n,:] = vms[n]*sdata['ions']['ilamc']

		diffFlux_e[n,:] = J0[n]*specFlux_factor_e*energies_e[n,:]*eeta_e[n,:]/sdata['electrons']['lamscl']
		diffFlux_i[n,:] = J0[n]*specFlux_factor_i*energies_i[n,:]*eeta_i[n,:]/sdata['ions'     ]['lamscl']

	# Package everything together
	sdata['electrons']['energies'] = energies_e*1E-3  # [eV -> keV]
	sdata['electrons']['eetas'   ] = eeta_e
	sdata['electrons']['diffFlux'] = diffFlux_e
	sdata['ions'     ]['energies'] = energies_i*1E-3  # [eV -> keV]
	sdata['ions'     ]['eetas'   ] = eeta_i
	sdata['ions'     ]['diffFlux'] = diffFlux_i
	
	result = {}
	result['T'        ] = scTs
	result['MJD'      ] = scMJDs
	result['MLAT'     ] = scMLATs
	result['MLON'     ] = scMLONs
	result['vm'       ] = vms
	result['nearest_i'] = nearest_i
	result['nearest_j'] = nearest_j
	result['xmin'     ] = xmin
	result['ymin'     ] = ymin
	result['zmin'     ] = zmin
	result['eqmin'    ] = np.sqrt(xmin**2+ymin**2)
	result['xeq'      ] = kh5.PullVar(trackf5, "xeq")
	result['yeq'      ] = kh5.PullVar(trackf5, "yeq")
	result['Req'      ] = np.sqrt(result['xeq']**2 + result['yeq']**2)
	result['electrons'] = sdata['electrons']
	result['ions'     ] = sdata['ions']

	if dojson: kj.dump(jfname, result)

	return result


#TODO: Energy grid mapping in a nice, jsonizable way
#      Right now, just need to call this whenever you want it
def consolidateODFs(scData, rcmTrackData, eGrid=None, doPlot=False):
	"""
	Prepare the spacecraft and RCM track data for comparison.

	Args:
		scData (dict): Dictionary containing spacecraft data.
		rcmTrackData (dict): Dictionary containing RCM track data.
		eGrid (ndarray, optional): Energy grid to use. If not provided, fixed bins spanning the full energy range will be created.
		doPlot (bool, optional): Whether to plot the data. Default is False.

	Returns:
		result (dict): Dictionary containing the consolidated data.

	"""
	#scData determines which species we're using
	species = scData['species']
	rcmSpec = rcmTrackData[species]

	scEGrid = scData['energies']  # Might be 1D (fixed bins) or 2D (different energies over time)
	rcmEGrid = rcmSpec['energies']

	Nt_sc = len(scData['epoch'])
	#Manipulate sc EGrid, if 2D, so that time dimension is first index
	sce_shape = scEGrid.shape
	if len(sce_shape) > 1:
		timeAxis = -1
		
		for i in range(len(sce_shape)):
			if sce_shape[i] == Nt_sc:
				timeAxis = i
		if timeAxis == -1:
			print("Finding right time axis in sc data didn't work")
			return
		scEGrid = np.rollaxis(scEGrid, timeAxis, 0)
		
	#If given grid to use, do that
	#Else, make fixed bins spanning full energy range
	if eGrid is None:
		eMax = np.max([scEGrid.max(), rcmEGrid.max()])
		eMin = np.min([scEGrid[scEGrid>0].min(), rcmEGrid[rcmEGrid>0].min()])
		numPoints = 150
		eGrid = np.logspace(np.log10(eMin), np.log10(eMax), numPoints, endpoint=True)
	Ne = len(eGrid)
	e_lower = np.zeros(Ne)
	e_upper = np.zeros(Ne)
	for iEG in range(Ne):
		#e_lower[iEG] = eGrid[0] if iEG == 0 else (eGrid[iEG-1] + eGrid[iEG])/2
		#e_upper[iEG] = eGrid[-1] if iEG == Ne-1 else (eGrid[iEG+1] + eGrid[iEG])/2
		e_lower[iEG] = 2*eGrid[0]-eGrid[1] if iEG == 0 else (eGrid[iEG-1] + eGrid[iEG])/2
		e_upper[iEG] = 2*eGrid[-1]-eGrid[-2] if iEG == Ne-1 else (eGrid[iEG+1] + eGrid[iEG])/2
	
	#Map sc odf on given energy grid new new eGrid
	sc_odf = np.zeros((Nt_sc, Ne))
	for n in range(Nt_sc):
		
		#Might not need this check
		if len(sce_shape) > 1:
			#Assume scEGrid's are evenly spaced	
			scEG_interfaces = (scEGrid[n,1:]+scEGrid[n,:-1])/2
			scEG_lower = np.append(scEGrid[n,0], scEG_interfaces)
			scEG_upper = np.append(scEG_interfaces, scEGrid[n,-1])
			#scEG_lower = np.append(2*scEGrid[n,0]-scEGrid[n,1], scEG_interfaces)
			#scEG_upper = np.append(scEG_interfaces, 2*scEGrid[n,-1]-scEGrid[n,-2])
			mapWeights = scutils.getWeights_ConsArea(scEGrid[n], scEG_lower, scEG_upper, eGrid, e_lower, e_upper)
		else:
			scEG_interfaces = (scEGrid[1:]+scEGrid[:-1])/2
			scEG_lower = np.append(scEGrid[0], scEG_interfaces)
			scEG_upper = np.append(scEG_interfaces, scEGrid[-1])
			#scEG_lower = np.append(2*scEGrid[0]-scEGrid[1], scEG_interfaces)
			#scEG_upper = np.append(scEG_interfaces, 2*scEGrid[-1]-scEGrid[-2])
			mapWeights = scutils.getWeights_ConsArea(scEGrid, scEG_lower, scEG_upper, eGrid, e_lower, e_upper)

		for e in range(Ne):
			fillAmt = 0  # Fraction of given energy cell that's been filled with oldGrid stuff
			for ik in range(len(mapWeights[e])):
				k = mapWeights[e][ik][0]
				weight = mapWeights[e][ik][1]
				sc_odf[n,e] += weight*scData['OmniDiffFlux'][n,k]
				fillAmt += weight
			if fillAmt > 0.2 and fillAmt < 1:
				#print('n={},e={},sc_odf={},fillAmt={}'.format(n,e,sc_odf[n,e],fillAmt))
				sc_odf[n,e] /= fillAmt  # Fill cell
	
	Nt_rcm = len(rcmTrackData['T'])
	rcm_odf = np.zeros((Nt_rcm, len(eGrid)))
	#TODO: For RCM, we should re-map eetas and then recalc diffFlux
	for n in range(Nt_rcm):
		vm = rcmSpec['energies'][n,-1]/rcmSpec['ilamc'][-1]
		rcme_lower = rcmSpec['ilami'][:-1]*vm
		rcme_lower[0] = rcmSpec['energies'][n,0]  # Otherwise its zero, and it makes odd lines
		rcme_upper = rcmSpec['ilami'][1:]*vm
		mapWeights = scutils.getWeights_ConsArea(rcmSpec['energies'][n,:], rcme_lower, rcme_upper, eGrid, e_lower, e_upper)

		if doPlot:
			plt.scatter(eGrid, np.zeros(len(eGrid)), c="black")
			plt.scatter(e_lower, np.zeros(len(eGrid)), c="red")
			plt.scatter(rcmSpec['energies'][n], np.ones(len(rcme_lower)), c="black")
			plt.scatter(rcme_lower, np.ones(len(rcme_lower)), c="red")
			plt.xscale('log')
			plt.show()

		for e in range(Ne):
			for ik in range(len(mapWeights[e])):
				k = mapWeights[e][ik][0]
				weight = mapWeights[e][ik][1]
				rcm_odf[n,e] += weight*rcmSpec['diffFlux'][n,k]

	result = {}
	result['energyGrid'] = eGrid
	result['sc'] = {
		'name' : scData['name'],
		'time' : scData['epoch'],
		'diffFlux' : sc_odf}
	result['rcm'] = {
		'time' : rcmTrackData['MJD'],
		'origEGrid' : rcmEGrid,
		'origODF' : rcmSpec['diffFlux'],
		'diffFlux' : rcm_odf}

	return result

def getIntensitiesVsL(rcmf5, mhdrcmf5, sStart, sEnd, sStride, species='ions', eGrid=None, jdir=None, forceCalc=False):
	"""Calculate rcm intensities (summed diff flux).
	
	Args:
		rcmf5 (str): RCM HDF5 filename to pull xmin, ymin, zmin from.
		mhdrcmf5 (str): MHDRCM HDF5 filename to pull IOpen from.
		AxLvT (matplotlib.axes.Axes, optional): If given, will plot the resulting L shell vs. time intensity.
		jdir (str, optional): Give JSON directory to enable JSON usage (read/write results to file).
		forceCalc (bool, optional): If dataset already found in file, re-calculate anyways and overwrite it.
	
	Returns:
		result (dict): Dictionary containing the calculated intensities.
	
	"""

	if jdir is not None:
		dojson = True
		rcmjfname = os.path.join(jdir, RCM_TKL_JFNAME)
	else:
		dojson = False

	if dojson and not forceCalc:
		if os.path.exists(rcmjfname):
			print("Grabbing RCM tkl data from " + rcmjfname)
			return kj.load(rcmjfname)

	#Setup

	rcmTimes = getRCMtimes(rcmf5,mhdrcmf5,jdir=jdir)
	iTStart = np.abs(rcmTimes['sIDs']-sStart).argmin()
	iTEnd = np.abs(rcmTimes['sIDs']-sEnd).argmin()

	sIDstrs = rcmTimes['sIDstrs'][iTStart:iTEnd+1:sStride]
	nSteps = len(sIDstrs)
	rcm5 = h5.File(rcmf5,'r')
	rcmS0 = rcm5[sIDstrs[0]]
	mhdrcm5 = h5.File(mhdrcmf5,'r')

	if species == 'electrons':
		sf_factor = specFlux_factor_e
	elif species == 'ions':
		sf_factor = specFlux_factor_i
	alamData = getSpecieslambdata(rcmS0, species)

	alams_kxx = alamData['ilamc'][:, np.newaxis, np.newaxis]
	kStart = alamData['kStart']
	kEnd   = alamData['kEnd']
	Nk = kEnd - kStart

	nlbins = 50
	lbins = np.linspace(2, 15, nlbins)

	if eGrid is None:
		eMin = 1E2  # [eV]
		eMax = 1E6  # [eV]
		Ne = 50
		eGrid = np.logspace(np.log10(eMin), np.log10(eMax), Ne, endpoint=True)
	Ne = len(eGrid)

	#Add endpoints to catch values outside of desired ranges
	lbins = np.concatenate(([lbins[0]-TINY], lbins, [lbins[-1]+TINY]))
	eGrid = np.concatenate(([eGrid[0]-TINY], eGrid, [eGrid[-1]+TINY]))

	rcmodf_tkl = np.zeros((nSteps, Ne+2, nlbins+2))
	rcmpress_tkl = np.zeros((nSteps, Ne+2, nlbins+2))

	if doProgressBar: bar = progressbar.ProgressBar(max_value=nSteps)
	for n in range(nSteps):
		if doProgressBar: bar.update(n)

		rcmS = rcm5[sIDstrs[n]]
		mhdrcmS = mhdrcm5[sIDstrs[n]]

		IOpen = mhdrcmS['IOpen']
		Ni, Nj = IOpen.shape  # Default should be (179, 90)
		#Shorten rcm data to match Ni, Nj of mhdrcm5
		vms = rcmS['rcmvm'][2:, :]
		xmins = rcmS['rcmxmin'][2:,:]
		ymins = rcmS['rcmymin'][2:,:]
		zmins = rcmS['rcmzmin'][2:,:]
		eetas = rcmS['rcmeeta'][kStart:kEnd,2:,:]

		vms_xij = vms[np.newaxis,:,:]

		#Calculate L shell for whole plane
		L_arr = scutils.xyz_to_L(xmins, ymins, zmins)  # [i,j]

		le_counts = np.zeros((Ne+2, nlbins+2))  # Keep track of how many entries into each bin so we can divide by it to get the average later

		energies = vms_xij * alams_kxx  # Should be [Nk, Ni, Nj]

		iL_arr = np.array([np.abs(lbins-i).argmin() for i in L_arr.flatten()]).reshape((Ni, Nj))
		iE_arr = np.array([np.abs(eGrid-e).argmin() for e in energies.flatten()]).reshape((Nk, Ni, Nj))

		diffFlux_Nk = sf_factor*energies*eetas/alamData['lamscl'][:, np.newaxis, np.newaxis]  # [k,i,j]

		pressure_kij = pressure_factor*alams_kxx*eetas*vms_xij**2.5 * 1E9  # [Pa -> nPa]

		for i in range(Ni):
			for j in range(Nj):
				for k in range(Nk):
					rcmodf_tkl[n, iE_arr[k,i,j], iL_arr[i,j]] += diffFlux_Nk[k,i,j]
					rcmpress_tkl[n, iE_arr[k,i,j], iL_arr[i,j]] += pressure_kij[k,i,j]
					le_counts[iE_arr[k,i,j], iL_arr[i,j]] += 1

		#Normalize to get avg. per count
		rcmodf_tkl[n,:,:] /= le_counts
		rcmpress_tkl[n,:,:] /= le_counts

	#Trim off the extra values
	lbins = lbins[1:-1]
	eGrid = eGrid[1:-1]
	rcmodf_tkl = rcmodf_tkl[:,1:-1,1:-1]
	rcmpress_tkl = rcmpress_tkl[:,1:-1,1:-1]
	
	# Sum across energy to get total avg. per L shell
	rcmpress_tl = np.ma.sum(np.ma.masked_invalid(rcmpress_tkl), axis=1)
	rcmodf_tl = np.ma.sum(np.ma.masked_invalid(rcmodf_tkl), axis=1)


	result = {}
	result['T']	         = rcmTimes['T'  ][iTStart:iTEnd+1:sStride]
	result['MJD']        = rcmTimes['MJD'][iTStart:iTEnd+1:sStride]
	result['lambda']     = alamData['ilamc']
	result['L_bins']     = lbins
	result['energyGrid'] = eGrid
	#result['nrg_tkl'] = rcmnrg_tkl
	result['odf_tkl']    = rcmodf_tkl
	result['odf_tl']    = rcmodf_tl
	result['press_tkl']  = rcmpress_tkl
	result['press_tl']  = rcmpress_tl

	if dojson: kj.dump(rcmjfname, result)

	return result

#TODO: Something odd with odf calculation
#TODO: Break out interpolator/mapper so they can be tested on their own and used by other functions
def getVarWedge(rcmf5, mhdrcmf5, sStart, sEnd, sStride, wedge_deg, species='ions', rcmTimes=None, eGrid=None, lGrid=None, jdir=None, forceCalc=False):
	"""Take a slice/wedge centered along the x axis (eq space), calculate average <var> vs. L and E.

	Args:
		rcmf5 (str): Filename of the RCM data in HDF5 format.
		mhdrcmf5 (str): Filename of the MHD RCM data in HDF5 format.
		sStart (float): Starting step value.
		sEnd (float): Ending step value.
		sStride (float): Step stride value.
		width_deg (float): Wedge width in degrees, centered around the x axis in the equatorial mapping of RCM data.
		rcmTimes (dict): Dictionary returned from the getRCMTimes() function.
		eGrid (ndarray, optional): Specific energy grid to map k energies to. Defaults to None.
		lbins (ndarray, optional): 1D array of L values to map to. Defaults to None.
		jdir (str, optional): Directory to find json files in. Defaults to None.

	Returns:
		dict: A dictionary containing the calculated results.

	Raises:
		None
	"""

	if jdir is not None:
		dojson = True
		rcmjfname = os.path.join(jdir, RCM_WEDGE_JFNAME)
	else:
		dojson = False

	if dojson and not forceCalc:
		if os.path.exists(rcmjfname):
			print("Grabbing RCM wedge-tkl data from " + rcmjfname)
			return kj.load(rcmjfname)

	#Setup
	if rcmTimes is None:
		rcmTimes = getRCMtimes(rcmf5,mhdrcmf5,jdir=jdir)
	iTStart = np.abs(rcmTimes['sIDs']-sStart).argmin()
	iTEnd = np.abs(rcmTimes['sIDs']-sEnd).argmin()

	sIDstrs = rcmTimes['sIDstrs'][iTStart:iTEnd+1:sStride]
	nSteps = len(sIDstrs)
	rcm5 = h5.File(rcmf5,'r')
	rcmS0 = rcm5[sIDstrs[0]]
	mhdrcm5 = h5.File(mhdrcmf5,'r')

	if species == 'electrons':
		sf_factor = specFlux_factor_e
	elif species == 'ions':
		sf_factor = specFlux_factor_i
	alamData = getSpecieslambdata(rcmS0, species)
	Nk = len(alamData['ilamc'])
	
	if lGrid is None:
		lMin = -10
		lMax = 8
		Nl = 52
		lGrid = np.linspace(lMin, lMax, Nl, endpoint=True)
	Nl = len(lGrid)
	
	if eGrid is None:
		eMin = 1E2  # [eV]
		eMax = 1E6  # [eV]
		Ne = 50
		eGrid = np.logspace(np.log10(eMin), np.log10(eMax), Ne, endpoint=True)
	Ne = len(eGrid)

	odf_tkl = np.zeros((nSteps, Ne, Nl))
	press_tkl = np.zeros((nSteps, Ne, Nl))

	if doProgressBar: bar = progressbar.ProgressBar(max_value=nSteps)
	for n in range(nSteps):
		if doProgressBar: bar.update(n)

		mhdrcmS = mhdrcm5[sIDstrs[n]]
		rcmS = rcm5[sIDstrs[n]]

		IOpen = mhdrcmS['IOpen']
		Ni, Nj = IOpen.shape  # Default should be (179, 90)
		#Shorten rcm data to match Ni, Nj of mhdrcm5
		vms = rcmS['rcmvm'][2:, :]
		xmins = rcmS['rcmxmin'][2:,:]
		ymins = rcmS['rcmymin'][2:,:]
		zmins = rcmS['rcmzmin'][2:,:]
		eetas = rcmS['rcmeeta'][alamData['kStart']:alamData['kEnd'],2:,:]

		L_arr = scutils.xyz_to_L(xmins, ymins, zmins) # [i,j]
		theta_arr = np.arctan2(ymins,xmins)*180/np.pi%360

		#Collect indicies of points within spatial bounds
		iPointCloud = []
		lPointCloud = []

		for i in range(Ni):
			for j in range(Nj):
				l_ij = L_arr[i,j]
				theta_ij = theta_arr[i,j]
				#Collect indicies of points within spatial bounds
				
				if IOpen[i,j] >= 0:
					continue

				if (theta_ij > 360-wedge_deg/2 or theta_ij < wedge_deg/2):  # dayside wedge
					if l_ij >= lGrid[0] and l_ij <= lGrid[-1] and abs(l_ij) > tkl_lInner:
						iPointCloud.append([i,j])
						lPointCloud.append(l_ij)	
				elif (theta_ij > 180-wedge_deg/2 and theta_ij < 180+wedge_deg/2): #nightside wedge
					if -l_ij >= lGrid[0] and -l_ij <= lGrid[-1] and abs(l_ij) > tkl_lInner:
						iPointCloud.append([i,j])
						lPointCloud.append(-l_ij)
	
		iPointCloud = np.array(iPointCloud)
		lPointCloud = np.array(lPointCloud)
		#Sort index array in order of ascending L values
		lSort = lPointCloud.argsort()
		iPointCloud = iPointCloud[lSort]
		lPointCloud = lPointCloud[lSort]
		Npc = len(iPointCloud)
		
		if Npc == 0:
			continue
		#Calc distance-based mapping weights for each lGrid index
		lMapWeights = [ [] for l in range(Nl)] # Nl x nx2 ,  nx[iPC, lWeight]
		wFlag = 9999
		for iLG in range(Nl):
			iLPC_lower = 0  # iPointCloud/lPointCloud index, 1 grid point below current lGrid index
			iLPC_upper = 0  # iPointCloud/lPointCloud index, 1 grid point above current lGrid index
			width_lower = abs(lGrid[iLG] - lGrid[iLG-1]) if iLG > 1 else wFlag
			width_upper = abs(lGrid[iLG+1] - lGrid[iLG]) if iLG < Nl-1 else wFlag
			if iLG != 0: 
				while iLPC_lower < Npc-1 and lPointCloud[iLPC_lower] < lGrid[iLG-1]:
					iLPC_lower += 1
			#Set upper iPLC index to next highest lbin value
			if iLG < Nl-1: 
				while iLPC_upper+1 < Npc and lPointCloud[iLPC_upper+1] < lGrid[iLG+1]:
					iLPC_upper += 1

			iLG_WM = []  # weight and points for current lGrid index
			
			for iLPC in range(iLPC_lower, iLPC_upper+1):
				# Is this a point below or above current lGrid point
				if lPointCloud[iLPC] < lGrid[iLG]:
					weight = 1-abs(lGrid[iLG] - lPointCloud[iLPC])/width_lower					
				elif lPointCloud[iLPC] > lGrid[iLG]:
					weight = 1-abs(lPointCloud[iLPC] - lGrid[iLG])/width_upper

				#Cheating !! this may be breaking something in odf calc
				if abs(weight)>1:
					weight = 0
					#print('iLG={} iLPC={} lG[iLG]={} lPC[iLPC]={} weight={}'.format(iLG, iLPC, lGrid[iLG], lPointCloud[iLPC], weight))
				iLG_WM.append([iLPC, weight])  # Save pointCloud index and associated lbin weight
			lMapWeights[iLG] = iLG_WM
		"""
		for iLG in range(Nl):
			print('iLG={}  L={} nP={}'.format(iLG, lGrid[iLG], len(lMapWeights[iLG])))
			for i in range(len(lMapWeights[iLG])):
				iPC = lMapWeights[iLG][i][0]
				print('  Lpc={} iPC={} weight={}'.format(lPointCloud[iPC], iPC, lMapWeights[iLG][i][1]))
		"""
		#Calc eGrid mapping weights
		eMapFracs = [ [ [] for e in range(Ne)] for pnt in range(Npc)] # Npc x Ne x (nx2)
		for iPC in range(Npc):
			i,j = iPointCloud[iPC]
			vm = vms[i,j]
			ke_center = alamData['ilamc']*vm
			ke_lower = alamData['ilami'][:-1]*vm
			ke_upper = alamData['ilami'][1:]*vm
			ke_width = ke_upper-ke_lower

			e_lower = np.zeros(Ne)
			e_upper = np.zeros(Ne)
			for iEG in range(Ne):
				e_lower[iEG] = eGrid[0] if iEG == 0 else (eGrid[iEG-1] + eGrid[iEG])/2
				e_upper[iEG] = eGrid[-1] if iEG == Ne-1 else (eGrid[iEG+1] + eGrid[iEG])/2

			eMapFracs[iPC] = scutils.getWeights_ConsArea(ke_center, ke_lower, ke_upper, eGrid, e_lower, e_upper)

			# Iterate over eGrid, grab weights for all overlapping k bins
			"""
			for iEG in range(Ne):
				e_lower = eGrid[0] if iEG == 0 else (eGrid[iEG-1] + eGrid[iEG])/2
				e_upper = eGrid[-1] if iEG == Ne-1 else (eGrid[iEG+1] + eGrid[iEG])/2

				frac_arr = []
				#Loop through k values, find k bind in 1D collision with current eGrid element
				for k in range(Nk):
					#Does it collide
					if e_lower < ke_upper[k] and e_upper > ke_lower[k]:
						#Get overlap bounds
						o_lower = ke_lower[k] if ke_lower[k] > e_lower else e_lower
						o_upper = ke_upper[k] if ke_upper[k] < e_upper else e_upper
						o_width = o_upper-o_lower
						#frac_arr.append([k, eetas[k,i,j]*o_width/ke_width])  # Fraction of k bin that coincides with eGrid bin
						frac_arr.append([k, o_width/ke_width[k]])  # Fraction of k bin that coincides with eGrid bin
				#print('n: {}  arr:{}'.format(len(frac_arr), frac_arr))
				eMapFracs[iPC][iEG] = frac_arr
			"""
		"""
		for iPC in range(Npc):
			print('iPC={}  L={}'.format(iPC, lPointCloud[iPC]))
			i,j = iPointCloud[iPC]
			vm = vms[i,j]
			for e in range(Ne):
				print('  e={}'.format(eGrid[e]))
				for i in range(len(eMapFracs[iPC][e])):
					k = eMapFracs[iPC][e][i][0]
					print('    k={}  E_k={} frac={}'.format(k, alamData['ilamc'][k]*vm,eMapFracs[iPC][e][i][1]))
		"""

		#Calculate the desired variables for all of our points
		odf_pc = np.zeros((Npc,Nk))
		press_pc = np.zeros((Npc,Nk))
		for npc in range(Npc):
			# Single point's values
			i,j = iPointCloud[npc]
			vm = vms[i,j]
			eeta_k = eetas[:,i,j]

			odf_pc[npc,:] =  sf_factor*alamData['ilamc']*vm*eeta_k/alamData['lamscl'] # [Nk]
			press_pc[npc,:] =  pressure_factor*alamData['ilamc']*eeta_k*vm**2.5 * 1E9  # [Pa -> nPa]
			#print('nPC={} vm={} odf={} press={}'.format(npc,vm,odf_pc[npc,:],press_pc[npc,:]))
			
		#Now we have all necessary info. Map all points to their corresponding L-E grid point
		odf_EL = np.zeros((Ne,Nl))
		press_EL = np.zeros((Ne,Nl))
		for iLG in range(Nl):
			#print('iLG=' + str(iLG))
			pointsToMap = lMapWeights[iLG]  # [nx2], nx[pc index, weight]
			numPoints = len(pointsToMap)

			for i in range(numPoints):
				iPC = pointsToMap[i][0]
				lWeight = pointsToMap[i][1]
				
				for e in range(Ne):
					for ik in range(len(eMapFracs[iPC][e])):
						k = eMapFracs[iPC][e][ik][0]
						eFrac = eMapFracs[iPC][e][ik][1]
						odf_EL[e,iLG] += lWeight/numPoints * eFrac * odf_pc[iPC,k]
						press_EL[e,iLG] += lWeight/numPoints * eFrac * press_pc[iPC,k]

		#Yay we made it
		odf_tkl[n,:,:] = odf_EL
		press_tkl[n,:,:] = press_EL

	#Destroy our hard work
	odf_tl = np.ma.sum(np.ma.masked_invalid(odf_tkl), axis=1)
	press_tl = np.ma.sum(np.ma.masked_invalid(press_tkl), axis=1)
		
	result = {}
	result['T'         ] = rcmTimes['T'  ][iTStart:iTEnd+1:sStride]
	result['MJD'       ] = rcmTimes['MJD'][iTStart:iTEnd+1:sStride]
	result['lambda'    ] = alamData['ilamc']
	result['L_bins'    ] = lGrid
	result['energyGrid'] = eGrid
	result['odf_tkl'   ] = odf_tkl
	result['odf_tl'    ] = odf_tl
	result['press_tkl' ] = press_tkl
	result['press_tl'  ] = press_tl

	if dojson: kj.dump(rcmjfname, result)

	return result

#TODO: Take list of variable strings to pull from rcm.h5 file
def getRCM_eqlatlon(mhdrcmf5, rcmTimes, sStart, sEnd, sStride, jdir=None, forceCalc=False):
	"""Grab certain variables along with equatorial and lat-lon grid.

	Args:
		rcmf5 (str): RCM HDF5 filename.
		rcmTimes (dict): Dictionary returned from getRCMTimes().
		evMJDs (float): MJD to evaluate cumulative fractions.
		evalLatLons (list): [lat, lon] location to perform calculation.
		species (str, optional): Species to calculate cumulative fractions for. Defaults to 'ions'.
		jdir (str, optional): Directory to find JSON files in. Defaults to None.
		forceCalc (bool, optional): Force calculation even if JSON file exists. Defaults to False.

	Returns:
		dict: A dictionary containing the calculated results.

	Raises:
		None
	"""
	if jdir is not None:
		dojson = True
		rcmjfname = os.path.join(jdir, RCM_EQLATLON_JFNAME)
	else:
		dojson = False

	if dojson and not forceCalc:
		if os.path.exists(rcmjfname):
			print("Grabbing RCM eq_lat-lon data from " + rcmjfname)
			return kj.load(rcmjfname)

	Nt = rcmTimes['Nt']
	iTStart = np.abs(rcmTimes['sIDs']-sStart).argmin()
	iTEnd = np.abs(rcmTimes['sIDs']-sEnd).argmin()
	sIDs = rcmTimes['sIDs'][iTStart:iTEnd+1:sStride]
	sIDstrs = rcmTimes['sIDstrs'][iTStart:iTEnd+1:sStride]

	mhdrcm5 = h5.File(mhdrcmf5,'r')
	
	mhdrcmS0 = mhdrcm5[sIDstrs[0]]
	Ni,Nj = mhdrcmS0['aloct'].shape

	mlat_rcm = 90.0-mhdrcmS0['colat'][0,:]*180/np.pi
	mlon_rcm = mhdrcmS0['aloct'][:,0]*180/np.pi
	mlatrcm_min = np.min(mlat_rcm)
	mlatrcm_max = np.max(mlat_rcm)

	#Get desired variables for all timesteps
	xmin_arr = np.ma.zeros((Nt, Ni, Nj))
	ymin_arr = np.ma.zeros((Nt, Ni, Nj))
	press_arr = np.ma.zeros((Nt, Ni, Nj))
	dens_arr = np.zeros((Nt, Ni, Nj))
	
	scLoc_eq = np.zeros((Nt, 2))
	scLoc_latlon = np.zeros((Nt, 2))

	rMin = 1.25
	rMax = 35.0
	ioCut = -0.5
	pCut = 1E-8
	
	print("Grabbing data...")
	"""
	for t in range(len(sIDstrs)):
		xm = mhdrcm5[sIDstrs[t]]['xMin'][:]
		ym = mhdrcm5[sIDstrs[t]]['yMin'][:]
		bmR = np.sqrt(xm*xm + ym*ym)
		pm = mhdrcm5[sIDstrs[t]]['P'][:]
		#Turn IOpen into a big true/false map
		iopen_t = mhdrcm5[sIDstrs[t]]['IOpen'][:] 

		Ir = (bmR < rMin) | (bmR > rMax)
		I_m = Ir | (iopen_t > ioCut) | (pm < pCut)
	"""
	#import kaipy.gamera.msphViz as mviz
	rcmdata = gampp.GameraPipe('',mhdrcmf5.split('.h5')[0])
	for t in range(0, len(sIDs)):
		bmX, bmY = rcmpp.RCMEq(rcmdata, sIDs[t], doMask=True)
		I = rcmpp.GetMask(rcmdata, sIDs[t])
		pm = rcmpp.GetVarMask(rcmdata, sIDs[t], 'P', I)

		xmin_arr[t,:,:] = np.transpose(bmX)
		ymin_arr[t,:,:] = np.transpose(bmY)
		press_arr[t,:,:] = np.transpose(pm)

	"""
	for t in range(tStart, len(sIDstrs)):
		#Linterp sc location based on time
		rcmTime = rcm_ut[t]
		if ut[0] > rcmTime or ut[-1] < rcmTime:
			scLoc_eq[t,:] = [0,0]
			scLoc_latlon[t,:] = [0,0]
		else:
			itime = 0
			while ut[itime] < rcmTime: itime += 1

			scLoc_eq[t,:] = [xeq_sc[itime], yeq_sc[itime]]
			scLoc_latlon[t,:] = [mlat_sc[itime], mlon_sc[itime]]
	"""
	result = {}
	result['T']     = rcmTimes['T'][iTStart:iTEnd+1:sStride]
	result['MJD']   = rcmTimes['MJD'][iTStart:iTEnd+1:sStride]
	result['MLAT']  = mlat_rcm
	result['MLON']  = mlon_rcm
	result['xmin']  = xmin_arr
	result['ymin']  = ymin_arr
	result['press'] = press_arr

	if dojson: kj.dump(rcmjfname, result)

	return result

def getRCM_CumulFrac(rcmf5, rcmTimes, evMJD, evalLatLon, species='ions', jdir=None,forceCalc=False):
	"""Calculate cumulative fractions of certain vars (ex. pressure) over energy channels, for each list of points for each given MJD.

	Args:
		rcmf5 (str): RCM HDF5 filename.
		rcmTimes (dict): Dictionary returned from getRCMTimes().
		evMJDs (list): List of MJDs to evaluate cumulative fractions.
		evalLatLons (list): List of [lat, lon] locations to perform calculation.

	Returns:
		dict: Result data structure containing the following keys:
			- 'MJD': MJD value.
			- 'i': i value.
			- 'j': j value.
			- 'lat': Latitude value.
			- 'lon': Longitude value.
			- 'xmin': xmin value.
			- 'ymin': ymin value.
			- 'vm': vm value.
			- 'energies': List of rcm energies [eV].
			- 'Ptot': Total pressure.
			- 'Ppar': List of each channel's individual pressure.
			- 'Pcum': List of each channel's cumulative pressure (summing from k=0).
	"""

	if jdir is not None:
		dojson = True
		rcmjfname = os.path.join(jdir, RCM_CUMULFRAC_JFNAME)
	else:
		dojson = False

	if dojson and not forceCalc:
		if os.path.exists(rcmjfname):
			print("Looking for RCM cumulative fraction data from " + rcmjfname)
			#In this case, saved variable is a list with all points saved so far in one big heap
			data = kj.load(rcmjfname)
			for point in data:
				if point['MJD'] == evMJD and point['lat'] == evalLatLon[0] and point['lon'] == evalLatLon[1]:
					return point

			print("Point not found in previous data, calculating")

	#Unpack RCM info
	sIDs = rcmTimes['sIDs']
	sIDstrs = rcmTimes['sIDstrs']
	rcmMJDs = rcmTimes['MJD']
	rcm5 = h5.File(rcmf5,'r')
	rcmS0 = rcm5[sIDstrs[0]]

	rcmLats = 90 - rcmS0['colat'][0,:]*180/np.pi
	rcmLons = rcmS0['aloct'][2:,0]*180/np.pi

	evLat = evalLatLon[0]
	evLon = evalLatLon[1]

	alamData = getSpecieslambdata(rcmS0, species)
	Nk = len(alamData['ilamc'])

	#Variable we're gonna return
	result = {}

	#TODO: Currently being lazy and getting the closest RCM point. Should implement TSC
	#Get RCM data closest to desired MJD, lat and lon
	ircmStep = np.abs(rcmMJDs-evMJD).argmin()
	i = np.abs(rcmLats-evLat).argmin()
	j = np.abs(rcmLons-evLon).argmin()
	stepStr = sIDstrs[ircmStep]
	rcmS5 = rcm5[stepStr]

	point = {}
	point['MJD'] = evMJD
	point['i'] = i
	point['j'] = j
	point['lat'] = rcmLats[i]
	point['lon'] = rcmLons[j]
	point['xmin'] = rcmS5['rcmxmin'][j,i]
	point['ymin'] = rcmS5['rcmymin'][j,i]
	vm = rcmS5['rcmvm'][j,i]
	point['vm'] = rcmS5['rcmvm'][j,i]
	point['energies'] = alamData['ilamc']*vm  # [eV]

	eetas = rcmS5['rcmeeta'][alamData['kStart']:alamData['kEnd'],j,i]

	pTot = 0
	pPar = np.zeros(Nk)
	pCum = np.zeros(Nk)

	pPar = pressure_factor*alamData['ilamc']*eetas*vm**2.5 * 1E9  # [Pa -> nPa], partial pressures for each channel
	pCum[0] = pPar[0] # First bin's partial press = cumulative pressure
	for k in range(1,Nk):
		pCum[k] = pCum[k-1] + pPar[k]  # Get current cumulative pressure by adding this bin's partial onto last bin's cumulative
	pTot = pCum[-1] # Last bin's cumulative = total pressure

	point['Ptot'] = pTot
	point['Ppar'] = pPar
	point['Pcum'] = pCum

	if dojson: 
		if os.path.exists(rcmjfname):
			data = kj.load(rcmjfname)
			for i in len(range(data)):
				#Remove old point if it exists
				pnt = data[i]
				if pnt['MJD'] == evMJD and pnt['lat'] == evalLatLon[0] and pnt['lon'] == evalLatLon[1]:
					data.remove(i)
				data.append(point)

			kj.dump(rcmjfname, data)
		else:
			data = [point]
			kj.dump(rcmjfname, data)

	return point
#======
#Plotting
#======

def plt_ODF_Comp(AxSC, AxRCM, AxCB, odfData, mjd=None, cmapName='CMRmap', norm=None, forcePop=False):
	"""
	Plot the ODF (Orientation Distribution Function) comparison between AxSC and AxRCM.

	Args:
		AxSC (matplotlib.axes.Axes): The subplot for the SC (Spacecraft) ODF.
		AxRCM (matplotlib.axes.Axes): The subplot for the RCM (Remote Component Manipulator) ODF.
		AxCB (matplotlib.axes.Axes): The colorbar axis.
		odfData (dict): A dictionary containing the ODF data.
		mjd (float): The Modified Julian Date (MJD) to highlight on the plot (optional).
		cmapName (str): The name of the colormap to use (default: 'CMRmap').
		norm (matplotlib.colors.Normalize): The normalization object for the colorbar (optional).
		forcePop (bool): A flag indicating whether to force population of the subplots (default: False).

	Returns:
		None
	"""
	axIsPopulated = not AxSC.get_ylabel() == ''

	eGrid   = odfData['energyGrid']
	scTime  = odfData['sc']['time']
	scODF   = odfData['sc']['diffFlux']
	rcmTime = odfData['rcm']['time']
	rcmODF  = odfData['rcm']['diffFlux']

	#ut = scutils.mjd_to_ut(rcmTime)
	ut = kT.MJD2UT(rcmTime)

	if norm is None:
		vMax = np.max([scODF.max(), rcmODF.max()])
		vMin = np.max([scODF[scODF>0].min(), rcmODF[rcmODF>0].min()])
		norm = kv.genNorm(vMin,vMax,doLog=True)

	if not axIsPopulated or forcePop:
		kv.genCB(AxCB,norm,r'Intensity [$cm^{-2} sr^{-1} s^{-1} keV^{-1}$]',cM=cmapName,doVert=True,cbSz=14)

		AxSC.pcolormesh(scTime, eGrid, np.transpose(scODF), norm=norm, shading='nearest', cmap=cmapName)

		AxSC.set_xlim([ut[0], ut[-1]])
		AxSC.set_ylabel("%s Energy [keV]"%odfData['sc']['name'])
		AxSC.set_yscale('log')
		#AxSC.xaxis.set_major_formatter(plt.NullFormatter())

		AxRCM.pcolormesh(ut, eGrid, np.transpose(rcmODF), norm=norm, shading='nearest', cmap=cmapName)
		#for n in range(len(rcmTime)):
		#    AxRCM.plot(odfData['rcm']['origEGrid'][n,:], odfData['rcm']['origODF'][n,:])
		AxRCM.set_ylabel("RCM Energy [keV]")
		AxRCM.set_yscale('log')

		#xlblFmt = mdates.DateFormatter('%H:%M')
		#AxRCM.xaxis.set_major_formatter(xlblFmt)

	if mjd is not None:
		if mjd < rcmTime[0] or mjd > rcmTime[-1]:
			print(str(mjd) + "not in rcm data, exiting")
			return
		iMJD = np.abs(rcmTime - mjd).argmin()
		lineUT = ut[iMJD]

		if len(AxSC.lines) != 0:
			AxSC.lines.pop(0)  #Remove previous mjd line
			AxSC.lines.pop(0)  #Remove previous mjd line
			AxRCM.lines.pop(0)
			AxRCM.lines.pop(0)
		yMin, yMax = AxSC.get_ylim()
		AxSC.plot([lineUT, lineUT], [yMin, yMax], '-k',linewidth=2)
		AxSC.plot([lineUT, lineUT], [yMin, yMax], '-w',linewidth=1)
		yMin, yMax = AxRCM.get_ylim()
		AxRCM.plot([lineUT, lineUT], [yMin, yMax], '-k',linewidth=2)
		AxRCM.plot([lineUT, lineUT], [yMin, yMax], '-w',linewidth=1)
		

def plt_tl(AxTL, tkldata, AxCB=None, mjd=None,cmapName='CMRmap',norm=None):
	"""
	Plot the total pressure as a function of L shell and time.

	Args:
		AxTL: The matplotlib Axes object for the L vs. Time plot.
		tkldata: A dictionary containing the data for the plot.
		AxCB: The matplotlib Axes object for the colorbar plot (optional).
		mjd: The Modified Julian Date for a specific time point (optional).
		cmapName: The name of the colormap to use (default: 'CMRmap').
		norm: The normalization object for the color mapping (optional).

	Returns:
		None
	"""

	L_arr = tkldata['L_bins']
	ut = kT.MJD2UT(tkldata['MJD'])
	press_tl = np.array(tkldata['press_tl'][:], dtype=float)  # Need to do this to handle masked stuff
	press_tl = np.ma.masked_invalid(press_tl)

	#Initialize static plots if hasn't been done yet
	doPopulateTL =  AxTL.get_ylabel() == ''
	doPopulateCB = AxCB is not None and AxCB.get_label() == ''

	if norm is None: 
		norm = genVarNorm(press_tl, doLog=True)

	if doPopulateTL:
		#L vs. Time
		AxTL.pcolormesh(ut, L_arr, np.transpose(press_tl), norm=norm, shading='nearest', cmap=cmapName)
		AxTL.set_xlim([ut[0], ut[-1]])
		AxTL.set_ylabel('L shell')
		AxTL.set_xlabel('UT')
		#Sim boundaries
		xmin,xmax = AxTL.get_xlim()
		AxTL.plot([xmin,xmax],[-tkl_lInner,-tkl_lInner], 'w--')
		AxTL.plot([xmin,xmax],[tkl_lInner,tkl_lInner], 'w--')
	if doPopulateCB:
		kv.genCB(AxCB, norm, r'Total pressure [$nPa$]', cM=cmapName, doVert=False)

	#Time-specific stuff
	if mjd is not None:
		if mjd < tkldata['MJD'][0] or mjd > tkldata['MJD'][-1]:
			return
		iMJD = np.abs(tkldata['MJD'] - mjd).argmin()
		lineUT = ut[iMJD]

		if len(AxTL.lines) > 2:
			AxTL.lines.pop(2)  #Remove previous mjd line
		yMin, yMax = AxTL.get_ylim()
		AxTL.plot([lineUT, lineUT], [yMin, yMax], '-k')


def plt_rcm_eqlatlon(AxLatlon, AxEq, rcmData, satTrackData=None, AxCB=None, mjd=None, norm=None, cmapName='viridis'):
	"""
	Plot RCM data on equal latitude-longitude and equatorial coordinate systems.

	Args:
		AxLatlon (matplotlib.axes.Axes): The axis object for the equal latitude-longitude plot.
		AxEq (matplotlib.axes.Axes): The axis object for the equatorial coordinate plot.
		rcmData (dict): A dictionary containing RCM data.
		satTrackData (dict, optional): A dictionary containing satellite track data. Default is None.
		AxCB (matplotlib.axes.Axes, optional): The colorbar axis object. Default is None.
		mjd (float, optional): The Modified Julian Date. Default is None.
		norm (matplotlib.colors.Normalize, optional): The normalization object for the color mapping. Default is None.
		cmapName (str, optional): The name of the colormap. Default is 'viridis'.

	Returns:
		None
	"""

	mjd_arr   = rcmData['MJD']
	xmin_arr  = rcmData['xmin']
	ymin_arr  = rcmData['ymin']
	mlat_arr  = rcmData['MLAT']
	mlon_arr  = rcmData['MLON']
	press_arr = rcmData['press']

	#ut = scutils.mjd_to_ut(rcmData['MJD'])
	ut = kT.MJD2UT(rcmData['MJD'])
	Nt = len(ut)

	if norm is None: 
		norm = genVarNorm(press_arr, doLog=True)

	#Initialize static plots if hasn't been done yet
	if AxCB is not None:
		AxCB = kv.genCB(AxCB, norm, r'Pressure [$nPa$]', cM=cmapName, doVert=False)

	if mjd is not None:
		if mjd < mjd_arr[0] or mjd > mjd_arr[-1]:
			print(str(mjd) + "not in rcm data, exiting")
			return
		iMJD = np.abs(mjd_arr - mjd).argmin()
		lineUT = ut[iMJD]

		AxLatlon.clear()
		AxEq.clear()

		#Prep rcm lat/lons for polar plotting
		riono = np.cos(mlat_arr*np.pi/180.)
		tiono = np.concatenate((mlon_arr, [mlon_arr[0]]))*np.pi/180.
		AxLatlon.pcolor(tiono, riono, to_center(np.transpose(press_arr[iMJD])),norm=norm, shading='auto', cmap=cmapName)
		AxLatlon.axis([0, 2*np.pi, 0, 0.7])

		AxEq.pcolor(xmin_arr[iMJD], ymin_arr[iMJD], to_center(press_arr[iMJD]), norm=norm, shading='auto', cmap=cmapName)

		#Draw satellite location
		if satTrackData is not None:
			x_sc = satTrackData['xeq']
			y_sc = satTrackData['yeq']
			req_sc = satTrackData['Req']
			iscMJD = np.abs(satTrackData['MJD'] - mjd).argmin()
			if req_sc[iscMJD] > 1E-8:
				leadMax = iscMJD
				while leadMax < min(iscMJD+80, Nt) and req_sc[leadMax] > 1E-8: leadMax += 1 #!!This isn't working as intended for some reason
				AxEq.plot(x_sc[iscMJD:leadMax], y_sc[iscMJD:leadMax], 'k-')
				
				satCircle = plt.Circle((x_sc[iscMJD], y_sc[iscMJD]), 0.15, color='black')
				AxEq.add_patch(satCircle)
		kv.addEarth2D(ax=AxEq)

		#Set bounds
		boxRatio = get_aspect(AxEq)
		xbMin = -10
		xbMax = 10
		dx = xbMax - xbMin
		dy = boxRatio*dx
		ybMin = -dy/2
		ybMax = dy/2
		AxEq.set_xlim([xbMin, xbMax])
		AxEq.set_ylim([ybMin, ybMax])







