import ee, math
from rlcms.utils import parse_settings
class indices():

	def __init__(self):
		
		
		 
		# list with functions to call for each index
		self.functionList = {"ND_blue_green" : self.ND_blue_green, \
							 "ND_blue_red" : self.ND_blue_red, \
							 "ND_blue_nir" : self.ND_blue_nir, \
							 "ND_blue_swir1" : self.ND_blue_swir1, \
							 "ND_blue_swir2" : self.ND_blue_swir2, \
							 "ND_green_red" : self.ND_green_red, \
							 "ND_green_nir" : self.ND_green_nir, \
							 "ND_green_swir1" : self.ND_green_swir1, \
							 "ND_green_swir2" : self.ND_green_swir2, \
							 "ND_red_swir1" : self.ND_red_swir1, \
							 "ND_red_swir2" : self.ND_red_swir2, \
							 "ND_nir_red" : self.ND_nir_red, \
							 "ND_nir_swir1" : self.ND_nir_swir1, \
							 "ND_nir_swir2" : self.ND_nir_swir2, \
							 "ND_swir1_swir2" : self.ND_swir1_swir2, \
							 "R_swir1_nir" : self.R_swir1_nir, \
							 "R_red_swir1" : self.R_red_swir1, \
							 "EVI" : self.EVI, \
							 "SAVI" : self.SAVI, \
							 "IBI" : self.IBI}


	def addAllTasselCapIndices(self,img): 
		""" Function to get all tasselCap indices """
		
		def getTasseledCap(img):
			"""Function to compute the Tasseled Cap transformation and return an image"""
			
			coefficients = ee.Array([
				[0.3037, 0.2793, 0.4743, 0.5585, 0.5082, 0.1863],
				[-0.2848, -0.2435, -0.5436, 0.7243, 0.0840, -0.1800],
				[0.1509, 0.1973, 0.3279, 0.3406, -0.7112, -0.4572],
				[-0.8242, 0.0849, 0.4392, -0.0580, 0.2012, -0.2768],
				[-0.3280, 0.0549, 0.1075, 0.1855, -0.4357, 0.8085],
				[0.1084, -0.9022, 0.4120, 0.0573, -0.0251, 0.0238]
			])
		
			bands=ee.List(['blue','green','red','nir','swir1','swir2'])
			
			# Make an Array Image, with a 1-D Array per pixel.
			arrayImage1D = img.select(bands).toArray()
		
			# Make an Array Image with a 2-D Array per pixel, 6x1.
			arrayImage2D = arrayImage1D.toArray(1)
		
			componentsImage = ee.Image(coefficients).matrixMultiply(arrayImage2D).arrayProject([0]).arrayFlatten([['brightness', 'greenness', 'wetness', 'fourth', 'fifth', 'sixth']]).float()
	  
			# Get a multi-band image with TC-named bands.
			return img.addBands(componentsImage);	
			
			
		def addTCAngles(img):

			""" Function to add Tasseled Cap angles and distances to an image. Assumes image has bands: 'brightness', 'greenness', and 'wetness'."""
			
			# Select brightness, greenness, and wetness bands	
			brightness = img.select('brightness')
			greenness = img.select('greenness')
			wetness = img.select('wetness')
	  
			# Calculate Tasseled Cap angles and distances
			tcAngleBG = brightness.atan2(greenness).divide(math.pi).rename(['tcAngleBG'])
			tcAngleGW = greenness.atan2(wetness).divide(math.pi).rename(['tcAngleGW'])
			tcAngleBW = brightness.atan2(wetness).divide(math.pi).rename(['tcAngleBW'])
			tcDistBG = brightness.hypot(greenness).rename(['tcDistBG'])
			tcDistGW = greenness.hypot(wetness).rename(['tcDistGW'])
			tcDistBW = brightness.hypot(wetness).rename(['tcDistBW'])
			img = img.addBands(tcAngleBG).addBands(tcAngleGW).addBands(tcAngleBW).addBands(tcDistBG).addBands(tcDistGW).addBands(tcDistBW)
			
			return img
	
		img = getTasseledCap(img)
		img = addTCAngles(img)
		return img

	def ND_blue_green(self,img):
		img = img.addBands(img.normalizedDifference(['blue','green']).rename(['ND_blue_green']))
		return img
	
	def ND_blue_red(self,img):
		img = img.addBands(img.normalizedDifference(['blue','red']).rename(['ND_blue_red']))
		return img
	
	def ND_blue_nir(self,img):
		img = img.addBands(img.normalizedDifference(['blue','nir']).rename(['ND_blue_nir']))
		return img
	
	def ND_blue_swir1(self,img):
		img = img.addBands(img.normalizedDifference(['blue','swir1']).rename(['ND_blue_swir1']))
		return img
	
	def ND_blue_swir2(self,img):
		img = img.addBands(img.normalizedDifference(['blue','swir2']).rename(['ND_blue_swir2']))
		return img

	def ND_green_red(self,img):
		img = img.addBands(img.normalizedDifference(['green','red']).rename(['ND_green_red']))
		return img
	
	def ND_green_nir(self,img):
		img = img.addBands(img.normalizedDifference(['green','nir']).rename(['ND_green_nir']))  # NDWBI
		return img
	
	def ND_green_swir1(self,img):
		img = img.addBands(img.normalizedDifference(['green','swir1']).rename(['ND_green_swir1']))  # NDSI, MNDWI
		return img
	
	def ND_green_swir2(self,img):
		img = img.addBands(img.normalizedDifference(['green','swir2']).rename(['ND_green_swir2']))
		return img
		
	def ND_red_swir1(self,img):
		img = img.addBands(img.normalizedDifference(['red','swir1']).rename(['ND_red_swir1']))
		return img
			
	def ND_red_swir2(self,img):
		img = img.addBands(img.normalizedDifference(['red','swir2']).rename(['ND_red_swir2']))
		return img

	def ND_nir_red(self,img):
		img = img.addBands(img.normalizedDifference(['nir','red']).rename(['ND_nir_red']))  # NDVI
		return img
	
	def ND_nir_swir1(self,img):
		img = img.addBands(img.normalizedDifference(['nir','swir1']).rename(['ND_nir_swir1']))  # NDWI, LSWI, -NDBI
		return img
	
	def ND_nir_swir2(self,img):
		img = img.addBands(img.normalizedDifference(['nir','swir2']).rename(['ND_nir_swir2'])) # NBR, MNDVI
		return img

	def ND_swir1_swir2(self,img):
		img = img.addBands(img.normalizedDifference(['swir1','swir2']).rename(['ND_swir1_swir2']))
		return img
  
	def R_swir1_nir(self,img):
		# Add ratios
		img = img.addBands(img.select('swir1').divide(img.select('nir')).rename(['R_swir1_nir']));  # ratio 5/4
		return img
			
	def R_red_swir1(self,img):
		img = img.addBands(img.select('red').divide(img.select('swir1')).rename(['R_red_swir1']));  # ratio 3/5
		return img

	def EVI(self,img):
		#Add Enhanced Vegetation Index (EVI)
		evi = img.expression(
			'2.5 * ((NIR - RED) / (NIR + 6 * RED - 7.5 * BLUE + 1))', {
			  'NIR': img.select('nir'),
			  'RED': img.select('red'),
			  'BLUE': img.select('blue')
		  }).float()
	
		img = img.addBands(evi.rename(['EVI']))

		return img
	  
	def SAVI(self,img):
		# Add Soil Adjust Vegetation Index (SAVI)
		# using L = 0.5;
		savi = img.expression(
			'(NIR - RED) * (1 + 0.5)/(NIR + RED + 0.5)', {
			  'NIR': img.select('nir'),
			  'RED': img.select('red')
		  }).float()
		img = img.addBands(savi.rename(['SAVI']))

		return img
	  
	def IBI(self,img):
		# Add Index-Based Built-Up Index (IBI)
		ibi_a = img.expression(
			'2*SWIR1/(SWIR1 + NIR)', {
			  'SWIR1': img.select('swir1'),
			  'NIR': img.select('nir')
			}).rename(['IBI_A'])
	

		ibi_b = img.expression(
			'(NIR/(NIR + RED)) + (GREEN/(GREEN + SWIR1))', {
			  'NIR': img.select('nir'),
			  'RED': img.select('red'),
			  'GREEN': img.select('green'),
			  'SWIR1': img.select('swir1')
			}).rename(['IBI_B'])
		
		ibi_a = ibi_a.addBands(ibi_b)
		ibi = ibi_a.normalizedDifference(['IBI_A','IBI_B'])
		img = img.addBands(ibi.rename(['IBI']))
		
		return img

	def addTopography(self,img): 
		"""  Function to add 30m SRTM elevation and derived slope, aspect, eastness, and 
		northness to an image. Elevation is in meters, slope is between 0 and 90 deg,
		aspect is between 0 and 359 deg. Eastness and northness are unitless and are
		between -1 and 1. """

		# Import SRTM elevation data
		elevation = ee.Image("USGS/SRTMGL1_003")
		
		# Calculate slope, aspect, and hillshade
		topo = ee.Algorithms.Terrain(elevation)
		
		# From aspect (a), calculate eastness (sin a), northness (cos a)
		deg2rad = ee.Number(math.pi).divide(180)
		aspect = topo.select(['aspect'])
		aspect_rad = aspect.multiply(deg2rad)
		eastness = aspect_rad.sin().rename(['eastness']).float()
		northness = aspect_rad.cos().rename(['northness']).float()
		
		# Add topography bands to image
		topo = topo.select(['elevation','slope','aspect']).addBands(eastness).addBands(northness)
		img = img.addBands(topo)
		return img

	def addJRC(self,img):
		""" Function to add JRC Water layers: 'occurrence', 'change_abs', 
			'change_norm', 'seasonality','transition', 'max_extent' """
		
		jrcImage = ee.Image("JRC/GSW1_0/GlobalSurfaceWater")
		
		img = img.addBands(jrcImage.select(['occurrence']).rename(['occurrence']))
		img = img.addBands(jrcImage.select(['change_abs']).rename(['change_abs']))
		img = img.addBands(jrcImage.select(['change_norm']).rename(['change_norm']))
		img = img.addBands(jrcImage.select(['seasonality']).rename(['seasonality']))
		img = img.addBands(jrcImage.select(['transition']).rename(['transition']))
		img = img.addBands(jrcImage.select(['max_extent']).rename(['max_extent']))
		
		return img


	def getIndices(self,img,covariates):	
		""" add indices to image"""
		# self = indices()
		# no need to add indices that are already there
		# see TODO below, can't use removeDuplicates in .map()
		# indices = self.removeDuplicates(covariates,img.bandNames().getInfo())
		indices = covariates
		
		for item in indices:
			img = self.functionList[item](img)

		return img

	def removeDuplicates(self,covariateList,bands):
		""" function to remove duplicates, i.e. existing bands do not need to be calculated """
		# TODO: this does not scale to being mappable server side (can't use getInfo in mapped functions)
		# would need to EEify this logic to use with ee.List()'s
		return [elem for elem in covariateList if elem not in bands]

	def renameBands(self,image,prefix):
		""" renames bands with prefix """
		
		bandnames = image.bandNames()

		def mapBands(band):
			band = ee.String(prefix).cat('_').cat(band)
			return band
			
		bandnames = bandnames.map(mapBands)
		
		image = image.rename(bandnames)

		return image


