#!/usr/bin/python
# The "new2" program for approach 2.  
# Creates csv annual sums from temperature adjusted and unadjusted onroad netCDFs.
# 10/21/08 James Beidler <beidler.james@epa.gov>
# Updated 4/16/09
# Update 7/13/09 by C. Allen - changed definition of dayPath so that this script works for different years
# Updated 12/23/09 by J. Beidler - Fixed input paths to work with unadjusted summaries.
# Updated 1/8/10 by J. Beidler - Changed parameter file loading.  Added calendar based dates.

from Numeric import *
from Scientific.IO.NetCDF import *
import sys, os, csv, calendar

def checkEV(evName):
	"""
	Checks if an environment variable is set.  If not, exits.  If it is, returns the variable.
	Takes the name of the environment variable.
	"""
	try: var = os.environ[evName]
	except:
		print "ERROR: Environment variable '%s' is not defined." %evName
		sys.exit(1)
	else: return var

baseYear = checkEV('BASE_YEAR') # Modeling year, used in definition of dayPath

# Path and filename of molecular weight conversion files
mwDict = {'cmaq_cb05_soa': '/garnet/oaqps/smoke/smoke2.7/scripts/annual_report/parameter_file_cmaq_cb05.txt', 'cmaq_cb05': '/garnet/oaqps/smoke/smoke2.7/scripts/annual_report/parameter_file_cmaq_cb05.txt', 'cmaq_cb05_tx': '/garnet/oaqps/smoke/smoke2.7/scripts/annual_report/parameter_file_cmaq_cb05.txt', 'saprc07t': '/garnet/work/bte/pyQA/parameter_file_saprc07t.txt'} 
mwDefault = 'cmaq_cb05'  #  Default molecular weight file 
# End User Config

def readDAYTABLE(dayPath, Year, Mon):
	"""
	Reads in the mrggrid date file for day conversion.
	"""
	dayFileName = 'smk_merge_dates_' + Year + Mon + '.txt'
	dayFileName = os.path.join(dayPath, dayFileName)
	dayTable = []
	for line in csv.reader(openFile(dayFileName).readlines()[1:]):
		dayTable.append(line)
	return dayTable 

def openFile(fileName, accessType = 'r'):
	"""
	Tests if a file is available for access.  If not it returns an error and exits.
	If it is available it returns an open file object.
	"""
	try: file = open(fileName, accessType)
	except:
		print "ERROR: %s not available for access." %fileName
		sys.exit(1)
	else: return file

def nameOutFile():
	"""
	Set the outfile name.
	"""
	if runType == "ADJ": 
		outFileName = sector + '_adj_ann_sum_' + grid + '_' + case + '.csv'
		outPath = os.path.join(imPath, sector + '_adj')
	else: 
		outFileName = sector + '_ann_sum_' + grid + '_' + case + '.csv'
		outPath = os.path.join(imPath, sector)
	return os.path.join(outPath, outFileName)

def fixDateLen(inDate):
	"""
	Return a numerical date as a two digit string
	"""
	if len(str(inDate)) == 1: inDate = '0' + str(inDate)
	else: inDate = str(inDate)
	return inDate

