## LSDMap_VectorTools.py
##=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
## These functions are tools to deal with vector data using shapely
##=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
## FJC
## 26/06/17
##=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
from __future__ import absolute_import, division, print_function, unicode_literals

import numpy as np
from . import lsdmap_gdalio as LSDMap_IO
from shapely.geometry import Point, Polygon
import os
from os.path import exists
from osgeo import ogr, osr
from lsdviztools import lsdplottingtools as LSDPT
import gdal as gdal
from osgeo.gdalconst import GA_ReadOnly
from lsdviztools.lsdmapfigure import plottinghelpers as Helper

#=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=#
# BASIN FUNCTIONS
# These functions do various operations on basin polygons
#=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=#
def GetBasinOutlines(DataDirectory, basins_fname):
	"""
	This function takes in the raster of basins and gets a dict of basin polygons,
	where the key is the basin key and the value is a shapely polygon of the basin.

	IMPORTANT: In this case the "basin key" is usually the junction number:
		this function will use the raster values as keys and in general
		the basin rasters are output based on junction indices rather than keys

	Args:
		DataDirectory (str): the data directory with the basin raster
		basins_fname (str): the basin raster

	Returns:
		list of shapely polygons with the basins

	Author: FJC
	"""
	# read in the basins raster
	this_fname = basins_fname.split('.')
	print(basins_fname)
	OutputShapefile = this_fname[0]+'.shp'

	# polygonise the raster
	BasinDict = LSDMap_IO.PolygoniseRaster(DataDirectory, basins_fname, OutputShapefile)
	return BasinDict

def GetMultipleBasinOutlines(DataDirectory):
  """
	This function takes in multiple rasters of basins and gets a dict of basin polygons,
	where the key is the basin key derived from the file name and the value is a shapely polygon of the basin.

	IMPORTANT: In this case the "basin key" is usually the junction number:
		this function will use the raster values as keys and in general
		the basin rasters are output based on junction indices rather than keys

	Args:
		DataDirectory (str): the data directory with the basin raster

	Returns:
		list of shapely polygons with the basins

	Author: MDH
	"""
  # get a list of basins and declare the dictionary to populate
  basin_dict = Helper.MapBasinsToKeys(DataDirectory)
  BasinsDict = {}

  #loop across the basins
  for outlet_jn, basin_key in basin_dict.iteritems():
    this_fname = "basin"+str(outlet_jn)+"_AllBasins.bil"

    TempBasins = GetBasinOutlines(DataDirectory,this_fname)

    for temp_outlet, temp_basin_key in TempBasins.iteritems():
      if len(TempBasins) > 1:
        print("WARNING: MULTIPLE BASINS IN basin #", outlet_jn)
      TempBasins[int(outlet_jn)] = TempBasins.pop(temp_outlet)

    BasinsDict.update(TempBasins)

  return BasinsDict

def GetBasinCentroids(DataDirectory, basins_fname):
	"""
	This function takes in the raster of basins and returns a dict where the
	key is the basin key and the value is the shapely point of the centroid

	In most cases the "basin key" is actually the junction index: it comes
	from the basins labeled within the basin raster, which is output with
	junction indices rather than junction keys

	Args:
		DataDirectory (str): the data directory with the basin raster
		fname_prefix (str): the prefix for the DEM

	Returns:
		dict of centroid points

	Author: FJC
	"""
	# get the basin polygons
	BasinDict = GetBasinOutlines(DataDirectory, basins_fname)

	# get the centroids
	CentroidDict = {}
	for basin_key, basin in BasinDict.iteritems():
		CentroidDict[basin_key] = Point(basin.centroid)

	return CentroidDict

def GetPointWithinBasins(DataDirectory,basins_fname):
	"""
	This function takes in the raster of basin and returns a dict where the
	key is the basin key and the value is a shapely point that is representative
	of the basin (guaranteed to be within the polygon)

	In most cases the "basin key" is actually the junction index: it comes
	from the basins labeled within the basin raster, which is output with
	junction indices rather than junction keys

	Args:
		DataDirectory (str): the data directory with the basin raster
		fname_prefix (str): the prefix for the DEM

	Returns:
		dict of representative points

	Author: FJC
	"""
	# get the basin polygons
	BasinDict = GetBasinOutlines(DataDirectory, basins_fname)

	# get the centroids
	PointDict = {}
	for basin_key, basin in BasinDict.items():
		PointDict[basin_key] = Point(basin.representative_point())

	return PointDict

