#pragma kernel TexConv2D

#include "Tensor.cginc"

TENSOR_DECL(X)
TENSOR_DECL(K)
TENSOR_DECL(B)
TENSOR_DECL(WBK)
TENSOR_DECL_RW(O)

uint4 _Pad;
uint4 _Stride;

struct TextureAsTensor : Tensor
{
	Texture2D<float4> tex;
	SamplerState smp;

	Texture2DArray<float4> texArray;
	SamplerState smpArray;

	void Init(uint4 nhwc, Texture2D<float4> tex_, SamplerState sampler_, Texture2DArray<float4> texArray_, SamplerState samplerArray_)
	{
		Tensor::Init(nhwc);
		tex = tex_;
		smp = sampler_;
		texArray = texArray_;
		smpArray = samplerArray_;
	}

	float4 Get(uint b, uint y, uint x)
	{
		float3 loc = float3((float)x / (float)width, (float)y / (float)height, b);
		if (batch > 1)
			return texArray.SampleLevel(smpArray, loc, 0);
		else
			return tex.SampleLevel(smp, loc.xy, 0);
	}
};

#define TENSOR_SHARED2_ARGS3(A, B, S, O) TENSOR_SHARED_ARG(A, S); TENSOR_SHARED_ARG(B, S); TENSOR_ARG_RW(O);
Texture2DArray<float4> Xtex2DArray;
Texture2D<float4> Xtex2D;
SamplerState samplerXtex2D { Filter = MIN_MAG_LINEAR_MIP_POINT; AddressU = Clamp; AddressV = Clamp; };
SamplerState samplerXtex2DArray { Filter = MIN_MAG_LINEAR_MIP_POINT; AddressU = Clamp; AddressV = Clamp; };

#define MAX_CHANNELS 4

NUMTHREADS((16,4,4), (16,4,2), (16,2,2))
void TexConv2D(uint3 dispatchThreadID : SV_DispatchThreadID)
{
// @TODO: currently it fails to compile, needs to be investigated
#if 0
	DISPATCH_ARGS(K.kernelCount, O.width, O.height);
	TextureAsTensor X; X.Init(Xdecl[0], Xtex2D, samplerXtex2D, Xtex2DArray, samplerXtex2DArray);

	TENSOR_SHARED_ARG(K, WBK);
	TENSOR_SHARED_ARG(B, WBK);
	TENSOR_ARG_RW(O);

	// ASSERT(X.channels <= MAX_CHANNELS)

	uint k = dispatchThreadID.x;
	uint x = dispatchThreadID.y;
	uint y = dispatchThreadID.z;

	if (k >= K.channels) return;
	if (x >= O.width) return;
	if (y >= O.height) return;

	for (uint n = 0; n < O.batch; ++n)
	{
		float acc = B.Get(k);
		for (uint dy = 0; dy < K.GetKernelHeight(); ++dy)
		{
			for (uint dx = 0; dx < K.GetKernelWidth(); ++dx)
			{
				uint oy = y * _Stride.y + dy;
				uint ox = x * _Stride.x + dx;

				// @TODO: investigate
				// WARNING: had to move both y check into the loop (as opposed to checking y in parent loop) - due to potential bug in Metal compiler
				if (oy < _Pad.y) continue;
				if (oy - _Pad.w >= X.height) continue;
				if (ox < _Pad.x) continue;
				if (ox - _Pad.z >= X.width) continue;

				float4 in4channels = X.Get(n, oy - _Pad.y, ox - _Pad.x);
				for (uint c = 0; c < X.channels && c < MAX_CHANNELS; ++c)
				{
					acc += in4channels[c] * K.Get(dy, dx, c, k);
				}
			}
		}

		O.Set(n, y, x, k, acc);
	}
#endif
}
