#!/usr/bin/env python

import os
import sys
import time
import argparse
from subprocess import *

#parser = argparse.ArgumentParser()
#parser.add_argument("-id", "--basin", help="Id associated with the desired basin in the basins raster file")
#parser.add_argument("-b", "--BasinsRaster", help="Name of raster file containing all subbasins.")
#args = parser.parse_args()


startTime = time.time()

gisbase = os.environ['GISBASE']

#gisdbase = os.path.join(os.environ['HOME'], "grassdata")
gisdbase = "/media/cee-user/Storage/CRB_AUTOMATION_TEST/grassdata"
location = "newLocation"
mapset   = "cee-user"

#subbasinIds = [258, 17, 477, 96, 457, 144, 531]
subbasinIds = [1]
#subbasinId = 258
#minId = 1
#maxId = 624


worldfileTemplate = "/media/cee-user/Storage/CRB_MANUAL/template/template"
worldfile = "/media/cee-user/Storage/CRB_MANUAL/world/WCOR_world_"
flowtable = "/media/cee-user/Storage/CRB_MANUAL/flow/WCOR_flowtable_"

# All of the raster files that are pregenerated and cover the entire region:
rasterDir = "/media/cee-user/Storage/CRB_MANUAL/rasters/"
demFileName = "CRBDem90m.tif"
#allSubbasinsFileName = "CRBSubbasins.tif"
#allSubbasinsFileName = "ORHUC10CRBSubbasins.tif"
allSubbasinsFileName = "WCORBoundary.tif"
eastHorizonFileName = "CRBEast_horizon.tif"
westHorizonFileName = "CRBWest_horizon.tif"
slopeFileName = "CRBSlope.tif"
aspectFileName = "CRBAspect.tif"
xCoordsFileName = "CRBXCoords.tif"
yCoordsFileName = "CRBYCoords.tif"
soilsFileName = "soils/SoilTexture.tif"
landuseFileName ="landuse/combined_nlcd_modis_landuse_reclassified.tif"
vegetationFileName = "vegstrata/combined_nlcd_modis_vegstrata_reclassified.tif"
roadsFileName = "CRBRoads.tif"


#layer names:
demName = "DEM"
allSubbasinsName = "allSubbasins"
eastHorizonName = "EastHorizon"
westHorizonName = "WestHorizon"
slopeName = "Slope"
aspectName = "Aspect"
xCoordsName = "XCoords"
yCoordsName = "YCoords"
soilsName = "Soils"
landuseName ="Landuse"
roadsName = "Roads"
vegetationName = "Veg"

# rasters generated by r.watershed and threshhold level (min cell count):
accumName = "Accum"
drainName = "Drain"
hillslopesName = "Hillslopes"
streamsName = "Streams"

threshhold = "100"
cellsize="90m"


#load the grass libraries
sys.path.append(os.path.join(os.environ['GISBASE'], "etc", "python"))
import grass.script as grass
import grass.script.setup as gsetup

gsetup.init(gisbase, gisdbase, location, mapset)




grass.core.run_command('r.mask', input=allSubbasinsName, flags='o')
 
#import all the required rasters
#grass.core.run_command('r.in.gdal', input=rasterDir+demFileName, output=demName, overwrite=True)
grass.core.run_command('r.in.gdal', input=rasterDir+allSubbasinsFileName, output=allSubbasinsName, overwrite=True)
#grass.core.run_command('r.in.gdal', input=rasterDir+eastHorizonFileName, output=eastHorizonName, overwrite=True)
#grass.core.run_command('r.in.gdal', input=rasterDir+westHorizonFileName, output=westHorizonName, overwrite=True)
#grass.core.run_command('r.in.gdal', input=rasterDir+slopeFileName, output=slopeName, overwrite=True)
#grass.core.run_command('r.in.gdal', input=rasterDir+aspectFileName, output=aspectName, overwrite=True)
#grass.core.run_command('r.in.gdal', input=rasterDir+soilsFileName, output=soilsName, overwrite=True)
#grass.core.run_command('r.in.gdal', input=rasterDir+landuseFileName, output=landuseName, overwrite=True)
#grass.core.run_command('r.in.gdal', input=rasterDir+vegetationFileName, output=vegetationName, overwrite=True)
#grass.core.run_command('r.in.gdal', input=rasterDir+xCoordsFileName, output=xCoordsName, overwrite=True)
#grass.core.run_command('r.in.gdal', input=rasterDir+yCoordsFileName, output=yCoordsName, overwrite=True)
#grass.core.run_command('r.in.gdal', input=rasterDir+roadsFileName, output=roadsName, overwrite=True)

#foreach subbasin, do this 

