'''

'''

from commands import getoutput

# WRF domain
dom = [143,209]


#################################################################
#																#
#						 NumPy Stand-Ins						#
#																#
# Condor does not contain the NumPy library so certain 			#
# functions need to be written for the sake of replacing those	#
# otherwise provided by NumPy.									#
#																#
#################################################################


def average(arr):
	# Finds the mean of an array or list of numbers
	return sum(arr)/float(len(arr))


def reshape(arr, size1, size2):
	# Reshapes a 1D array or list into a 2D one
	if len(arr) != size1*size2:
		raise RuntimeError('Array length does not match values...')
	
	new_arr = []
	k = 0
	for i in range(size1):
		x_slice = []
		for j in range(size2):
			x_slice.append(arr[k])
			k+=1
		new_arr.append(x_slice)
		
	return new_arr


def flatten(arr):
	# Reshapes a 2D array or list into a 1D one
	new_arr = []
	for i in range(len(arr)):
		for j in range(len(arr[0])):
			new_arr.append(arr[i][j])
			
	return new_arr


def arr_sum(arr1, arr2):
	# Adds the corresponding elements of two 2D lists to create a new one
	shape = [len(arr1), len(arr1[0])]
	
	sum_arr = []
	a1 = flatten(arr1)
	a2 = flatten(arr2)
	
	for i in range(len(a1)):
		sum_arr.append(a1[i]+a2[i])
	
	return reshape(sum_arr, shape[0], shape[1])


def sum_1D(lst1, lst2):
	# Adds the corresponding elements of two 1D lists to create a new one
	sum_lst = []
	for i in range(len(lst1)):
		sum_lst.append(lst1[i]+lst2[i])
		
	return sum_lst


def moving_avg(list, points):
	# Takes the x-point moving average of a list
	if points%2 == 0:
		padding = (points)/2
	else:
		padding = (points - 1)/2
	
	moving_avgs = []
	for i in range(padding):
		moving_avgs.append(list[i])
	
	for i in range(padding, len(list)-padding):
		
		sample = list[i-padding:i+padding]
		moving_avgs.append(average(sample))
		
	for i in range(padding):
		moving_avgs.append(float('nan'))
			
	return moving_avgs


def no_neg(SL, dSL):
	# Shifts projected snow lines up so that they do not dip below the non-terrain-dependent line
	for i in range(len(dSL)):
		if dSL[i] < SL[i]:
			dSL[i] = SL[i]
			
	return dSL


#################################################################
#																#
#					   Snow Line Alteration						#
#																#
# Functions which are to be used for altering NetCDF files with	#
# new snow variables.  Find the snow line, calculate the new	#
# one, and delete everything south of it.  To execute:			#
#																#
#	Run extract_snow with the met_em file you want to alter.	#
#	Run delta_writer with the SNW_file.txt produced by step 1.	#
#	Run rewrite_snow with the binary_file.txt produced by step	#
#		2 and the met_em file you want to alter.				#
#																#
# For snowless sensitivity experiments, snow_to_0 will set all	#
# snow variables to zero throughout the domain.					#
#																#
#################################################################


def extract_hgt(file):
	# Gets terrain height data from a NetCDF (met_em) file and writes it to a text file 
	ncl_content = '\
	begin							\n\
	in = addfile("%s","r")			\n\
	hgt = in->HGT_M					\n\
	asciiwrite("HGT_file.txt", hgt)	\n\
	end' % file

	f = open("temp_ncl.ncl","w")
	f.write(ncl_content)
	f.close()

	getoutput('ncl temp_ncl.ncl')
	getoutput('rm temp_ncl.ncl')


def extract_LD(file):
	# Gets lake depth data from a NetCDF (met_em) file and writes it to a text file 
	ncl_content = '\
	begin							\n\
	in = addfile("%s","r")			\n\
	ld = in->LAKE_DEPTH				\n\
	asciiwrite("LD_file.txt", ld)	\n\
	end' % file

	f = open("temp_ncl.ncl","w")
	f.write(ncl_content)
	f.close()

	getoutput('ncl temp_ncl.ncl')
	getoutput('rm temp_ncl.ncl')
	

