#ifndef MEDIANFILTERKERNEL_H__
#define MEDIANFILTERKERNEL_H__

#define PixelType float


// Texture that holds the input image
//----------------------------------------------------------------------------
texture<float, 2, cudaReadModeElementType> InputImageTexture;


// Division. If division remainder is neq zero then the result is ceiled
//----------------------------------------------------------------------------
#define DivUp(a,b) ((a % b != 0) ? (a/b + 1) : (a/b))


// Linear (M) and 2D (S) access to the shared memory. Kernel radius is automatically incorporated
//----------------------------------------------------------------------------
#define M(i)	SharedMem[tv+i%(KernelRadius*2+1)][tu+i/(KernelRadius*2+1)]
#define S(u,v)	SharedMem[v+KernelRadius][u+KernelRadius]


// Defines for the 3x3 kernel
//----------------------------------------------------------------------------
#define s2(a,b)            {float tmp = a; a = min(a,b); b = max(tmp,b);}
#define mn3(a,b,c)         s2(a,b); s2(a,c);
#define mx3(a,b,c)         s2(b,c); s2(a,c);
#define mnmx3(a,b,c)       mx3(a,b,c); s2(a,b);                               // 3 exchanges
#define mnmx4(a,b,c,d)     s2(a,b); s2(c,d); s2(a,c); s2(b,d);                // 4 exchanges
#define mnmx5(a,b,c,d,e)   s2(a,b); s2(c,d); mn3(a,c,e); mx3(b,d,e);          // 6 exchanges
#define mnmx6(a,b,c,d,e,f) s2(a,d); s2(b,e); s2(c,f); mn3(a,b,c); mx3(d,e,f); // 7 exchanges


// Defines for the 5x5 kernel
//----------------------------------------------------------------------------
#define t2(a, b)								s2(w[a], w[b]);
#define t24(a, b, c, d, e, f, g, h)				t2(a, b); t2(c, d); t2(e, f); t2(g, h); 
#define t25(a, b, c, d, e, f, g, h, i, j)		t24(a, b, c, d, e, f, g, h); t2(i, j);



// Brute force implementation
//----------------------------------------------------------------------------
__global__ void
d_MedianFilter(PixelType* id, PixelType* od, int w, int h, int KernelRadius)
{
	// Compute kernel size
	int KernelSize = (KernelRadius*2+1)*(KernelRadius*2+1);

	// Set up kernels (they're allocated from extern)
	extern __shared__ float window[];

	// Get position in flattened image
	int pos = blockIdx.x*blockDim.x + threadIdx.x;

	// Outside => return
	if ( pos >= w*h )
		return;

	// Position
	int y = (int)(pos/w);
	int x = (int)(pos%w);
	 
	// Thread id
	unsigned int tid=threadIdx.y*blockDim.y+threadIdx.x;

	// Boundary check
	if ( y-KernelRadius < 0 || y+KernelRadius >= h || x-KernelRadius < 0 || x+KernelRadius >= w )
	{
		od[y*w + x] = id[y*w + x] ;
		return;
	}
	
	// Initialize the kernel
	int KernelPos = 0;
	for ( int row = -KernelRadius; row <= KernelRadius; row++ )
	{
		for ( int col = -KernelRadius; col <= KernelRadius; col++ )
		{
			//window[tid*KernelSize + KernelPos] = id[(y+row)*w + (x+col)];
			window[tid*KernelSize + KernelPos] = tex2D(InputImageTexture,x+col,y+row);
			KernelPos++;
		}
	}
 
	// Order elements (only half of them)
	int pivot = ceil(KernelSize/2.f);
	for ( unsigned int j = 0; j < pivot; j++ )
	{
		// Find position of minimum element
		int min = j;
		for ( unsigned int l = j+1; l < KernelSize; l++ )
		{
			if ( window[tid*KernelSize + l] < window[tid*KernelSize + min] )
				min = l;
		}
				
		// Put found minimum element in its place
		const float temp =  window[tid*KernelSize + j];
		window[tid*KernelSize + j] = window[tid*KernelSize + min];
		window[tid*KernelSize + min] = temp;
	}

	// Set the output
	od[y*w + x]=window[tid*KernelSize + pivot-1];
}


