uniform float4x4 ViewProj;
uniform texture2d image;

uniform float4 cb_v4 = { -0.100644, -0.338572,  0.439216, 0.501961 };
uniform float4 cr_v4 = {  0.439216, -0.398942, -0.040274, 0.501961 };

uniform float opacity;
uniform float contrast;
uniform float brightness;
uniform float gamma;

uniform float2 chroma_key;
uniform float2 pixel_size;
uniform float similarity;
uniform float smoothness;
uniform float spill;

sampler_state textureSampler {
	Filter    = Linear;
	AddressU  = Clamp;
	AddressV  = Clamp;
};

struct VertData {
	float4 pos : POSITION;
	float2 uv  : TEXCOORD0;
};

VertData VSDefault(VertData v_in)
{
	VertData vert_out;
	vert_out.pos = mul(float4(v_in.pos.xyz, 1.0), ViewProj);
	vert_out.uv  = v_in.uv;
	return vert_out;
}

float4 CalcColor(float4 rgba)
{
	return float4(pow(rgba.rgb, float3(gamma, gamma, gamma)) * contrast + brightness, rgba.a);
}

float GetChromaDist(float3 rgb)
{
	float cb = dot(rgb.rgb, cb_v4.xyz) + cb_v4.w;
	float cr = dot(rgb.rgb, cr_v4.xyz) + cr_v4.w;
	return distance(chroma_key, float2(cb, cr));
}

float GetNonlinearChannel(float u)
{
	return (u <= 0.0031308) ? (12.92 * u) : ((1.055 * pow(u, 1.0 / 2.4)) - 0.055);
}

float3 GetNonlinearColor(float3 rgb)
{
	return float3(GetNonlinearChannel(rgb.r), GetNonlinearChannel(rgb.g), GetNonlinearChannel(rgb.b));
}

float3 SampleTexture(float2 uv)
{
	float3 rgb = image.Sample(textureSampler, uv).rgb;
	return GetNonlinearColor(rgb);
}

float GetBoxFilteredChromaDist(float3 rgb, float2 texCoord)
{
	float2 h_pixel_size = pixel_size / 2.0;
	float2 point_0 = float2(pixel_size.x, h_pixel_size.y);
	float2 point_1 = float2(h_pixel_size.x, -pixel_size.y);
	float distVal = GetChromaDist(SampleTexture(texCoord-point_0));
	distVal += GetChromaDist(SampleTexture(texCoord+point_0));
	distVal += GetChromaDist(SampleTexture(texCoord-point_1));
	distVal += GetChromaDist(SampleTexture(texCoord+point_1));
	distVal *= 2.0;
	distVal += GetChromaDist(GetNonlinearColor(rgb));
	return distVal / 9.0;
}

float4 ProcessChromaKey(float4 rgba, VertData v_in)
{
	float chromaDist = GetBoxFilteredChromaDist(rgba.rgb, v_in.uv);
	float baseMask = chromaDist - similarity;
	float fullMask = pow(saturate(baseMask / smoothness), 1.5);
	float spillVal = pow(saturate(baseMask / spill), 1.5);

	rgba.a *= opacity;
	rgba.a *= fullMask;

	float desat = dot(rgba.rgb, float3(0.2126, 0.7152, 0.0722));
	rgba.rgb = lerp(float3(desat, desat, desat), rgba.rgb, spillVal);

	return CalcColor(rgba);
}

float4 PSChromaKeyRGBA(VertData v_in) : TARGET
{
	float4 rgba = image.Sample(textureSampler, v_in.uv);
	rgba.rgb = max(float3(0.0, 0.0, 0.0), rgba.rgb / rgba.a);
	rgba = ProcessChromaKey(rgba, v_in);
	rgba.rgb *= rgba.a;
	return rgba;
}

technique Draw
{
	pass
	{
		vertex_shader = VSDefault(v_in);
		pixel_shader  = PSChromaKeyRGBA(v_in);
	}
}