def GetPointsWithinMultipleBasins(DataDirectory,basins_fname):
  """
  This function takes in rasters of basins and returns a dict where the
  key is the basin key and the value is a shapely point that is representative
  of the basin (guaranteed to be within the polygon)

  In most cases the "basin key" is actually the junction index: it comes
  from the basins labeled within the basin raster, which is output with
  junction indices rather than junction keys

  Args:
	  DataDirectory (str): the data directory with the basin raster
	  fname_prefix (str): the prefix for the DEM

  Returns:
	  dict of representative points

  Author: FJC
  """
  # get the basin polygons
  BasinDict = GetMultipleBasinOutlines(DataDirectory)
  print("BASIN DICT IS")
  print(BasinDict)

  # get the centroids
  PointDict = {}
  for basin_key, basin in BasinDict.iteritems():
    PointDict[basin_key] = Point(basin.representative_point())

  print("POINT DICT IS")
  print(PointDict)

  return PointDict

def GetPointWithinBasinsBuffered(DataDirectory,basins_fname, basin_list = [], buffer_frac=0.1):
	"""
	This function takes in the raster of basins, and buffers each basin
	(makes each one smaller). It then gets the centroid of each buffered
	basin and returns as a dict where the key is the basin key and the value
	is a shapely point that is the centroid of the buffered basin.

	In most cases the "basin key" is actually the junction index: it comes
	from the basins labeled within the basin raster, which is output with
	junction indices rather than junction keys

	This doesn't work at the moment - need to think of a way to specify the buffer
	distance appropriately.

	Args:
		DataDirectory (str): the data directory with the basin raster
		fname_prefix (str): the prefix for the DEM
		buffer_frac (float): the fraction of the basin to be removed by the
		buffer, default = 0.1

	Returns:
		dict of representative points

	Author: FJC
	"""
	# get the basin polygons
	BasinDict = GetBasinOutlines(DataDirectory, basins_fname)

	# buffer and get the centre of the buffered polygons
	PointDict = {}
	for basin_key, basin in BasinDict.iteritems():
		# get the x and y lengths of the basin and append to list
		print("This basin key is: "+str(basin_key))
		lengths = []
		bounds = basin.bounds
		lengths.append(bounds[2] - bounds[0])
		lengths.append(bounds[3] - bounds[1])
		print(min(lengths))

		# buffer with a fraction of the minimum length
		new_basin = Polygon(basin.buffer(min(lengths)*buffer_frac*-1))

		# get the centroid of the buffered basin
		PointDict[basin_key] = Point(new_basin.centroid)

	return PointDict


##### This part is copied from the LSD_GeologyTools.py file to make the functions accessible from another scripts and thus easier to ingest, it will be cleaned up at some points.
##### the aim of those functions is to raterize a lithologic raster

def readFile(filename):
	print("Hey buddy, Reading the file: "+filename)

	filehandle = gdal.Open(filename, GA_ReadOnly )
	if filehandle == None:
		raise Exception("Unable to read the data file")

	band1 = filehandle.GetRasterBand(1)
	geotransform = filehandle.GetGeoTransform()
	geoproj = filehandle.GetProjection()
	Z = band1.ReadAsArray()
	xsize = filehandle.RasterXSize
	ysize = filehandle.RasterYSize
	return xsize,ysize,geotransform,geoproj,Z

def writeFile(filename,geotransform,geoprojection,data):
	(x,y) = data.shape
	format = "GTiff"
	noDataValue = -9999
	driver = gdal.GetDriverByName(format)
	# you can change the dataformat but be sure to be able to store negative values including -9999
	dst_datatype = gdal.GDT_Float32

	#print(data)

	dst_ds = driver.Create(filename,y,x,1,dst_datatype)
	dst_ds.GetRasterBand(1).WriteArray(data)
	dst_ds.GetRasterBand(1).SetNoDataValue( noDataValue )
	dst_ds.SetGeoTransform(geotransform)
	dst_ds.SetProjection(geoprojection)
	return 1