// Adaptive filter size kernel
//----------------------------------------------------------------------------
template<int KernelRadius, unsigned int BlockSizeX, unsigned int BlockSizeY>
__global__ void
MedianFilterKernel(unsigned int NU, unsigned int NV, float* Output, int NormalizedOutputPitch)
{
	// 2D index and linear index within this thread block
	int tu = threadIdx.x;
	int tv = threadIdx.y;
	
	// Global 2D index and linear index.
	int gu = blockIdx.x*BlockSizeX+tu;
	int gv = blockIdx.y*BlockSizeY+tv;
	int gl = gv*NormalizedOutputPitch + gu;


	// Note that we cannot use the following termination criteria as threads processing these
	// regions might have to pull data
	/*if ( gu >= NU || gv >= NV )
		return;*/

	// Shared memory. We're using matrix and NOT texture notation (i.e. row first).
	// This scheme avoids bank conflicts
	const int SharedMemSizeX = BlockSizeX+2*KernelRadius;
	const int SharedMemSizeY = BlockSizeY+2*KernelRadius;
	__shared__ float SharedMem[SharedMemSizeX][SharedMemSizeY];

	// Successively pull data. First along the v-direction
	int sv = tv; // Index in shared memory
	int v = gv-KernelRadius; // Index in global memory
	for ( int j = 0; j < DivUp(SharedMemSizeY,BlockSizeY); j++ )
	{
		// Termination criteria
		if ( sv >= SharedMemSizeY )
			break; 

		// Loop over the u-direction
		int su = tu; // Index in shared memory
		int u = gu-KernelRadius; // Index in global memory
		for ( int i = 0; i < DivUp(SharedMemSizeX,BlockSizeX); i++ )
		{
			// Termination criteria
			if ( su >= SharedMemSizeX  )
				break; 

			// Pull data for this thread
			SharedMem[sv][su] = tex2D(InputImageTexture,u,v);

			// Update u-coordinate in shared and global memory
			su += BlockSizeX;
			u += BlockSizeX;
		}

		// Update v-coordinate in shared memory
		sv += BlockSizeY;
		v += BlockSizeY;
	}

	// Wait until all threads have pulled their data
	__syncthreads();
	
	// These threads do not process a pixel
	if ( gl >= (NU-1)*NormalizedOutputPitch+NV ) // if ( gu >= NX || gv >= NY )
		return;
	/*Output[gl] = SharedMem[tv+KernelRadius][tu+KernelRadius];
	return;*/

	// Depending on the kernel radius we apply some special variants
	if ( KernelRadius == 1 )
	{
		// Exchange trick: Morgan McGuire, ShaderX 2008

		// Initialize the window
		float w[6] = 
		{ 
			S(tu-1,	tv-1), 
			S(tu,	tv-1), 
			S(tu+1,	tv-1),       
			S(tu-1,	tv  ), 
			S(tu,	tv  ), 
			S(tu+1,	tv  ) 
		};

		// Median search
		mnmx6(w[0], w[1], w[2], w[3], w[4], w[5]);
		w[5] = S(tu-1, tv+1);
		mnmx5(w[1], w[2], w[3], w[4], w[5]);
		w[5] = S(tu  , tv+1);
		mnmx4(w[2], w[3], w[4], w[5]);
		w[5] = S(tu+1, tv+1);
		mnmx3(w[3], w[4], w[5]);

		// pick the middle one
		Output[gl] = w[4];
	}
	else if ( KernelRadius == 2 )
	{
		// Exchange trick: Morgan McGuire, ShaderX 2008
		// http://graphics.cs.williams.edu/papers/MedianShaderX6/median5.pix

		// Initialize the window
		float w[25];
		w[0]	= S(tu-2,	tv-2); // 1st row
		w[1]	= S(tu-1,	tv-2);
		w[2]	= S(tu,		tv-2);
		w[3]	= S(tu+1,	tv-2);
		w[4]	= S(tu+2,	tv-2); 
		w[5]	= S(tu-2,	tv-1); // 2nd row
		w[6]	= S(tu-1,	tv-1);
		w[7]	= S(tu,		tv-1);
		w[8]	= S(tu+1,	tv-1);
		w[9]	= S(tu+2,	tv-1); 
		w[10]	= S(tu-2,	tv  ); // 3rd row
		w[11]	= S(tu-1,	tv  );
		w[12]	= S(tu,		tv  );
		w[13]	= S(tu+1,	tv  );
		w[14]	= S(tu+2,	tv  ); 
		w[15]	= S(tu-2,	tv+1); // 4th row
		w[16]	= S(tu-1,	tv+1);
		w[17]	= S(tu,		tv+1);
		w[18]	= S(tu+1,	tv+1);
		w[19]	= S(tu+2,	tv+1); 
		w[20]	= S(tu-2,	tv+2); // 5th row
		w[21]	= S(tu-1,	tv+2);
		w[22]	= S(tu,		tv+2);
		w[23]	= S(tu+1,	tv+2);
		w[24]	= S(tu+2,	tv+2); 
		
		// Median search
		t25(0, 1,			3, 4,		2, 4,		2, 3,		6, 7);
		t25(5, 7,			5, 6,		9, 7,		1, 7,		1, 4);
		t25(12, 13,		11, 13,		11, 12,		15, 16,		14, 16);
		t25(14, 15,		18, 19,		17, 19,		17, 18,		21, 22);
		t25(20, 22,		20, 21,		23, 24,		2, 5,		3, 6);
		t25(0, 6,			0, 3,		4, 7,		1, 7,		1, 4);
		t25(11, 14,		8, 14,		8, 11,		12, 15,		9, 15);
		t25(9, 12,		13, 16,		10, 16,		10, 13,		20, 23);
		t25(17, 23,		17, 20,		21, 24,		18, 24,		18, 21);
		t25(19, 22,		8, 17,		9, 18,		0, 18,		0, 9);
		t25(10, 19,		1, 19,		1, 10,		11, 20,		2, 20);
		t25(2, 11,		12, 21,		3, 21,		3, 12,		13, 22);
		t25(4, 22,		4, 13,		14, 23,		5, 23,		5, 14);
		t25(15, 24,		6, 24,		6, 15,		7, 16,		7, 19);
		t25(3, 11,		5, 17,		11, 17,		9, 17,		4, 10);
		t25(6, 12,		7, 14,		4, 6,		4, 7,		12, 14);
		t25(10, 14,		6, 7,		10, 12,		6, 10,		6, 17);
		t25(12, 17,		7, 17,		7, 10,		12, 18,		7, 12);
		t24(10, 18,		12, 20,		10, 20,		10, 12);

		// Pick the middle one
		Output[gl] = w[12];
	}
	else
	{
		// Median search without modifying elements
		int n = (KernelRadius*2+1)*(KernelRadius*2+1);
		int i, less, greater, equal;
		
		float  min, max, guess, maxltguess, mingtguess;
		min = max = M(0);

		for ( i = 1 ; i < n ; i++ ) 
		{
			if ( M(i) < min ) 
				min = M(i);
			if ( M(i) > max ) 
				max = M(i);
		}

		while (1) 
		{
			guess = (min+max)/2;
			less = 0; greater = 0; equal = 0;
			maxltguess = min ;
			mingtguess = max ;
			for ( i = 0; i < n; i++ ) 
			{
				if ( M(i) < guess) 
				{
					less++;
					if ( M(i) > maxltguess ) 
						maxltguess = M(i);
				} 
				else if ( M(i) > guess ) 
				{
					greater++;
					if ( M(i) < mingtguess ) 
						mingtguess = M(i);
				} 
				else
				{
					equal++;
				}
			}

			if ( less <= (n+1)/2 && greater <= (n+1)/2 ) 
			{
				break; 
			}
			else if ( less>greater ) 
			{
				max = maxltguess ;
			}
			else 
			{
				min = mingtguess;
			}
		}

		if ( less >= (n+1)/2 ) 
		{
			Output[gl] = maxltguess;
		}
		else if ( less+equal >= (n+1)/2 ) 
		{
			Output[gl] = guess;
		}
		else
		{
			Output[gl] = mingtguess;
		}
	}
}


#if 0
else if ( KernelRadius == 3 )
	{
		// Number of elements to scan
		int N = (KernelRadius*2+1)*(KernelRadius*2+1);

		// Find minimum and maximum
		float min = 1e16;
		float max =  -1e16;
		for ( int i = 0; i < N ; i++ ) 
		{
			if ( M(i) < min ) 
				min = M(i);
			if ( M(i) > max ) 
				max = M(i);
		}

		// Iterate
		float LowValue = max;
		float HighValue = min;
		float Median = 0;
		int MedianIndex = (N+1)/2;
		while ( HighValue < LowValue )
		{
			// Estimate median 
			Median = (HighValue + LowValue)/2;

			// Count the number of elements leq to the approximated median
			int K = 0;
			for ( int i = 0; i < N; i++ )
			{
				if ( M(i) <= Median )
					K++;
			}

			// Termination criterion or update of low and high value
			if ( K == MedianIndex )
			{
				break;
			}
			else if ( K > MedianIndex )
			{
				HighValue = Median;
			}
			else
			{
				LowValue = Median;
			}
		}

		// Store the output
		Output[gl] = Median;
	}
#endif

#endif // MEDIANFILTERKERNEL_H__