class createReport(object):
	"""
	"""

	def __init__(self, Mon = ""):
		self.Mon = Mon

	def nameInFile(self, inDate):
		"""
		Set the infile name based on the SMOKE conventions.
		""" 
		if runType == "ADJ": 
			inFileName = 'emis_mole_' + sector + '_adj_%s_' %inDate + grid + '_' + spec + '_' + case + '.ncf'
			inPath = os.path.join(imPath, sector + '_adj')
		else: 
			inFileName = 'emis_mole_' + sector + '_%s_' %inDate + grid + '_' + spec + '_' + case + '.ncf'
			inPath = os.path.join(imPath, sector)
		return os.path.join(inPath, inFileName)

	def moleToMassFactor(self, speciesName):
		"""
		Get the moles to mass conversion factor for the specified species.
		"""
		if speciesName == 'NAPHTH_72': speciesName = 'NAPHTHALENE'
		if '_72' in speciesName:
			tmp = speciesName.split('_72')
			speciesName = tmp[0]
		if speciesName in molecDct: factor = molecDct[speciesName]
		else: factor = 1
		return factor

	def sumSpeciesDay(self, species):
		"""
		Sum a species over a day
		"""
		dataIn = species.getValue()
		speciesDaySum = zeros([species.shape[2], species.shape[3]], 'f')
		for hour in range(species.shape[0] - 1):
			for row in range(species.shape[2]):
				for col in range(species.shape[3]):
					speciesDaySum[row][col] = speciesDaySum[row][col] + dataIn[hour][0][row][col]
		return speciesDaySum

	def convValue(self, monSum, speciesName):
		"""
		Convert monthly value from moles per second to tons per hour.
		"""
		monSum = monSum * self.moleToMassFactor(speciesName)   # Convert moles per second to grams per second
		monSum = monSum * (0.00000110231131)  # Convert grams per second to tons per second
		monSum = monSum * 3600   # Convert tons per second to tons per hour
		return monSum  

	def outputSpecies(self, vNames, species, speciesMonSum):
		"""
		Output the species to the outfile.
		"""
		for speciesNum, speciesName in enumerate(vNames):
			if speciesName == 'TFLAG': continue
			for row in range(species.shape[2]):
				for col in range(species.shape[3]):
					outFile.write(self.Mon + "," + str(col + 1) + "," + str(row + 1) + "," + speciesName + "," + str(self.convValue(speciesMonSum[speciesNum][row][col], speciesName)) + '\n')

	def sumSpeciesMonth(self):
		"""
		Sum a species over a month.
		"""
		mon = int(self.Mon) 

		for day in range(1, calendar.monthrange(int(Year), mon)[1] + 1):
			Day = fixDateLen(day)

			# Set in date and out date and open the in and out files.
			if runType == "UNADJ": inDate = dayTable[int(Day) - 1][6].strip()
			else: inDate = Year+self.Mon+Day

			# Open the infile for reading
			inFileName = self.nameInFile(inDate)
			print "In File: " + inFileName
			inFile = NetCDFFile(inFileName, 'r')

			# Get the list of variables from the in file
			variableNames = inFile.variables.keys()
			vNames = [species for species in variableNames]

			for speciesNum, speciesName in enumerate(variableNames):
				# Skip TFLAG 
				if speciesName == 'TFLAG': continue

				species = inFile.variables[speciesName]
				# Sum a species over a day
				speciesDaySum = self.sumSpeciesDay(species)

				# Calculate and populate monthly sums
				if (Day == '01'):
					speciesMonSum = variableNames
					speciesMonSum[speciesNum] = speciesDaySum
				elif (day == calendar.monthrange(int(Year), mon)[1]):
					speciesMonSum[speciesNum] = speciesMonSum[speciesNum] + speciesDaySum
				else:
					speciesMonSum[speciesNum] = speciesMonSum[speciesNum] + speciesDaySum
		self.outputSpecies(vNames, species, speciesMonSum)	


Year = checkEV('BASE_YEAR')
grid = checkEV('GRID')
sector = checkEV('SECTOR')
spec = checkEV('EMF_SPC')
case = checkEV('CASE')
imPath = checkEV('IMD_ROOT')
runType = checkEV('RUN_SECT')   # ADJ for adjusted or UNADJ for unadjusted
dayPath = checkEV('MRGDATE_ROOT')

monList = checkEV('MONTHS_LIST') # Months to run as reported by the set_months script
monList = [ int(mon.strip()) for mon in monList.split(' ') ]
monList.sort

# Find the species used in this case and load the molecular weight conversion dictionary
if spec in mwDict:
	mwSpec = spec.lower()
else:
	mwSpec = mwDefault
execfile(mwDict[mwSpec])  # Load up the molecular weight conversion dictionary from the external file

outFileName = nameOutFile()
print "Out file: " + outFileName
# Open the outfile for writing
outFile = file(outFileName, 'w')
outFile.write( '# Summed %s %s merged reports for months %s-%s\n' %(sector, runType, monList[0], monList[ len(monList) - 1 ]) )
outFile.write('# Month, Col, Row, Species Name, Value\n')
 
# Main loop
for mon in monList:
	Mon = fixDateLen(mon)
	print "Processing month: " + Mon
	dayTable = readDAYTABLE(dayPath, Year, Mon)
	reportMonth = createReport(Mon) 
	reportMonth.sumSpeciesMonth()

print 'Outputting to ' + outFileName
outFile.close()
