#!/usr/bin/python
# Allocates the gridded reports to annual and monthly county and state reports.
#
# 10/27/2008 James Beidler <beidler.james@epa.gov>
# Updated 4/16/09
# 1/8/2010 J. Beidler - This file replaces onroad_state_county.py and onroad_state_county_byscc.py.  Both reports will be generated from this script. 

import os, csv, sys

# Dictionary that translates digit month string to three character month name 
monDict = {'01': 'jan', '02': 'feb', '03': 'mar', '04': 'apr', '05': 'may', '06': 'jun', '07': 'jul', '08': 'aug', '09': 'sep', '10': 'oct', '11': 'nov', '12': 'dec'}  

def openOutFile(areaType, dateType, bySCC):
	"""
	Set the outfile name.
	"""
	if bySCC == True:
		outFileName = 'rep_SCC7_' + areaType + '_' + dateType + '_adj_' + sector + '_' + case + '_' + grid + '.txt'
		byType = 'emissions by SCC'
	if bySCC == False:
		outFileName = 'rep_' + areaType + '_' + dateType + '_adj_' + sector + '_' + case + '_' + grid + '.txt'
		byType = 'emissions'
	outPath = os.path.join(home, case + '/reports/programs')
	outFileName = os.path.join(outPath, outFileName)
	outFile = openFile(outFileName, 'w')
	outFile.write( '# %s adjusted %s %s %s for months %s-%s\n' %(sector, areaType, dateType, byType, monList[0], monList[ len(monList) - 1 ]) )

	print 'Writing %s %s adjusted %s to: %s' %(areaType, dateType, byType, outFileName)

	if dateType == 'monthly': headerLine = '# Month,FIPS,State Name,'
	else: headerLine = '# FIPS,State Name,'
	if areaType == 'county': headerLine = headerLine + 'County Name,'
	if bySCC == True: headerLine = headerLine + 'SCC7,'
	outFile.write(headerLine + ','.join(speciesList) + '\n')

	return outFile 

