// Anime4K Bilateral Mode Denoiser
// Based on Anime4K project by bloc97

//!BGFX EFFECT
//!VERSION 1
//!NAME Anime4K Denoise Bilateral Mode
//!CATEGORY Anime4K
//!DESCRIPTION Bilateral mode filter for noise reduction. Selects the most common value in the neighborhood for optimal edge preservation.


//!PARAMETER
//!LABEL Intensity Sigma
//!DESC Controls denoising strength. Higher values produce stronger smoothing but may blur details.
//!DEFAULT 0.1
//!MIN 0.01
//!MAX 5
//!STEP 0.01
float intensitySigma;

//!TEXTURE
Texture2D INPUT;

//!TEXTURE
//!WIDTH INPUT_WIDTH
//!HEIGHT INPUT_HEIGHT
Texture2D OUTPUT;

//!SAMPLER
//!FILTER POINT
SamplerState sam;


//!PASS 1
//!IN INPUT
//!OUT OUTPUT
//!BLOCK_SIZE 16
//!NUM_THREADS 64

// Intensity window size - higher values give stronger denoising
#define INTENSITY_SIGMA intensitySigma
// Spatial window size - controls blur radius
#define SPATIAL_SIGMA 1.0
// Histogram regularization window size - approximates closest-to-mean filter
#define HISTOGRAM_REGULARIZATION 0.2
// Power curve for intensity weighting
#define INTENSITY_POWER_CURVE 1.0

// Kernel dimension calculations
#define KERNELSIZE uint(max(uint(SPATIAL_SIGMA), 1) * 2 + 1)
#define KERNELHALFSIZE (int(KERNELSIZE/2))
#define KERNELLEN (KERNELSIZE * KERNELSIZE)

#define GETOFFSET(i) int2(int(i % KERNELSIZE) - KERNELHALFSIZE, int(i / KERNELSIZE) - KERNELHALFSIZE)


// Calculate luminance from RGB
float GetLuma(float3 rgba) {
	return dot(float3(0.299, 0.587, 0.114), rgba);
}

// Standard Gaussian function
float GaussianFunc(float x, float s, float m) {
	float scaled = (x - m) / s;
	return exp(-0.5 * scaled * scaled);
}


void Pass1(uint2 blockStart, uint3 threadId) {
	uint2 gxy = (TileSwizzle8x8(threadId.x) << 1) + blockStart;

	const uint2 outputSize = GetOutputSize();
	if (gxy.x >= outputSize.x || gxy.y >= outputSize.y) {
		return;
	}

	float2 inputPt = GetInputPt();
	uint i, j, k, m;

	// Gather source pixels with precomputed luminance
	float4 src[KERNELSIZE + 1][KERNELSIZE + 1];
	[unroll]
	for (i = 0; i <= KERNELSIZE - 1; i += 2) {
		[unroll]
		for (j = 0; j <= KERNELSIZE - 1; j += 2) {
			float2 tpos = (gxy + int2(i, j) - KERNELHALFSIZE + 1) * inputPt;
			const float4 sr = INPUT.GatherRed(sam, tpos);
			const float4 sg = INPUT.GatherGreen(sam, tpos);
			const float4 sb = INPUT.GatherBlue(sam, tpos);

			// Gather returns: w z
			//                 x y
			src[i][j] = float4(sr.w, sg.w, sb.w, GetLuma(float3(sr.w, sg.w, sb.w)));
			src[i][j + 1] = float4(sr.x, sg.x, sb.x, GetLuma(float3(sr.x, sg.x, sb.x)));
			src[i + 1][j] = float4(sr.z, sg.z, sb.z, GetLuma(float3(sr.z, sg.z, sb.z)));
			src[i + 1][j + 1] = float4(sr.y, sg.y, sb.y, GetLuma(float3(sr.y, sg.y, sb.y)));
		}
	}

	// Process 2x2 output block
	[unroll]
	for (i = 0; i <= 1; ++i) {
		[unroll]
		for (j = 0; j <= 1; ++j) {
			const uint2 destPos = gxy + uint2(i, j);

			// Build histogram arrays
			float3 histogram_v[KERNELLEN];
			float histogram_l[KERNELLEN];
			float histogram_w[KERNELLEN];
			float histogram_wn[KERNELLEN];

			// Get center pixel luminance
			float vc = src[KERNELHALFSIZE + i][KERNELHALFSIZE + j].a;

			// Compute sigma values
			float is = pow(vc + 0.0001, INTENSITY_POWER_CURVE) * INTENSITY_SIGMA;
			float ss = SPATIAL_SIGMA;

			// Sample neighborhood and compute initial weights
			[unroll]
			for (k = 0; k < KERNELLEN; k++) {
				const int2 ipos = GETOFFSET(k);
				const uint2 idx = uint2(i, j) + ipos.yx + KERNELHALFSIZE;
				histogram_v[k] = src[idx.x][idx.y].rgb;
				histogram_l[k] = src[idx.x][idx.y].a;
				histogram_w[k] = GaussianFunc(histogram_l[k], is, vc) * GaussianFunc(length(ipos), ss, 0.0);
				histogram_wn[k] = 0.0;
			}

			// Apply histogram regularization
			[unroll]
			for (k = 0; k < KERNELLEN; k++) {
				histogram_wn[k] += GaussianFunc(0.0, HISTOGRAM_REGULARIZATION, 0.0) * histogram_w[k];
				[unroll]
				for (uint m = (k + 1); m < KERNELLEN; m++) {
					float d = GaussianFunc(histogram_l[m], HISTOGRAM_REGULARIZATION, histogram_l[k]);
					histogram_wn[m] += d * histogram_w[k];
					histogram_wn[k] += d * histogram_w[m];
				}
			}

			// Find mode (most weighted value)
			float3 maxv = 0;
			float maxw = 0;

			[unroll]
			for (k = 0; k < KERNELLEN; ++k) {
				if (histogram_wn[k] >= maxw) {
					maxw = histogram_wn[k];
					maxv = histogram_v[k];
				}
			}

			OUTPUT[destPos] = float4(maxv, 1);
		}
	}
}