def extract_snow(file):
	# Gets snow water equivalent data from a NetCDF (met_em) file and writes it to a text file 
	ncl_content = '\
	begin							\n\
	in = addfile("%s","r")			\n\
	swe = in->SNOW					\n\
	asciiwrite("SNW_file.txt", swe)	\n\
	end' % file

	f = open("temp_ncl.ncl","w")
	f.write(ncl_content)
	f.close()

	getoutput('ncl temp_ncl.ncl')
	getoutput('rm temp_ncl.ncl')


def overwrite_snow(use_file, t0_file):
	# Overwrites snow values for spin-up days to T-0 day
	if use_file != t0_file:

		ncl_content = '\
		begin														\n\
		in = addfile("%(t0)s","r")									\n\
		out = addfile("%(uf)s","w")									\n\
		swe = in->SNOW												\n\
		snd = in->SNOWH												\n\
		out->SNOW=(/swe/)											\n\
		out->SNOWH=(/snd/)											\n\
		end' % {"t0": t0_file, "uf": use_file}	
		
		f = open("temp_ncl.ncl","w")
		f.write(ncl_content)
		f.close()
		
		g = getoutput('ncl temp_ncl.ncl')
		print g
		

def snoLine(file, threshold=5., lake_fill=True, hgt_threshold=2000):
	# Finds the snow line of a 2D snow dataset via the southern extent of the threshold
	f = open(file, "r")
	raw = map(float, f)
	f.close()
	
	if hgt_threshold < 1e20:
		# Import terrain height data
		f = open("HGT_file.txt", "r")
		hgt = map(float, f)
		f.close()
		hgt = reshape(hgt,dom[0],dom[1])
	
	# Import lake depth data
	f = open("LD_file.txt", "r")
	ld = map(float, f)
	f.close()
	
	for i in range(len(raw)):
		if raw[i] > 6.:
			raw[i] = 6.
		elif lake_fill == True and ld[i] > 10.0:
			raw[i] = 5.
			
	fiveDeg_y = 10#18	# From the bottom to the top of the WRF domain is about 39 degrees
					# Divide that by 143 and you get .27 degrees per y-position
					# So that's about 18 grid spaces per 5 degrees latitude
					
	fiveDeg_x = 12	# Similar process which looks at longitudes at either end of domain
					# at 45 degrees latitude
	
	#sno = array(raw).reshape(dom[0],dom[1])
	sno = reshape(raw,dom[0],dom[1])
	
	snow_line = []
	tru_SL = []
	for x in range(len(sno[0])):
		
		y_slice = []
		hgt_y_slice = []
		for y in range(len(sno)):
			y_slice.append(sno[y][x])
			hgt_y_slice.append(hgt[y][x])
			
		for y in range(len(y_slice)):
			if y > len(y_slice)-fiveDeg_y:
				if y_slice[y] >= threshold:
					snow_line.append(y)
					found = True
					break
			elif y_slice[y] >= threshold and average(y_slice[y:y+fiveDeg_y]) >= threshold and hgt_y_slice[y] < hgt_threshold:
				snow_line.append(y)
				found = True
				break
			else:
				found = False
				
		for y in range(len(y_slice)):
			if y > len(y_slice)-fiveDeg_y:
				if y_slice[y] >= threshold:
					tru_SL.append(y)
					tfound = True
					break
			elif x < 90:
				if y_slice[y] >= threshold and average(y_slice[y:y+1]) >= threshold:
					tru_SL.append(y)
					tfound = True
					break
				else:
					tfound = False
			else:
				if y_slice[y] >= threshold and average(y_slice[y:y+fiveDeg_y]) >= threshold:
					tru_SL.append(y)
					tfound = True
					break
				else:
					tfound = False
				
		if found == False:
			snow_line.append(float('nan'))
		if tfound == False:
			tru_SL.append(float('nan'))
	
	return moving_avg(snow_line, fiveDeg_x), moving_avg(tru_SL, fiveDeg_x)