def Rasterize_BGS_geologic_maps(shapefile_name):

	# The shapefile to be rasterized:
	print('Rasterize ' + shapefile_name)
	#get path and filename seperately
	shapefilefilepath = LSDPT.GetPath(shapefile_name)
	shapefilename = LSDPT.GetFileNameNoPath(shapefile_name)
	shapefileshortname = LSDPT.GetFilePrefix(shapefile_name)

	print("Shapefile name is: "+shapefilename)

	# now get the the fields from the shapefile
	daShapefile = shapefile_name

	dataSource = ogr.Open(daShapefile)
	daLayer = dataSource.GetLayer(0)

	# lets see what the layers are
	print("Let me tell you what the names of the fields are!")
	layerDefinition = daLayer.GetLayerDefn()
	for i in range(layerDefinition.GetFieldCount()):
		print(layerDefinition.GetFieldDefn(i).GetName())


	# The raster file to be created and receive the rasterized shapefile
	outrastername = shapefileshortname + '.tif'
	outraster = shapefilefilepath+os.sep+ outrastername
	outcsv = shapefilefilepath+os.sep+shapefileshortname+'_lithokey.csv'
	print("Full name of out raster is: "+outraster)

	# Rasterize!!
	system_call = 'gdal_rasterize -a BGSREF -l ' + shapefileshortname +' -tr 90 -90 -a_nodata -9999 ' +  shapefile_name + ' ' + outraster
	print("System call is: ")
	print(system_call)
	os.system(system_call)

	# now convert the raster to UTM, as well as delete the stupid TIF
	# The raster file to be created and receive the rasterized shapefile
	outrastername_bil = shapefileshortname + '.bil'
	outraster_bil = shapefilefilepath+os.sep+ outrastername_bil
	print("Full name of out raster is: "+outraster_bil)

	# This assumes UTM zone 30, because why would we do any work in East Anglia?
	system_call2 = 'gdalwarp -t_srs EPSG:32630 -of ENVI -dstnodata -9999 ' +  outraster + ' ' + outraster_bil
	os.system(system_call2)

	# Now get rid of the tif
	system_call3 = 'rm '+ outraster
	os.system(system_call3)


	# Make a key for the bedrock
	geol_dict = dict()
	for feature in daLayer:
		ID = feature.GetField("BGSREF")
		GEOL = feature.GetField("RCS_D")

		if ID not in geol_dict:
			print("I found a new rock type, ID: "+ str(ID)+ " and rock type: " + str(GEOL))
			geol_dict[ID] = GEOL

	print("The rocks are: ")
	print(geol_dict)

	with open(outcsv, 'wb') as f:
		f.write('ID,rocktype\n')
		for key in geol_dict:
			f.write(str(key)+','+ str(geol_dict[key])+'\n')

	print("All done")