def openFile(fileName, accessType = 'r'):
	"""
	Tests to see 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 fixLen(x, y = 2):
	while len(x) < y:
		x = '0' + x
	return x 

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

class createReport(object):
	"""
	Create a class for each month.
	"""
	def __init__(self, Mon = ""):
		self.Mon = Mon
		self.countyDescDict = {}
		self.stateDescDict = {}
		self.rowList = [] 
		self.colList = [] 
		self.speciesList = []

	def nameRepFile(self):
		"""
		Set the infile name for the REP file based on the SMOKE conventions.
		"""
		inFileName = 'rep_' + sector + '_' + case + '_invgrid_cell_county_' + grid + '.txt' 
		inPath = os.path.join(home, case + '/reports/inv')
		return os.path.join(inPath, inFileName)

	def readRepFile(self, repFileName):
		"""
		Read the rep file and put into a list.
		"""
		repDict = {}
		for line in csv.reader(openFile(repFileName, 'r'), delimiter='|'):
			line = [col.strip() for col in line]
			# On the line that starts with #LABEL...
			if '#Label' in line[0]:
				# Create the specieslist starting at the 8th col 
				self.createSpeciesList(line[7:])
				# Create the reference dictionary of the columns
				colDict = dict(zip(line, range(len(line))))

			# If the line is a comment, the skip it 
			if line[0][0] == "#": continue

			col = line[colDict['X cell']]
			row = line[colDict['Y cell']]
			cell = col+','+row
			fips = line[colDict['Region']]
			scc = line[colDict['SCC Tier 3']]
			state = line[colDict['State']]
			county = line[colDict['County']]

			for speciesName in self.speciesList:	
				value = float(line[colDict[speciesName]])
				if fips not in repDict:
					repDict[fips] = { cell: { scc: { speciesName: value } } }
				elif cell not in repDict[fips]: 
					repDict[fips][cell] = { scc: { speciesName: value } } 
				elif scc not in repDict[fips][cell]: 
					repDict[fips][cell][scc] = { speciesName: value } 
				elif speciesName not in repDict[fips][cell][scc]: 
					repDict[fips][cell][scc][speciesName] = value	
					 
			if fips not in self.countyDescDict: 
					self.countyDescDict[fips] = (state, county)
			if fips[1:3] not in self.stateDescDict: 
					self.stateDescDict[fips[1:3]] = state 
		return repDict

	def createSpeciesList(self, speciesLine):
		"""
		Create a list of species from the report file
		"""
		for speciesName in speciesLine:
			if 'S-' in speciesName[:3]: continue
			# Skip PM10 and PM2_5
			if 'PM10' in speciesName: continue
			if 'PM2_5' in speciesName: continue

			self.speciesList.append(speciesName)

	def nameCmSum(self):
		"""
		Set the infile name for the CMSUM file.
		"""
		inFileName = sector + '_adj_ann_sum_' + grid + '_' + case + '.csv'
		inPath = os.path.join(imPath, sector + '_adj')
		return os.path.join(inPath, inFileName)

	def readCmSum(self, cmSumName):
		"""
		Read the CMSUM file and put into a list.
		"""
		cmSpecList = []	
		cmSumDict = {}
		for line in csv.reader(openFile(cmSumName, 'r')):
			if line[0][0] == '#': continue

			month = line[0]
			col = line[1]
			row = line[2]
			cell = col+','+row
			species = line[3]
			value = float(line[4])

			if species not in cmSpecList:
				cmSpecList.append(species)

			if month not in cmSumDict: 
				cmSumDict[month] = { cell: {species: value} }
			if cell not in cmSumDict[month]: 
				cmSumDict[month][cell] = {species: value}
			if species not in cmSumDict[month][cell]: 
				cmSumDict[month][cell][species] = value

			if col not in self.colList: 
				self.colList.append(col)
			if row not in self.rowList: 
				self.rowList.append(row)

		# Species error checking
		for cmSpec in cmSpecList:
			if cmSpec not in self.speciesList:
				print "WARNING: Species %s in CMSUM file but not in inventory report" %cmSpec

		# More species error checking
		for invSpec in self.speciesList:
			if invSpec not in cmSpecList:
				print "WARNING: Species %s in inventory report but not in CMSUM file" %invSpec

		return cmSumDict
	

	def calcCountyFac(self, repDict):
		"""
		Calculate the county allocation factors.
		"""
		crSum = {}
		fipsCRSum = {}
		for fips in self.countyDescDict:
			fipsCRSum[fips] = {} 
			for cell in repDict[fips]:
				if cell not in crSum: 
					crSum[cell] = {} 
				if cell not in fipsCRSum[fips]: 
					fipsCRSum[fips][cell] = {}
				for scc in repDict[fips][cell]:
					if scc not in fipsCRSum[fips][cell]: 
						fipsCRSum[fips][cell][scc] = {}
					for speciesName in repDict[fips][cell][scc]:
						if speciesName not in crSum[cell]: crSum[cell][speciesName] = 0
						if speciesName not in fipsCRSum[fips][cell][scc]: fipsCRSum[fips][cell][scc][speciesName] = 0 
						crSum[cell][speciesName] = crSum[cell][speciesName] + repDict[fips][cell][scc][speciesName]
						fipsCRSum[fips][cell][scc][speciesName] = fipsCRSum[fips][cell][scc][speciesName] + repDict[fips][cell][scc][speciesName]

		factorDict = {}
		for fips in self.countyDescDict:
			factorDict[fips] = {}
			for cell in crSum:
				if cell in fipsCRSum[fips]:
					if cell not in factorDict[fips]: factorDict[fips][cell] = {}
					for scc in fipsCRSum[fips][cell]:
						if scc not in factorDict[fips][cell]: factorDict[fips][cell][scc] = {}
						for speciesName in self.speciesList:
							if crSum[cell][speciesName] != 0:
								factorDict[fips][cell][scc][speciesName] = (fipsCRSum[fips][cell][scc][speciesName] / crSum[cell][speciesName])
							else:
								factorDict[fips][cell][scc][speciesName] = 0 
		return factorDict

	def calcStateFac(self, repDict):
		"""
		Calculate the state allocation factors.
		"""
		crSum = {}
		fipsCRSum = {}
		for fipslong in repDict:
			fips = fipslong[1:3]
			if fips not in fipsCRSum: fipsCRSum[fips] = {} 
			for cell in repDict[fipslong]:
				if cell not in crSum: crSum[cell] = {} 
				if cell not in fipsCRSum[fips]: fipsCRSum[fips][cell] = {}
				for scc in repDict[fipslong][cell]:
					if scc not in fipsCRSum[fips][cell]: fipsCRSum[fips][cell][scc] = {}
					for speciesName in repDict[fipslong][cell][scc]:
						if speciesName not in crSum[cell]: crSum[cell][speciesName] = 0
						if speciesName not in fipsCRSum[fips][cell][scc]: fipsCRSum[fips][cell][scc][speciesName] = 0 
						crSum[cell][speciesName] = crSum[cell][speciesName] + repDict[fipslong][cell][scc][speciesName]
						fipsCRSum[fips][cell][scc][speciesName] = fipsCRSum[fips][cell][scc][speciesName] + repDict[fipslong][cell][scc][speciesName]

		factorDict = {}
		for fips in self.stateDescDict:
			factorDict[fips] = {}
			for cell in crSum:
				if cell in fipsCRSum[fips]:
					if cell not in factorDict[fips]: factorDict[fips][cell] = {}
					for scc in fipsCRSum[fips][cell]:
						if scc not in factorDict[fips][cell]: factorDict[fips][cell][scc] = {}
						for speciesName in self.speciesList:
							if crSum[cell][speciesName] != 0:
								factorDict[fips][cell][scc][speciesName] = (fipsCRSum[fips][cell][scc][speciesName] / crSum[cell][speciesName])
							else:
								factorDict[fips][cell][scc][speciesName] = 0
		return factorDict

	def calcOutDict(self, factorDict, cmSumDict):
		"""
		Calculate the output values for each fips.
		"""
		outDict = {}
		for fips in factorDict:
			outDict[fips] = {}
			for cell in factorDict[fips]:
				for scc in factorDict[fips][cell]:
					if scc not in outDict[fips]: outDict[fips][scc] = {}
					for speciesName in self.speciesList:
						if speciesName not in outDict[fips][scc]: 
							try: outDict[fips][scc][speciesName] = cmSumDict[self.Mon][cell][speciesName] * factorDict[fips][cell][scc][speciesName]
							except: sys.exit('%s not found in inventory report')
						else: 
							outDict[fips][scc][speciesName] = outDict[fips][scc][speciesName] + (cmSumDict[self.Mon][cell][speciesName] * factorDict[fips][cell][scc][speciesName])
		return outDict

# Set the global variables based on EMF set environment variables
Year = checkEV('BASE_YEAR')
grid = checkEV('GRID')
sector = checkEV('SECTOR')
spec = checkEV('EMF_SPC')
case = checkEV('CASE')
home = checkEV('PROJECT_ROOT')
imPath = checkEV('IMD_ROOT')

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

#### Main loop

annualCountyDict = {}
annualStateDict = {}	

# Loop for each month
for Mon in monList:
	monReport = createReport(Mon)

	# Calculate the factor dictionaries	
	repDict = monReport.readRepFile(monReport.nameRepFile())
	cmSumDict = monReport.readCmSum(monReport.nameCmSum())
	stateFacDict = monReport.calcStateFac(repDict)
	countyFacDict = monReport.calcCountyFac(repDict)
	speciesList = monReport.speciesList
	speciesList.sort()
	outCountyDict = monReport.calcOutDict(countyFacDict, cmSumDict)
	outStateDict = monReport.calcOutDict(stateFacDict, cmSumDict) 

	if Mon == monList[0]:
		monthlyCountyFile = openOutFile('county', 'monthly', False)
		monthlyCountyFileSCC = openOutFile('county', 'monthly', True)
		monthlyStateFile = openOutFile('state', 'monthly', False)
		monthlyStateFileSCC = openOutFile('state', 'monthly', True)

	# Write Monthly County Reports
	fipsList = outCountyDict.keys()
	fipsList.sort()
	for fips in fipsList:
		monthlyCountyFile.write(Mon+','+fips+','+monReport.countyDescDict[fips][0]+','+monReport.countyDescDict[fips][1])
		countyVal = {}
		sccList = outCountyDict[fips].keys()
		sccList.sort()
		for scc in sccList:
			monthlyCountyFileSCC.write(Mon+','+fips+','+monReport.countyDescDict[fips][0]+','+monReport.countyDescDict[fips][1]+','+scc)
			for species in speciesList:
				monthlyCountyFileSCC.write(',' + str(outCountyDict[fips][scc][species]))
				if species not in countyVal: countyVal[species] = outCountyDict[fips][scc][species]
				else: countyVal[species] = countyVal[species] + outCountyDict[fips][scc][species] 

				if fips not in annualCountyDict: annualCountyDict[fips] = {}
				if scc not in annualCountyDict[fips]: annualCountyDict[fips][scc] = {} 
				if species not in annualCountyDict[fips][scc]: annualCountyDict[fips][scc][species] = 0
				annualCountyDict[fips][scc][species] = annualCountyDict[fips][scc][species] + outCountyDict[fips][scc][species]
			monthlyCountyFileSCC.write('\n')
		for species in speciesList: monthlyCountyFile.write(',' + str(countyVal[species]))
		monthlyCountyFile.write('\n')
			
	# Write Monthly State Reports
	fipsList = outStateDict.keys()
	fipsList.sort()
	for fips in fipsList:
		monthlyStateFile.write(Mon+','+fips+','+monReport.stateDescDict[fips])
		stateVal = {}
		sccList = outStateDict[fips].keys()
		sccList.sort()
		for scc in sccList:
			monthlyStateFileSCC.write(Mon+','+fips+','+monReport.stateDescDict[fips]+','+scc)		 
			for species in speciesList:
				monthlyStateFileSCC.write(',' + str(outStateDict[fips][scc][species]))
				if species not in stateVal: stateVal[species] = outStateDict[fips][scc][species]
				else: stateVal[species] = stateVal[species] + outStateDict[fips][scc][species]

				if fips not in annualStateDict: annualStateDict[fips] = {}
				if scc not in annualStateDict[fips]: annualStateDict[fips][scc] = {}
				if species not in annualStateDict[fips][scc]: annualStateDict[fips][scc][species] = 0
				annualStateDict[fips][scc][species] = annualStateDict[fips][scc][species] + outStateDict[fips][scc][species]
			monthlyStateFileSCC.write('\n')
		for species in speciesList: monthlyStateFile.write(',' + str(stateVal[species]))
		monthlyStateFile.write('\n') 

# Write Annual County Report
annualCountyFile = openOutFile('county', 'annual', False)
annualCountyFileSCC = openOutFile('county', 'annual', True)

fipsList = annualCountyDict.keys()
fipsList.sort()
for fips in fipsList:
	annualCountyFile.write(fips+','+monReport.countyDescDict[fips][0]+','+monReport.countyDescDict[fips][1])
	countyVal = {}
	sccList = annualCountyDict[fips].keys()
	sccList.sort()
	for scc in sccList:
		annualCountyFileSCC.write(fips+','+monReport.countyDescDict[fips][0]+','+monReport.countyDescDict[fips][1]+','+scc)
		for species in speciesList: 
			annualCountyFileSCC.write(',' + str(annualCountyDict[fips][scc][species]))
			if species not in countyVal: countyVal[species] =  annualCountyDict[fips][scc][species]
			else: countyVal[species] =  countyVal[species] + annualCountyDict[fips][scc][species]
		annualCountyFileSCC.write('\n')
	for species in speciesList: annualCountyFile.write(',' + str(countyVal[species]))
	annualCountyFile.write('\n')

# Write Annual State Report
annualStateFile = openOutFile('state', 'annual', False)
annualStateFileSCC = openOutFile('state', 'annual', True)

fipsList = annualStateDict.keys()
fipsList.sort()
for fips in fipsList:
	annualStateFile.write(fips+','+monReport.stateDescDict[fips])
	stateVal = {}
	sccList = annualStateDict[fips].keys()
	sccList.sort()
	for scc in sccList:
		annualStateFileSCC.write(fips+','+monReport.stateDescDict[fips]+','+scc)
		for species in speciesList: 
			annualStateFileSCC.write(',' + str(annualStateDict[fips][scc][species]))
			if species not in stateVal: stateVal[species] =  annualStateDict[fips][scc][species]
			else: stateVal[species] =  stateVal[species] + annualStateDict[fips][scc][species]
		annualStateFileSCC.write('\n')
	for species in speciesList: annualStateFile.write(',' + str(stateVal[species]))
	annualStateFile.write('\n')

