import sys
import math as M

import autoplot as A

# These are commonly needed, good candidates for exporting into the Autoplot 
# namespace
from org.das2.fsm import FileStorageModel
from org.das2.qds.util.Reduction import reducex

##############################################################################
# Constants

DESC = 0
FIRST_COL = 1
NUM_BINS = 2
DROP_LEFT = 3
FILES = 4
g_sRoot = """file:/opt/project/juno/mirror/jed/JNOJED_3001/DATA"""




















g_dOffset = {'counts':0, 'counts_per_sec':1, 'flux':0}





g_dLabel  = {'counts':'counts',
				'counts_per_sec':'Count Rate (s!a-1!n)',
				'flux':'Differential Intensity (cm!a-2!n keV!a-1!n s!a-1!n sr!a-1!n)'}
				 
perr = sys.stderr.write
 

##############################################################################

def getFileList(sType, lFans, tr):
	"""Get a list of files of homogeneous files to read, the data produced
	from reading all these files will be a qube and may thus be time averaged
	easily.
	
	   sType - one of HIERSTOFXER, HIERSTOFXPHR, LOERSTOFXER, LOERSTOFXPHR
		lFans - A list of JEDI sensors (hockey pucks)
		tr - An Autoplot timerange object
	"""
	
	lFiles = []
	for sFan in lFans:
		sPtrn = '%s/$Y/$j/JED_%s_%s_CDR_$Y$j_V$v.TAB'%(g_sRoot, sFan, sType)
		
		i = FileStorageModel.splitIndex( sPtrn )

		sBase = sPtrn[0:i]  # the static part of the name
		sTplt = sPtrn[i:]   # the templated part of the name

		#fs = FileSystem.create(sBase)

		#fsm = FileStorageModel.create(fs, sTplt)

		# Get best names only returns relative paths from the root...
		mon = monitor.getSubtaskMonitor("Finding %s files"%sType)
		#aNames = fsm.getBestNamesFor(DatumRangeUtil.parseTimeRange(tr), mon)

		# ...so glue them back together
		#lFiles += [ "%s%s"%(sBase, s) for s in aNames]
		
	return lFiles
	
##############################################################################

def getYTags(sFile, iStart, nBins):
	"""Open a TOF x E/PH file and get the energy bin centers"""

	if sFile.startswith('file:'):
		sFile = sFile[5:]
	if sFile.startswith('///'):
		sFile = sFile[2:]

	fIn = file(sFile, 'rb')
	fIn.readline()
	sBeg = fIn.readline()
	sEnd = fIn.readline()
	fIn.close()
	lBeg = sBeg.strip().split(',')
	lEnd = sEnd.strip().split(',')
	lCenters = [None]*nBins
	
	for i in xrange(0, nBins):
		lCenters[i] = M.sqrt( float(lBeg[iStart +i*3]) * float(lEnd[iStart + i*3]))
	
	dsYTags = A.dataset(lCenters, A.Units.lookupUnits("keV"))
	dsYTags.putProperty(QDataSet.SCALE_TYPE,'log' )
	
	return dsYTags


##############################################################################
# main