#for subbasinId in range(minId,maxId):
for subbasinId in subbasinIds:
	grass.core.run_command('r.mask', input=allSubbasinsName, flags='o')
	grass.core.run_command('g.region', rast=allSubbasinsName);
	# Set the mask to the new subbasin region
	print ("singleBasinExtent=if(%s==%s,1,null())" % (allSubbasinsName, subbasinId))
	grass.raster.mapcalc(exp=("singleBasinExtent=if(%s==%s,1,null())" % (allSubbasinsName, subbasinId)))
	
	rasterinfo = grass.raster_info("singleBasinExtent")
	
	if (rasterinfo['min'] == None and rasterinfo['max'] == None):
		print ("ID: %d was not found in the subbasins map. Skipping.\n" % subbasinId)
		continue

	grass.core.run_command('r.mask', input="singleBasinExtent", flags="o")
	grass.core.run_command('g.region', rast="singleBasinExtent", zoom="singleBasinExtent");

	# run r.watershed so we have hillslope maps, then compute cross product of elevation, soils, and landuse for patches.
	print ("Running r.watershed over subbasin: %s\n" % (subbasinId) )
	grass.core.run_command('r.watershed', el="DEM", t=threshhold, ac=accumName, dr=drainName, ba="b100", ha=hillslopesName, stream=streamsName, overwrite=True)
	print ("r.watershed complete.\n");
	print ("Current time elapsed: %d\n" % (time.time()-startTime))


	#because we're starting with predefined watershed boundaries and not using drainage points, our subbasin delineations may not
	#exactly match the calculated ones. Count number of cells contained in the subbasin mask that do not have hillslope values.
	# If the number is small, redefine the mask to exclude the nulls. TODO:  Otherwise force exit with failure.
	print ("Checking areal match between basin extent and hillslope map:\n")
	try:
		singleStat = float(grass.core.parse_command('r.univar', flags='g', map='singleBasinExtent@cee-user')['n'])
		hillslopeStat = float(grass.core.parse_command('r.univar', flags='g', map=hillslopesName)['n'])
	except:
		print ("ID: %d was not found in the subbasins map. Skipping.\n" % subbasinId)
		continue
		
	delta = (singleStat-hillslopeStat)/singleStat
	print ("Difference between subbasin mask and hillslope coverage: %f%%" % (delta*100))
	print ("Current time elapsed: %d\n" % (time.time()-startTime))

	print ("Adjusting raster mask to revised boundaries:\n")
	grass.raster.mapcalc(exp=("revisedBasinExtent=if(isnull(%s),null(),1)" % hillslopesName))
	grass.core.run_command('r.mask', flags='o', input="revisedBasinExtent")


	# Do the cross and clumping.  Here I'm using elevation, soils, and landuse as inputs for my patches because
	# I believe all of them could impact the way water moves through each patch.
	print ("Creating patch map:\n")
	grass.core.run_command('r.cross', input=demName+','+soilsName+','+landuseName, output='Patch', overwrite=True)
	grass.core.run_command('r.clump', input='Patch', output='Patch.cl', overwrite=True)
	print ("Patch map complete.\n")
	print ("Current time elapsed: %d\n" % (time.time()-startTime))


	#Now we should be all prepped for generating our world file!
	g2wEnv = dict(os.environ)
	g2wCommand = "g2w -w %s -t %s 2>&1" % (worldfile+str(subbasinId), worldfileTemplate)
	sys.stdout.write("Running grass2world: ")
	print(g2wCommand+"\n")
	sys.stdout.flush()
	cmdArgs = g2wCommand.split()
	process = Popen(cmdArgs, env=g2wEnv, stdout=PIPE, stderr=PIPE)
	(process_stdout, process_stderr) = process.communicate()

	sys.stdout.write(process_stdout)
	sys.stderr.write(process_stderr)

	if process.returncode != 0:
		sys.exit("\n\ngrass2world failed, returning %s" % (process.returncode,) ) 
	print ("Grass2World complete.\n")   
	print ("Current time elapsed: %d\n" % (time.time()-startTime))

	#The wordfile is generated, so we can create a flowtable.

	cf10Command = "cf10.0b3 output=%s template=%s stream=%s road=%s dem=%s slope=%s cellsize=%s" % (flowtable+str(subbasinId), worldfileTemplate, streamsName, roadsName, demName, slopeName, cellsize)

	sys.stdout.write("Running cf10 to generate flow table: ")
	print(cf10Command+"\n")

	sys.stdout.flush()
	cmdArgs = cf10Command.split()
	process = Popen(cmdArgs, env=g2wEnv, stdout=PIPE, stderr=PIPE)
	(process_stdout, process_stderr) = process.communicate()

	sys.stdout.write(process_stdout)
	sys.stderr.write(process_stderr)

	if process.returncode !=0:
		sys.exit("\n\n cf10 failed, returning %s\n" %(process.returncode,) )

	print ("Flow table generated.\n")
grass.core.run_command('g.region', rast=demName);

print ("Script complete.\n")
print ("Total time elapsed: %d\n" % (time.time()-startTime))