def returnCovariates(img):
	"""Workflow for computing Landsat and covariates. bands and covariates are hardcoded inside the function."""
	# hard coded for now
	bands = ['blue','green','red','nir','swir1', 'swir2']	
	bandLow = ['p20_blue','p20_green','p20_red','p20_nir','p20_swir1', 'p20_swir2']
	bandHigh = ['p80_blue','p80_green','p80_red','p80_nir','p80_swir1', 'p80_swir2']

	"""Calculate the urban, builtup cropland rice and barren primitives """
	covariates = ["ND_blue_green","ND_blue_red","ND_blue_nir","ND_blue_swir1","ND_blue_swir2", \
				  "ND_green_red","ND_green_nir","ND_green_swir1","ND_green_swir2","ND_red_swir1",\
				  "ND_red_swir2","ND_nir_red","ND_nir_swir1","ND_nir_swir2","ND_swir1_swir2",\
				  "R_swir1_nir","R_red_swir1","EVI","SAVI","IBI"]
		
	index = indices()
	
	def scaleLandsat(img):
		"""Landast is scaled by factor 0.0001 """
		thermalBand = ee.List(['thermal'])
		thermal = ee.Image(img).select(thermalBand).divide(10)
					
		otherBands = ee.Image(img).bandNames().removeAll(thermalBand)
		
		scaled = ee.Image(img).select(otherBands).multiply(0.0001)
		image = ee.Image(scaled.addBands(thermal))        
	    
		return ee.Image(image.copyProperties(img))

	img = scaleLandsat(img)
		
	def addIndices(img,prefix):
		img = index.addAllTasselCapIndices(img)
		img = index.getIndices(img,covariates)
		img = index.addJRC(img).unmask(0)
		img = index.addTopography(img).unmask(0)	
		if len(prefix) > 0:	
			img = index.renameBands(img,prefix)
		return img
			
		
	down = addIndices(img.select(bandLow,bands),"p20")
	middle = addIndices(img.select(bands),"")
	up = addIndices(img.select(bandHigh,bands),"p80")
		
	img = down.addBands(middle).addBands(up)
	
	return img

# TODO: can't figure out how to detangle getIndices() and addAllTasselCapIndices() 
# so that they can be directly passed as func to hf.Dataset.apply_func()
# so we have this function that takes unnamed kwargs, passed to apply_func()..
def returnCovariatesFromOptions(img,**kwargs):
	"""
	Computes and adds image covariates according to user settings
	args:
		img (ee.Image): image to compute covariates 
		kwargs (dict): a settings dictionary 
	returns:
		img (ee.Image): multi-band image with all desired covariates
	"""
	settings = kwargs
	if 'indices' in kwargs.keys():
		if len(kwargs['indices']) > 0:
			covariates = settings['indices']
			index = indices()
		
		img = ee.Image(img)
		img = index.getIndices(img,covariates)
	
	if 'addTasselCap' in kwargs.keys():
		if kwargs['addTasselCap']:
			img = index.addAllTasselCapIndices(img)
	
	return img