# method to calculate the vegetation resilience counts through time 

# NOTES:
#	0 - Not Modeled
#	1 - Black Spruce
#	2 - White Spruce
#	3 - Deciduous
#	4 - Shrub Tundra
#	5 - Graminoid  Tundra
#	6 - Wetland Tundra
#	7 - Barren lichen-moss
#	8 - Temperate Rainforest
#	255 - out of bounds

def get_rep_num( x ):
	'''return rep number from firescar filename'''
	base = os.path.basename( x )
	return base.split('_')[1]
def count_transitions( arr, axis ):
	''' 
	takes an array of integers and returns the number of 
	shifts in values in the series
	array should be 3d and ordered:
	time, x, y [0,1,2]
	'''
	arr = np.diff( arr, axis=axis )
	arr[ arr != 0 ] = 1
	return np.sum( arr, axis=axis )
def veg_resilience( veg_list, axis ):
	''' 
	stacks a list of vegetation rasters into a numpy
	ndarray with 3 dimensions ( time, x, y ) and runs the 
	count_transitions function to return the transition
	counts per pixels in the series.
	'''
	arr = np.dstack([ rasterio.open( i ).read_band( 1 ).data for i in veg_list ])
	arr = np.rollaxis( arr, -1 )
	return count_transitions( arr, axis )

if __name__ == '__main__':
	from itertools import groupby
	import glob, os, sys, re, rasterio, collections
	from scipy.stats import itemfreq
	import pathos.multiprocessing as mp
	import numpy as np
	import scipy as sp
	import argparse

	parser = argparse.ArgumentParser( description='program to calculate Relative Vegetation Change from ALFRESCO Veg outputs' )
	parser.add_argument( '-p', '--input_path', action='store', dest='input_path', type=str, help='path to ALFRESCO output Maps directory' )
	parser.add_argument( '-o', '--output_path', action='store', dest='output_path', type=str, help='path to output directory' )
	parser.add_argument( '-m', '--model', action='store', dest='model', type=str, help='model name' )
	parser.add_argument( '-s', '--scenario', action='store', dest='scenario', type=str, help='scenario' )
	parser.add_argument( '-n', '--ncpus', action='store', dest='ncpus', type=int, help='number of cores to utilize' )
	parser.add_argument( '-by', '--begin_year', action='store', dest='begin_year', type=int, help='beginning year in the range' )
	parser.add_argument( '-ey', '--end_year', action='store', dest='end_year', type=int, help='ending year in the range' )
	
	args = parser.parse_args()

	if not os.path.exists( args.output_path ):
		os.mkdir( args.output_path )

	print( '    running %s : scenario %s' % ( args.model, args.scenario ) )
	
	# list, sort, group by replicate
	veg_list = glob.glob( os.path.join( args.input_path, 'Veg*.tif' ) )
	year_list = range( args.begin_year, args.end_year + 1 )
	veg_list = [ i for i in veg_list if int( os.path.basename( i ).split('_')[ len( os.path.basename( i ).split( '_' ) )-1 ].split( '.' )[0] ) in year_list ]
	veg_sorted = sorted( veg_list, key=lambda x: get_rep_num( x ) )
	veg_grouped = [ list( g ) for k, g in groupby( veg_sorted, key=lambda x: get_rep_num( x ) ) ]
	
	pool = mp.Pool( args.ncpus )
	final = pool.map( lambda x: veg_resilience( x, axis=0 ), veg_grouped )
	pool.close()

	final = np.sum( final, axis=0 ) / np.float( len(veg_list) )
	
	meta = rasterio.open( veg_list[ 0 ] ).meta
	meta.update( compress='lzw', dtype=np.int32, crs={ 'init':'EPSG:3338' } )

	output_filename = os.path.join( args.output_path, 'alfresco_relative_vegetation_change_counts_' + args.model + '_' + args.scenario + '_' + str(args.begin_year) + '_' + str(args.end_year) + '.tif' )

	with rasterio.open( output_filename, 'w', **meta ) as out:
		out.write_band( 1, final2.astype( np.int32 ) )

	del final,out

	# # # example
	# import os
	# models = [ 'cccma_cgcm3_1', 'mpi_echam5' ]
	# scenarios = [ 'sresa1b']
	# year_ranges = [ ('1900', '1999'), ('2000', '2099') ]
	# ncpus = 32
	# for model in models:
	# 	print( 'MODEL : %s' % model )
	# 	for scenario in scenarios:
	# 		for year_range in year_ranges:
	# 			print( '  SCENARIO : %s' % scenario )
	# 			input_path = os.path.join( '/big_scratch/apbennett/Calibration/FinalCalib', model + '.' + scenario, 'Maps' )
	# 			output_path = os.path.join( '/workspace/Shared/Tech_Projects/AK_LandCarbon/project_data/ALFRESCO/relative_vegetation_change_counts', year_range[0] + '_thru_' + year_range[1] )
	# 			if not os.path.exists( output_path ):
	# 				os.mkdir( output_path )
	# 			os.system( 'python /big_scratch/malindgren/ALFRESCO_PostProcessing/alfresco-calibration/alfresco_vegetation_resilience_counts.py -p ' + input_path + ' -o ' + output_path + ' -m ' + model + ' -s ' + scenario + ' -n ' + str( ncpus ) + ' -by ' + year_range[0] + ' -ey ' + year_range[1] )
	
	# # example clip the maps to the iem domain
	# l = glob.glob( '/workspace/Shared/Tech_Projects/Alaska_IEM/project_data/ALFRESCO/relative_vegetation_change_counts/2000_thru_2099/*.tif' )
	# cutline_iem = '/workspace/Shared/Tech_Projects/Alaska_IEM/project_data/ALFRESCO/aiem_domain_fixed/AIEM_domain.shp'

	# # for in_rst in l:
	# 	print 'clipping ' + in_rst
	# 	out_rst = in_rst.replace( '.tif', '_iem.tif' )
	# 	os.system( 'gdalwarp -overwrite -q -tr 1000 1000 -co "COMPRESS=LZW" -cutline ' + cutline_iem + ' -crop_to_cutline -of GTiff ' + in_rst + ' ' + out_rst )

	# # example mask the clipped images to the iem extent ( rasterized )
	# filepath = '/workspace/Shared/Tech_Projects/Alaska_IEM/project_data/ALFRESCO/relative_vegetation_change_counts/2000_thru_2099'
	# files = glob.glob( os.path.join( filepath, '*_iem.tif' ) )

	# mask = rasterio.open( '/workspace/Shared/Tech_Projects/Alaska_IEM/project_data/ALFRESCO/aiem_domain_fixed/AIEM_domain.tif' )
	# mask_arr = mask.read_band( 1 )

	# for f in files:
	# 	rst = rasterio.open( f )
	# 	arr = rst.read_band( 1 )
	# 	meta = rst.meta
	# 	meta.update( compress='lzw', nodata=255 )
	# 	arr[ mask_arr < 1 ] = 255

	# 	masked = np.ma.masked_where( mask_arr < 1, arr )
	# 	output_filename = f.replace( 'iem', 'iem_masked' )

	# 	with rasterio.open( output_filename, 'w', **meta ) as out:
	# 		out.write_band( 1, masked )