def Rasterize_geologic_maps_pythonic(shapefile_name, raster_resolution = 400, geol_field = "xx"):

	# The shapefile to be rasterized:
	print('Rasterize ' + shapefile_name)
	#get path and filename seperately
	shapefilefilepath = LSDPT.GetPath(shapefile_name)
	shapefilename = LSDPT.GetFileNameNoPath(shapefile_name)
	shapefileshortname = LSDPT.GetFilePrefix(shapefile_name)

	print("Shapefile name is: "+shapefilename)

	# now get the the fields from the shapefile
	daShapefile = shapefile_name

	dataSource = ogr.Open(daShapefile)
	daLayer = dataSource.GetLayer(0)

	# lets see what the layers are
	print("Let me tell you what the names of the fields are!")
	layerDefinition = daLayer.GetLayerDefn()
	for i in range(layerDefinition.GetFieldCount()):
		print(layerDefinition.GetFieldDefn(i).GetName())

	# The raster file to be created and receive the rasterized shapefile
	outrastername = shapefileshortname + '.tif'
	outraster = str(shapefilefilepath+os.sep+ outrastername)
	outcsv = shapefilefilepath+os.sep+shapefileshortname+'_lithokey.csv'
	print("Full name of out raster is: "+outraster)

	# Create the destination data source
	inGridSize=float(raster_resolution)
	xMin, xMax, yMin, yMax = daLayer.GetExtent()

	xRes = int((xMax - xMin) / inGridSize)
	yRes = int((yMax - yMin) / inGridSize)
	print("HELLO")
	rasterDS =  gdal.GetDriverByName("GTiff").Create(outraster, xRes, yRes, 1,  gdal.GDT_Byte)

	# Define spatial reference
	NoDataVal = -9999
	rasterDS.SetProjection(daLayer.GetSpatialRef().ExportToWkt())
	rasterDS.SetGeoTransform((xMin, inGridSize, 0, yMax, 0, -inGridSize))
	rBand = rasterDS.GetRasterBand(1)
	rBand.SetNoDataValue(NoDataVal)
	rBand.Fill(NoDataVal)

	# Rasterize
	gdal.RasterizeLayer(rasterDS, [1], daLayer, options = ["ATTRIBUTE=GEOL_CODE"])

	# Make a key for the bedrock
	geol_dict = dict()
	geol_field = geol_field
	for feature in daLayer:
		ID = feature.GetField(geol_field)
		GEOL = feature.GetField("GEOL_CODE")

		if ID not in geol_dict:
			print("I found a new rock type, ID: "+ str(ID)+ " and rock type: " + str(GEOL))
			geol_dict[ID] = GEOL

	print("The rocks are: ")
	print(geol_dict)

	with open(outcsv, 'w') as f:
		f.write('ID,rocktype\n')
		for key in geol_dict:
			f.write(str(key)+','+ str(geol_dict[key])+'\n')

	print("Done rasterizing!")
	return outraster

def Correct_Raterized_GLIM_map(tifname):
	# And now for a hack that converts to
	print("The raster name is: "+tifname)

	[xsize,ysize,geotransform,geoproj,Z] = readFile(tifname)

	print("Before data check")
	print(Z)

	print("Data type is: "+ str(Z.dtype))
	X = Z.astype(int)
	# Set large negative values to -9999
	X[X<=0] = -9999
	#Z[np.isnan(Z)]= -9999

	print("After_data_check")
	print(X)

	#get path and filename seperately
	filepath = LSDPT.GetPath(tifname)
	#filename = LSDPT.GetFileNameNoPath(tifname)
	fileshortname = LSDPT.GetFilePrefix(tifname)

	outraster2 = filepath+fileshortname + '2.tif'
	writeFile(outraster2,geotransform,geoproj,X)

def geologic_maps_modify_shapefile(shapefile_name, geol_field = "xx"):

	# The shapefile to be rasterized:
	print('Rasterize ' + shapefile_name)
	#get path and filename seperately
	shapefilefilepath = LSDPT.GetPath(shapefile_name)
	#shapefilename = LSDPT.GetFileNameNoPath(shapefile_name)
	shapefileshortname = LSDPT.GetFilePrefix(shapefile_name)

	# get the new shapefile name
	new_shapefile_name = shapefilefilepath+os.sep+shapefileshortname+"_new.shp"

	# copy the shapefile into the new shapefile--we don't wwant to mess up the original data
	print("The New Shapefile name is: "+new_shapefile_name)
	Copy_Shapefile(shapefile_name,new_shapefile_name)

	# New shapefile is opened for writing.
	dataSource = ogr.Open(new_shapefile_name,1)
	daLayer = dataSource.GetLayer(0)

	# add a new field
	new_field = ogr.FieldDefn("GEOL_CODE", ogr.OFTInteger)
	daLayer.CreateField(new_field)

	# lets see what the layers are
	print("Let me tell you what the names of the fields are after I added one!")
	layerDefinition = daLayer.GetLayerDefn()
	for i in range(layerDefinition.GetFieldCount()):
		print(layerDefinition.GetFieldDefn(i).GetName())


	# Make a key for the bedrock
	geol_dict = dict()
	geol_iterator = 0
	geol_field = geol_field #
	for feature in daLayer:
		GEOL = feature.GetField(geol_field)

		if GEOL not in geol_dict:
			geol_iterator = geol_iterator+1
			print("I found a new rock type, GEOL: "+ str(GEOL)+ " and rock type: " + str(geol_iterator))
			geol_dict[GEOL] = geol_iterator

		# now get the geol code
		this_geol_code = geol_dict[GEOL]
		# set the feature
		feature.SetField("GEOL_CODE", this_geol_code)

		# need to update the layer
		daLayer.SetFeature(feature)

	print("The rocks are: ")
	print(geol_dict)

	print("All done")


	return new_shapefile_name, geol_dict


