#define WIDTH 512
#define HEIGHT 424

#pragma kernel CSMain

#include "UnityCG.cginc"

RWStructuredBuffer<float4> kinect_motion_buffer;
RWStructuredBuffer<float4> position_buffer;
RWStructuredBuffer<float4> velocity_buffer;

Texture2D<float2> depth_texture;
Texture2D<float2> velocity_texture;

float4x4 kinect_matrix;
float clip;
float3 clip_box;


[numthreads(32, 32, 1)]
void CSMain(uint3 id : SV_DispatchThreadID)
{
	int index = id.y * WIDTH + id.x;
	int high = (int)(depth_texture[id.xy].y * 255.0);
	int low  = (int)(depth_texture[id.xy].x * 255.0);
	float depth = (float)((high << 8) + low) * 0.001;

	float2 vel = velocity_texture[id.xy].xy * 255.0 - 128.0;

	float pre_depth = kinect_motion_buffer[index].z;
	kinect_motion_buffer[index] = float4(vel, depth, pre_depth);

	///////////////////////////////////

	float2 img_size = float2(512.0, 424.0);
	float fovy = 60.0;
	float PI = UNITY_PI;

	float focal = img_size.y / 2.0 / tan(fovy / 2.0 * PI / 180.0);
	float2 vel2d = vel;

	float2 f_pos = float2(id.xy) + 0.5;
	float2 t_pos = f_pos + vel2d;

	float f_depth = pre_depth;
	float t_depth = depth;

	float3 flip = float3(1.0, -1.0, -1.0);
	float3 f_vec = float3((f_pos - 0.5 * img_size) / focal, -1.0) * flip;
	float3 t_vec = float3((t_pos - 0.5 * img_size) / focal, -1.0) * flip;

	float3 f_p3d = float3(f_vec * f_depth);
	float3 t_p3d = float3(t_vec * t_depth);

	float3 vel3d = t_p3d - f_p3d;

	float depth_diff_threshold = 0.14;
	float cam_dot_threshold = 0.955;
	if (f_depth == 0.0 || t_depth == 0.0 || abs(f_depth - t_depth) > depth_diff_threshold)
	{
		f_p3d = float3(0, 0, 0);
		vel3d = float3(0, 0, 0);
	}
	else if (abs(dot(normalize(vel3d), normalize(f_vec))) > cam_dot_threshold)
	{
		f_p3d = float3(0, 0, 0);
		vel3d = float3(0, 0, 0);
	}

	float3 kpos = mul(kinect_matrix, float4(t_p3d, 1.0)).xyz;
	float3 kvel = mul(kinect_matrix, float4(vel3d, 0.0)).xyz;

	if ((kpos.x < -clip_box.x * 0.5 || kpos.x > clip_box.x * 0.5 ||
		kpos.z < -clip_box.z * 0.5 || kpos.z > clip_box.z * 0.5 ||
		kpos.y < 0.0 || kpos.y > clip_box.y) && clip > 0.5)
	{
		kpos = float3(0, 0, 0);
		kvel = float3(0, 0, 0);
	}

	position_buffer[index] = float4(kpos, 1.0);
	velocity_buffer[index] = float4(-kvel, 0.0);
}