'''
Module Name:  Meanshift pseudo code

Created on May 24, 2016

@author: Bo Yang, Moritz Lennert, Markus Metz

GRASS GSoC2016 project, All Rights Reserved
'''

#==================================================================================
#               input parameters
#==================================================================================

NODATA= -3.40282346639e+038
count = 0
range_bandwidth = 27103 #temporally set as the mean value of the image. later apply adaptive
inputRaster= GetParameter(0)
threshold_factor= float(GetParameter(1))#convergence factor, normally 0.01
outputSTARFMraster= GetParameter(2)
windowSize = int(GetParameter(3))


#===================================================================================
#                       input  Raster
#===================================================================================


input_raster=Raster(inputRaster)
#input image

last_iteration_raster = array([[input_raster[j][i] for i in range(input_raster.shape[1])] for j in range(input_raster.shape[0])])
# image sotre for last iteration

iteration_indicator_raster = array([[0 for i in range(input_raster.shape[1])] for j in range(input_raster.shape[0])])
# if this pixel reach convergence, 0 for not, and 1 for convergence

current_iteration_raster = array([[0.0 for i in range(input_raster.shape[1])] for j in range(input_raster.shape[0])])
#current iteration raster

pixel_number = (input_raster.shape[0]-windowSize) * (input_raster.shape[1]-windowSize)
#Count how many pixel are convergence

#===================================================================================
#               gaussian weight difference (S=||B1-A1||)
#===================================================================================

def Mean_shift_filter(Raster=input_raster, w = windowSize):
    for j in range(w/2,input_raster.shape[0]-w/2):
        for i in range(w/2,input_raster.shape[1]-w/2):
			if iteration_indicator_raster[j][i] == 0: #skip those pixel already get convergence
				totalWeight=0.0
				for m in range(w):
					for n in range(w):
						totalWeight = totalWeight + math.exp(-0.5*((input_raster[j][i] - last_iteration_raster[j-w/2+m][i-w/2+n])/(range_bandwidth))**2)
						# pre-calculate the total weight 
				for m1 in range(w):
					for n1 in range(w):
						current_iteration_raster[j][i] = current_iteration_raster[j][i] + input_raster[j-w/2+m1][i-w/2+n1] * (math.exp(-0.5*((input_raster[j][i] - last_iteration_raster[j-w/2+m1][i-w/2+n1])/(range_bandwidth))**2/totalWeight))
			if abs(current_iteration_raster[j][i] - last_iteration_raster[j][i]) < threshold_factor:
				iteration_indicator_raster[j][i] = 1
				count += 1
				#if the pixel convergence

void main():
	while (count < pixel_number):
		Mean_shift_filter()
	saveRaster()