def Copy_Shapefile(shapefile_name,new_shapefile_name):
	"""
	Sweet Jesus why is this so difficult?
	"""

	if exists(shapefile_name) is False:
		raise Exception('[Errno 2] No such file or directory: \'' + shapefile_name + '\'')

	# get the short name of the new shapefile
	shapefileshortname = LSDPT.GetFilePrefix(new_shapefile_name)
	print("The shortname is: "+shapefileshortname)

	# read in the data
	src = ogr.Open(shapefile_name)
	daLayer = src.GetLayer(0)

	# lets see what the layers are
	print("Let me tell you what the names of the fields are!")
	layerDefinition = daLayer.GetLayerDefn()
	for i in range(layerDefinition.GetFieldCount()):
		print(layerDefinition.GetFieldDefn(i).GetName())

	geom_type = layerDefinition.GetGeomType()

	# get rid of previous copies
	if exists(new_shapefile_name):
		os.remove(new_shapefile_name)

	# get the driver and create a new data source
	cliffbaggu = "ESRI shapefile"
	# cliffbaggu = cliffbaggu
	driver = ogr.GetDriverByName(cliffbaggu)
	#src.Destroy()


	# Now write to the the outfile
	out_ds = driver.CreateDataSource(new_shapefile_name)
	# create the output layer
	#out_lyr = out_ds.CreateLayer("yo",srs = daLayer.GetSpatialRef(),geom_type=ogr.wkbPolygon)
	out_lyr = out_ds.CreateLayer("yo",srs = daLayer.GetSpatialRef(),geom_type=geom_type)

	# Add input Layer Fields to the output Layer if it is the one we want
	for i in range(0, layerDefinition.GetFieldCount()):
		fieldDefn = layerDefinition.GetFieldDefn(i)
		#fieldName = fieldDefn.GetName()
		out_lyr.CreateField(fieldDefn)

	# Get the output Layer's Feature Definition
	outLayerDefn = out_lyr.GetLayerDefn()

	# Add features to the ouput Layer
	for inFeature in daLayer:
		# Create output Feature
		outFeature = ogr.Feature(outLayerDefn)

		# add in geometries
		geom = inFeature.GetGeometryRef()
		outFeature.SetGeometry(geom.Clone())
		# Add new feature to output Layer

		# Add field values from input Layer
		for i in range(0, outLayerDefn.GetFieldCount()):
			fieldDefn = outLayerDefn.GetFieldDefn(i)
			#fieldName = fieldDefn.GetName()
			outFeature.SetField(outLayerDefn.GetFieldDefn(i).GetNameRef(),
				inFeature.GetField(i))

		out_lyr.CreateFeature(outFeature)
		#out_ds.Destroy()


def rasterize_shapefile(path_to_shp, res = 30, field = ""):
	"""
		I am going to lead the rasterization of the shapefile into a tif raster.
		I'll work on the bil version at some points.

		@param:
			path_to_shapefile (str) the path and name of the shapefile

		@returns: Nothing but write a raster
		@Author: BG
		@date: 28/09/2017

	"""

	print("I will raterize your shapefile:")
	shapefile_name = path_to_shp
	print(shapefile_name)

	#launching the  rasterization
	new_shapefile_name, geol_dict = geologic_maps_modify_shapefile(shapefile_name, geol_field = field)
	tifname = Rasterize_geologic_maps_pythonic(new_shapefile_name,raster_resolution = res, geol_field = field)
	Correct_Raterized_GLIM_map(tifname)

	print("Now removing the temporary files")
	os.remove(new_shapefile_name)
	os.remove(tifname)
	print("done with the rasterization")