def main(trQuery, sFans, sTimeBinWidth, sOutType):
		
	nColOffset = g_dOffset[sOutType]
		
	# Need to use up to 8 datasets here, list entries for each set are:
	# 0. Description
	# 1. 1st Data Column
	# 2. Number of Energy Birs
	# 3. Number of Low enery channels to drop
	# 4. File list 
	
	lFans = [s.strip() for s in sFans.split(',')]
	
	if len(lFans) == 0:
		raise ValueError("At least one JEDI fan must be specified")
	elif len(lFans) == 1:
		sFanInfo = "Fan %s"%sFans
	else:
		sFanInfo = "Fans %s"%", ".join(lFans)

		
	llDataSets = []
	
	llDataSets += [
		
		# For this one the lower 4-bins overlap with the upper bins of the 
		# low-energy set, and the calibration for each sensor is different.  To
		# provide a smooth spectrogram the bottom 4 are dropped here
		['High Resolution, High Energy, %s'%sFanInfo, 14, 24,  4,
		  getFileList('HIERSTOFXER', lFans, trQuery) ],
		
		# This one is odd because the first 7 energy levels for each SSD are for 
		# IONS even though this is described as a proton only file
		['High Resolution, Low Energy, %s'%sFanInfo,  14, 16,  8, 
		  getFileList('HIERSTOFXPHR', lFans, trQuery) ],
		
		['Low Resolution, High Energy, %s'%sFanInfo,  14,  6,  0,
		 getFileList('LOERSTOFXER', lFans, trQuery) ],
		
		['Low Resolution, Low Energy, %s'%sFanInfo,   14,  4,  0,
		  getFileList('LOERSTOFXPHR', lFans, trQuery) ]
	]
			
	dsOut = None
	for lSet in llDataSets:          # Over all file series
		
		#perr("INFO: %s files for range: %d\n"%(lSet[DESC], lSet[FILES] ) )
		
		dsFileSet = None
		for sFile in lSet[FILES]:     # Over all files
				
			dsFile = None
			
			nEBins = lSet[NUM_BINS]
			iStartCol = lSet[FIRST_COL]
			iEndCol   = iStartCol + nEBins*3*6  # Get data for all SSDs
				
			sUri = 'vap+dat:%s?depend0=field0&skip=5&rank2=%d:%d'%(sFile, iStartCol, iEndCol)
						
			dsFile = A.getDataSet(sUri)
			dsFile.putProperty(QDataSet.UNITS, g_dUnits[sOutType])
			
			# Take every third item starting from the offset, the offset picks the
			# output type, either counts, counts/sec, or flux
			dsFile = dsFile[:,nColOffset::3]
			
			# Sum over all 6 SSDs
			dsSum = dsFile[:, 0:nEBins]
			for i in xrange(1, 6):
				dsSum = dsSum + dsFile[:, i*nEBins:(i+1)*nEBins ]
			
			# Drop un-wanted low energy channels
			if lSet[DROP_LEFT] > 0:
				dsSum = dsSum[:, lSet[DROP_LEFT] :]
				
			# If we are doing flux output, need to divide by the number of SSDs to
			# keep the units constant
			if sOutType == 'flux':
				dsSum = dsSum / 6.0
			
			# Add the energy bin centers from the file
			dsYTags = getYTags(sFile, lSet[FIRST_COL]+lSet[DROP_LEFT]*3, 
			                          lSet[NUM_BINS]-lSet[DROP_LEFT]  )
			
			sBaseName = sUri.split('?')[0].split('/')[-1]
			perr("INFO: Extracted %2d energy bins and %5d records from %s\n"%(
			     len(dsYTags), dsFile.length(), sBaseName))
			
			dsSum.putProperty(QDataSet.DEPEND_1, dsYTags)
			
			dsSum.putProperty(QDataSet.LABEL, g_dLabel[sOutType])
			
			if dsFileSet == None:
				dsFileSet = dsSum
			else:
				#This...
				#dsFileSet.append(dsSum)
				
				#...is not this
				dsFileSet = A.append(dsFileSet, dsSum)
		
		if dsFileSet != None:
		
			# Sort the resulting data as it may have come from both 090 and 270 sensors
			indices = A.sort(dsFileSet.property(QDataSet.DEPEND_0))
			dsSorted = dsFileSet[indices]
			
			# Reduce in time, if requested
			sAvgInfo = ''
			if sTimeBinWidth != '0 s':
				dsFileSet = reducex(dsFileSet, A.dataset(sTimeBinWidth))
				sAvgInfo = "!c(%s averages)"%sTimeBinWidth
			
			if dsOut == None:
				dsOut = dsFileSet
			else:
				dsOut = join(dsOut, dsFileSet)
				

	# Add a title here
	if dsOut != None:
		if sOutType == 'flux':
			sNote = "Average over all look directions"
		else:
			sNote = "Sum over all look directions"
			
		dsOut.putProperty(QDataSet.TITLE, "JEDI-%s Proton Spectrum - %s%s"%(
		                  ", ".join(lFans), sNote, sAvgInfo))
	
	return dsOut

##############################################################################

trQuery = A.getParam( 'timerange', '2016-181', 'timerange to load' )

sFans = A.getParam( 'fan', 'spin_plane', 'three individual fans, or all spin-plane', 
                   [ '090','180','270','spin_plane' ] )

if sFans == 'spin_plane':
	sFans = """090,270"""
sTimeBinSeconds = A.getParam(
		'timeBinSeconds', '30', 'averaging time in seconds, 0 is intrisinic',
	   [ '600','60', '30', '0' ] 
)

sTimeBinWidth__CLASSTYPE = QDataSet  # (spot line1014 b)








sOutType = A.getParam(
	'outputType','flux',"Data output type", 
	['counts','counts_per_sec','flux']
)

# Autoplot uses magic variables, so you have to set one of them to output the
# dataset, kinda strange