def delta_writer(alter_file, d_file, threshold=5.):
	# Uses a delta file to alter the snow line of a dataset
	#  Writes information to a binary file of 1's and 0's which are used to multiply to the original data
	f = open(alter_file, "r")
	raw = map(float, f)
	f.close()
	
	f = open(d_file, "r")
	dlt = map(float, f)
	f.close()
	
	f = open("HGT_file.txt", "r")
	hgt = map(float, f)
	f.close()
	hgt = reshape(hgt,dom[0],dom[1])
	
	for i in range(len(raw)):
		if raw[i] > 1e+10:
			raw[i] = 0.
		elif raw[i] > threshold*2:
			raw[i] = threshold*2	
	
	sno = reshape(raw,dom[0],dom[1])
	
	SL, tSL = snoLine(alter_file)
	d_SL = no_neg(SL, sum_1D(tSL, dlt))
	
	binary_file = sno
	for i in range(len(binary_file)):
		for j in range(len(binary_file[i])):
			if binary_file[i][j] > 0.:
				binary_file[i][j] = 1.
	
	for i in range(len(d_SL)):
		if d_SL[i] > -1:
			d_SL[i] = int(d_SL[i])
			
			for j in range(int(d_SL[i])):
				if hgt[j][i] < 2000.:
					binary_file[j][i] = 0.
		
		else:
			# Remove snow beyond boundaries of deltas by border snow extent value
			if i < len(d_SL)/2.:
				for j in d_SL:
					if j > -1:
						y_val = j
						break
			else:
				for j in range(len(d_SL)):
					if d_SL[-j] > -1:
						y_val = d_SL[-j]
						break
			for j in range(int(y_val)):
				binary_file[j][i] = 0.
	
	#flat_binary = binary_file.reshape(dom[0]*dom[1])
	flat_binary = flatten(binary_file)
	f = open("binary_file.txt", "w")
	for i in flat_binary:
		f.write(str(i))
		f.write('\n')
	f.close()


def rewrite_snow(file):
	# Writes new snow data to NetCDF file by multiplying snow mass and depth by binary values
	ncl_content = '\
	begin														\n\
	in = addfile("%s","w")										\n\
	swe = in->SNOW												\n\
	snd = in->SNOWH												\n\
	flags = asciiread("binary_file.txt",(/143,209/),"float")	\n\
	sizes = dimsizes(swe)										\n\
	SWE = new(sizes,float)										\n\
	SND = new(sizes,float)										\n\
	SWE(0,:,:) = swe(0,:,:)*flags								\n\
	SND(0,:,:) = snd(0,:,:)*flags								\n\
	in->SNOW=(/SWE/)											\n\
	in->SNOWH=(/SND/)											\n\
	asciiwrite("dSNW_file.txt", SWE)							\n\
	end' % file
	
	f = open("temp_ncl.ncl","w")
	f.write(ncl_content)
	f.close()
	
	getoutput('ncl temp_ncl.ncl')
	getoutput('rm temp_ncl.ncl; rm binary_file.txt')


def snow_to_0(file):
	# Rewrites values for snow depth and mass to zeros (for sensitivity testing)
	ncl_content = '\
	begin														\n\
	in = addfile("%s","w")										\n\
	swe = in->SNOW												\n\
	snd = in->SNOWH												\n\
	sizes = dimsizes(swe)										\n\
	SWE = new(sizes,float)										\n\
	SND = new(sizes,float)										\n\
	SWE(:,:,:) = 0												\n\
	SND(:,:,:) = 0												\n\
	in->SNOW=(/SWE/)											\n\
	in->SNOWH=(/SND/)											\n\
	end' % file
	
	f = open("temp_ncl.ncl","w")
	f.write(ncl_content)
	f.close()
	
	getoutput('ncl temp_ncl.ncl')
	getoutput('rm temp_ncl.ncl')