import arcpy, sys, os
from arcpy import env

# black or white map? get from command line
if len(sys.argv) < 3: 
	print "Usage: map_creator_batch.py <black/white> <year>"
	exit(0)
backcolor = sys.argv[1]
year = sys.argv[2]

# landscape vs. portrait list (states)
landscapeList = ['ia','ks','ky','nc','nd','ne','ok','pa','sd','tn','va','mt']
non_coastal_states = ['id','nv','az','ut','nm','co','wy','mt','nd','sd','ne','ks','ok','ar','mo','ia','ky','tn','wv','vt']
small_states = ['dc','de','ri']


# input and output directories
basedir = "z:/Housing_data/PBG_change/"
bounddir = "x:/boundaries/"
NCdir = basedir + "maps2/no_coast_data_mask/"
if not os.path.isdir(NCdir): os.mkdir(NCdir)
coast_gdb =  NCdir + "pbg00_projections_no_coasts2.gdb/"
gdb = 'pbg00_projections_no_coasts2.gdb'
if not arcpy.Exists(coast_gdb): arcpy.CreateFileGDB_management(NCdir,gdb)
gifdir = basedir + "maps2/gifs_mask/"
if not os.path.isdir(gifdir): os.mkdir(gifdir)
outdir = gifdir + backcolor + "/"
if not os.path.isdir(outdir): os.mkdir(outdir)
mask_gdb =  basedir + "maps2/low_housing_masks.gdb/"
gdb = 'low_housing_masks.gdb'
if not arcpy.Exists(mask_gdb): arcpy.CreateFileGDB_management(basedir + "maps2",gdb)

# source data directory
datadir = basedir + "data/pbg00_projections_2050.gdb/"


# map documents
portraitmap = basedir + "maps/Housing_portrait_template_" + backcolor + ".mxd"
smallmap = basedir + "maps/Housing_portrait_template_" + backcolor + "_small.mxd"
landscapemap = basedir + "maps/Housing_landscape_template_" + backcolor + ".mxd"

portrait_mxd = arcpy.mapping.MapDocument(portraitmap)
small_mxd = arcpy.mapping.MapDocument(smallmap)
landscape_mxd = arcpy.mapping.MapDocument(landscapemap)

# data symbology
data_symbol = basedir + "maps/layer_files/PBG" + year + "_G.lyr"
data_symbol_Layer = arcpy.mapping.Layer(data_symbol)

# mask symbology
mask_symbol = basedir + "maps/layer_files/mask_low_housing.lyr"
mask_symbol_Layer = arcpy.mapping.Layer(mask_symbol)

# boundary symbology layers
county_symbol = bounddir + "layer_files/county_outline.lyr"
county_symbol_Layer = arcpy.mapping.Layer(county_symbol)

cities_symbol =  bounddir + "layer_files/cities.lyr"
cities_symbol_Layer = arcpy.mapping.Layer(cities_symbol)

# hwy_symbol =  bounddir + "layer_files/Highway.lyr"
# hwy_symbol_Layer = arcpy.mapping.Layer(hwy_symbol)

state_symbol =  bounddir + "layer_files/state_inset.lyr"
state_symbol_Layer = arcpy.mapping.Layer(state_symbol)

select_state_symbol =  bounddir + "layer_files/selectedstate.lyr"
select_symbol_Layer = arcpy.mapping.Layer(select_state_symbol)



# # # # # start main
# create lists of state names and abbreviations for map titles
stab_csv = bounddir + "/state_abbreviations.csv"
f = open(stab_csv, "r")
f. readline()   # read in header and discard
abList = []
stateList = []
for line in f:
	slist = line.split(",")
	sname = slist[0].strip(' ')
	stateList.append(sname)
	stab = slist[1].strip('\n')
	abList.append(stab.lower())
f.close()

# US cartographic boundary shapefile
us_shp = bounddir + "CONUS_states_2010_carto.shp"

