# read and plot SGPS Sept. 2017 event data
# B. Kress, Jan. 2021

# standard packages
import sys
import numpy as np
from netCDF4 import Dataset as NCDataset
#from mpl_toolkits.axes_grid1 import Divider, Size
#from mpl_toolkits.axes_grid1.mpl_axes import Axes
import datetime as dt
import dateutil.parser as dtp

import matplotlib.pyplot as pyplot
import matplotlib.gridspec as gridspec
import matplotlib
matplotlib.rcParams.update({'font.size': 10})
import matplotlib.cm as cm

# constants
UNITS = 2
CHANNELS = 14
DIFF_CHANNELS = 13
RECORDS = 8640

# input files
filenamelist = []
filenamelist.append('se_sgps-l2-avg5m_g16_s20172440000000_e20172732355000_v2_0_0.nc')
NUM_INPUT_FILES = len(filenamelist)

#dimensions:
#        record_number = UNLIMITED ; // (8640 currently)
#        diff_channels = 13 ;
#        sensor_units = 2 ;
#variables:
#        double L2_SciData_TimeStamp(record_number) ;
#        float AvgDiffProtonFlux(record_number, sensor_units, diff_channels) ;
#        float AvgIntProtonFlux(record_number, sensor_units) ;

# arrays for timestamps and fluxes
TimeStamp = np.zeros([RECORDS], dtype=np.float64)
ProtonFluxes = np.zeros([RECORDS, UNITS, CHANNELS], dtype=np.float32)

# open file and read data set
filenum = 0
print 'reading ' + filenamelist[filenum]
ds = NCDataset(filenamelist[filenum], mode='r' )
num_records = len(ds.variables['L2_SciData_TimeStamp'][:])
TimeStamp[0:num_records] = ds.variables['L2_SciData_TimeStamp'][:]
ProtonFluxes[0:num_records,:,0:DIFF_CHANNELS] = ds.variables['AvgDiffProtonFlux'][:,:,:]
ProtonFluxes[0:num_records,:,DIFF_CHANNELS] = ds.variables['AvgIntProtonFlux'][:,:]
ds.close()

# plotting
plot=1
if (plot==1):

	# colormap
	numcolors = 13
	x = np.arange(numcolors)
	rgb = pyplot.get_cmap('jet',numcolors)(x)[:,:3]
	# adjust end colors a little lighter
	rgb[0,2] = .6
	rgb[numcolors-1,0] = .6
	chan_color = np.flip(rgb, axis=0) 

	# nominal channel energies
	chan_label = ['P1 (1.0-1.9 MeV)','P2A (1.9-2.3 MeV)','P2B (2.3-3.4 MeV)','P3 (3.4-6.5 MeV)','P4 (6.5-12 MeV)',\
	    'P5 (12-25 MeV)','P6 (25-40 MeV)','P7 (40-80 MeV)','P8A (83-99 MeV)','P8B (99-118 MeV)','P8C (118-150 MeV)',\
	    'P9 (150-275 MeV)','P10 (275-500 MeV)','P11 (>500 MeV)']

	# set tic locations and labels
	first_j2000_sec = TimeStamp[0]
	last_j2000_sec = TimeStamp[len(TimeStamp)-1]
	xmin = first_j2000_sec
	xmax = last_j2000_sec
	NumInFiles = 30
	monthstr = '9/'
	firstday = 1
	emptystr = []
	ticloc = []
	ticstr = []
	ticloc.append(first_j2000_sec)
	ticstr.append(monthstr + str(firstday))
	for i in xrange(1,NumInFiles):
		ticloc.append(ticloc[i-1]+86400)
		ticstr.append(monthstr + str(firstday+i))

	# plot
	pyplot.figure(1, figsize=[16,6])

	# plot panel 2
	gridspec.GridSpec(5,1)

	ax2 = pyplot.subplot2grid((5,1), (2,0), colspan=1, rowspan=2)
	k=1
	for j in xrange(0,13):
		pyplot.plot(TimeStamp[:], ProtonFluxes[:,k,j], color=chan_color[j,:], label=chan_label[j])
	pyplot.yscale('log')
	pyplot.xlabel('2017')
	pyplot.ylabel('protons/cm$^2$-s-sr-keV')
	pyplot.xlim([xmin, xmax])
	pyplot.ylim([6.e-8, 100.])
	pyplot.xticks(ticloc,ticstr)
	# no x-tic labels 
	labels = [item.get_text() for item in ax2.get_xticklabels()]
	empty_string_labels = ['']*len(labels)
	ax2.set_xticklabels(empty_string_labels)
	textstr = 'SGPS+X (east)'
	ax2.text(0.05, 0.97, textstr, transform=ax2.transAxes, fontsize=14, verticalalignment='top')

	# plot panel 1
	ax1 = pyplot.subplot2grid((5,1), (0,0), colspan=1, rowspan=2)
	k=0
	for j in xrange(13):
		pyplot.plot(TimeStamp[:], ProtonFluxes[:,k,j], color=chan_color[j,:], label=chan_label[j])
	pyplot.yscale('log')
	pyplot.ylabel('protons/cm$^2$-s-sr-keV')
	pyplot.xlim([xmin, xmax])
	pyplot.ylim([6.e-8, 100.])
	pyplot.xticks(ticloc,ticstr)
	# no x-tic labels 
	labels = [item.get_text() for item in ax1.get_xticklabels()]
	empty_string_labels = ['']*len(labels)
	ax1.set_xticklabels(empty_string_labels)
	textstr = 'SGPS-X (west)'
	ax1.text(0.05, 0.97, textstr, transform=ax1.transAxes, fontsize=14, verticalalignment='top')

	pyplot.legend(loc='upper right', prop={'size': 10})

	# plot panel 3
	ax3 = pyplot.subplot2grid((5,1), (4,0), colspan=1, rowspan=1)
	j=13
	k=0
	pyplot.plot(TimeStamp[:], ProtonFluxes[:,k,j], color='k', label='SGPS-X (west) P11 (>500 MeV)')
	k=1
	pyplot.plot(TimeStamp[:], ProtonFluxes[:,k,j], color='gray', label='SGPS+X (east) P11 (>500 MeV)')
	pyplot.yscale('log')
	pyplot.xlabel('2017')
	pyplot.ylabel('protons/cm$^2$-s-sr')
	pyplot.xlim([xmin, xmax])
	pyplot.xticks(ticloc,ticstr)
	textstr = 'SGPS P11'
	ax3.text(0.05, 0.94, textstr, transform=ax3.transAxes, fontsize=14, verticalalignment='top')
	pyplot.legend(loc='upper right', prop={'size': 10})

	# save and/or show figure
	pyplot.savefig('sgps_sept2017.png', bbox_inches='tight')
	#pyplot.savefig('sgps_sept2017.eps', format='eps', bbox_inches='tight', dpi=600)
	pyplot.show()