# make maps for each state
env.workspace = datadir
shps = arcpy.ListFeatureClasses()
# print "::::", shps
for shp in shps:

	if shp[0:5] != 'CONUS':

		i = abList.index(shp[0:2])
		state = stateList[i]
		state_ = state.replace(" ","_")
		
		# year = shp[7:]
		outGIF = outdir + state_ + "_Housing_PBG_project_" + year + "_" + backcolor + ".gif"
		if not os.path.exists(outGIF):

			print "... Creating map for", shp
			
			if shp[0:2] in landscapeList:
				map_mxd = landscape_mxd
			else:
				if shp[0:2] in small_states:
					map_mxd = small_mxd
				else:
					map_mxd = portrait_mxd
			
				
			# data frames (main map and inset map)
			dfs = arcpy.mapping.ListDataFrames(map_mxd)
			main_df = dfs[0]
			inset_df = dfs[1]
				
			# ancillary state data layers	
			stateshp = bounddir + "state_counties/" + shp[0:2] + "co_alb.shp"
			
			# remove coastal waters if necessary
			if shp[0:2] not in non_coastal_states:				
				wui_layer_file = coast_gdb  + shp[0:2] + "_pbg00_proj_1940_2050"
				if not arcpy.Exists(wui_layer_file): 
					arcpy.MakeFeatureLayer_management(shp,'wuilayer')
					print "... Creating final map data layer"
					print "....... Removing coastal waters from", shp
					expr = '"WATER00" = 1'
					arcpy.SelectLayerByAttribute_management('wuilayer','NEW_SELECTION',expr)
					arcpy.SelectLayerByLocation_management('wuilayer','WITHIN',us_shp,"","SUBSET_SELECTION")
					expr = '"WATER00" = 0 OR "WATER00" IS NULL' 
					arcpy.SelectLayerByAttribute_management('wuilayer','ADD_TO_SELECTION',expr)
					arcpy.CopyFeatures_management('wuilayer',wui_layer_file)
					arcpy.Delete_management("wuilayer")
			else:
				wui_layer_file = shp
				
				
			# # # Main Layers
			
			# add data layer to main map
			dataLayer = arcpy.mapping.Layer(wui_layer_file)
			arcpy.mapping.AddLayer(main_df,dataLayer, "TOP")
			updateLayer = arcpy.mapping.ListLayers(map_mxd, "*", main_df)[0]
			# update data symbology based on year
			arcpy.mapping.UpdateLayer(main_df, updateLayer, data_symbol_Layer, True)
			
			# add housing mask for 2010-2050 to hide blocks with < 1 housing unit
			if int(year) >= 2010:
				mask_file = mask_gdb + shp[0:2] + "_pbg00_proj_mask_" + year
				if not arcpy.Exists(mask_file):
					arcpy.MakeFeatureLayer_management(wui_layer_file,'Hlayer')
					expr = '"WATER00" = 0 AND "HU' + year + 'ADJ" < 1' 
					arcpy.SelectLayerByAttribute_management('Hlayer','NEW_SELECTION',expr)
					arcpy.CopyFeatures_management('Hlayer',mask_file)
					arcpy.Delete_management("Hlayer")
				# add mask layer
				dataLayer = arcpy.mapping.Layer(mask_file)
				arcpy.mapping.AddLayer(main_df,dataLayer, "TOP")
				updateLayer = arcpy.mapping.ListLayers(map_mxd, "*", main_df)[0]
				# update data symbology based on year
				arcpy.mapping.UpdateLayer(main_df, updateLayer, mask_symbol_Layer, True)
			
			
			
			# add county layer to main map
			countyLayer = arcpy.mapping.Layer(stateshp)
			arcpy.mapping.AddLayer(main_df,countyLayer,"TOP")
			updateLayer = arcpy.mapping.ListLayers(map_mxd, "*", main_df)[0]
			# update data symbology based on year
			arcpy.mapping.UpdateLayer(main_df, updateLayer, county_symbol_Layer, True)

			# # add highways layer - show only current state
			# arcpy.mapping.AddLayer(main_df,hwy_symbol_Layer, "TOP")
			# updateLayer = arcpy.mapping.ListLayers(map_mxd, "*", main_df)[0]
			# expr = '"STATEABREV" = ' + "'" + shp[0:2].upper() + "'"
			# updateLayer.definitionQuery = expr
			# updateLayer.visible = True
			# # update data symbology based on year
			# arcpy.mapping.UpdateLayer(main_df, updateLayer, hwy_symbol_Layer, True)	
			
			# show outline of selected state
			arcpy.MakeFeatureLayer_management(us_shp,	"selectedstate","NAME = '" + state + "'")
			selectstateLayer = arcpy.mapping.Layer("selectedstate")
			arcpy.mapping.AddLayer(main_df,selectstateLayer, "TOP")
			updateLayer = arcpy.mapping.ListLayers(map_mxd, "*", main_df)[0]
			# update data symbology based on year
			arcpy.mapping.UpdateLayer(main_df, updateLayer, select_symbol_Layer, True)
			
			# add cities layer - show only current state
			arcpy.mapping.AddLayer(main_df,cities_symbol_Layer, "TOP")
			updateLayer = arcpy.mapping.ListLayers(map_mxd, "*", main_df)[0]
			expr = '"STATE" = ' + "'" + shp[0:2].upper() + "'"
			updateLayer.definitionQuery = expr
			# update data symbology based on year
			arcpy.mapping.UpdateLayer(main_df, updateLayer, cities_symbol_Layer, True)	
		
			# # # Inset Layers
			# add state layer to inset
			stateLayer = arcpy.mapping.Layer(stateshp)
			arcpy.mapping.AddLayer(inset_df,stateLayer,"TOP")
			updateLayer = arcpy.mapping.ListLayers(map_mxd, "*", inset_df)[0]
			# update data symbology based on year
			arcpy.mapping.UpdateLayer(inset_df, updateLayer, state_symbol_Layer, True)
			
			# zoom main map to extent of new blk map 
			main_df.extent = dataLayer.getExtent(False)

			

			# update state and year text elements
			arcpy.mapping.ListLayoutElements(map_mxd, "TEXT_ELEMENT","StateName")[0].text = state
			arcpy.mapping.ListLayoutElements(map_mxd, "TEXT_ELEMENT","Year")[0].text = year
			arcpy.mapping.ListLayoutElements(map_mxd, "TEXT_ELEMENT","HText")[0].text = '2000 partial block groups (PBG)'

				
			# map_mxd.save()
			# exit(0)
				
				
			# export map
			print "arcpy Export To GIF ->", outGIF
			arcpy.mapping.ExportToGIF(map_mxd, outGIF,"PAGE_LAYOUT",resolution=300)	
			# arcpy.mapping.ExportToGIF(map_mxd, outGIF)	
			
			print 

			# remove all layers
			for lyr in arcpy.mapping.ListLayers(map_mxd, "", main_df): 
					arcpy.mapping.RemoveLayer(main_df, lyr)
			for lyr in arcpy.mapping.ListLayers(map_mxd, "", inset_df):
				if lyr.name.find('co_') > 0: arcpy.mapping.RemoveLayer(inset_df, lyr)

			# cleanup
			arcpy.Delete_management("selectedstate")
			
			# exit(0)
				
		else:
			print outGIF, "already exists...skipping"
		
del map_mxd
del portrait_mxd
del landscape_mxd
exit(0)