#include "Affine2DMatrix.fx"
#include "MatrixDecomposition.fx"

const float2x3 Camera : Camera;
const float2x3 View : View;
const float2x3 Projection : Projection;
const float2   ViewportSize : ViewportSize;

static const float2 UnitSquareVertices[4] = {
	float2(-1, -1),
	float2(1, -1),
	float2(1, 1),
	float2(-1, 1),
};

static const float2 SpriteTextureCoordinates[4] = {
	float2(0, 1),
	float2(1, 1),	
	float2(1, 0),
	float2(0, 0),
};

static const float2x3 Viewport = {
	ViewportSize.x/2, 0,  floor(ViewportSize.x / 2),
	0, -ViewportSize.y/2, floor(ViewportSize.y / 2),
};

///
/// Abreviations:
///	MS - Model Space (origin at center of object)
/// LS - Local Space (same as Model Space)
/// WS - World Space
/// PS - Pixel Space
/// CS - Clip Space

static const float SHAPE_BOX = 1;
static const float SHAPE_CIRCLE = 2;
static const float FIRST_NON_BASIC_PRIMITIVE = 3;
static const float SHAPE_SPRITE = 3;
static const float SHAPE_TRIANGLE = 4;
static const float SHAPE_QUADRATIC_BEZIER_CURVE = 5;

static const float INVALID_TEXTURE_SLOT = 255;

SamplerState s0 : register(s0);
SamplerState s1 : register(s1);
SamplerState s2 : register(s2);
SamplerState s3 : register(s3);
SamplerState s4 : register(s4);
SamplerState s5 : register(s5);
SamplerState s6 : register(s6);
SamplerState s7 : register(s7);
SamplerState s8 : register(s8);
SamplerState s9 : register(s9);
SamplerState s10 : register(s10);
SamplerState s11 : register(s11);
SamplerState s12 : register(s12);
SamplerState s13 : register(s13);
SamplerState s14 : register(s14);
SamplerState s15 : register(s15);

const float4 TextureSizes[16] : TextureSizes;

struct Vertex_IN
{	
	float4 TranslationAndPosition : POSITION0;
	float4 ScaleAndOrientation : POSITION1;

	// Border in pixels, Border Percentage, BorderStyle, Type
	// Vector drawables data:
	// x = Border in pixels
	// y = Border Percentage
	// z = BorderStyle
	// w = Type
	//
	// Sprite (and related type) data:
	// xy = Sprite origin
	// z = unused
	// w = Type
	float4 BorderStyle : TEXCOORD0;
	
	float4 InteriorColor : COLOR0;
	float4 BorderColor : COLOR1;

	// Circle, box user data:
	// x = index of vertex
	// 
	// Triangle user data:
	// (xy) = first point of border plane
	// (zw) = second point of border plane
	//
	// Sprite user data:
	// x = index of vertex
	// (zw) = sprite size in pixels
	float4 UserData_1 : TEXCOORD1;

	// Clip window is a local space Rect.  Pixels can only be drawn inside this window.
	// x = minX clip window
	// y = minY clip window
	// z = maxX clip window
	// w = maxY clip window
	float4 ClipWindowLS : TEXCOORD2;

	// x = border texture slot
	// y = interior texture slot
	// z = sprite texture slot
	float4 TextureSlots : TEXCOORD3;	
};

typedef struct //Vertex_OUT, Pixel_IN
{
	float4 PositionCS : POSITION0;
	
	float4 InteriorColor : COLOR0;
	float4 BorderColor : COLOR1;
	
	// Vector drawables data:
	// x = Border offset in pixel space (not circle or square),
	// y = unused,
	// z = BorderStyle (0 = solid, 1 = smooth),
	// w = Type
	//
	// Sprite (and related type) data:
	// xy = Sprite origin
	// w = Type
	float4 BorderInformation : TEXCOORD0;

	// Circle user data:
	// _1.xy = UV coords in border circle space (interp)
	// _1.zw = UV coords in interior circle space (interp)

	// Box user data:
	// _1.xy = interior box UV coords (interp)

	// Triangle user data:
	// _1.xy = clip space coordinate of border line point A (invariant)
	// _1.zw = clip space coordinate of border line point B (invariant)

	// Sprite user data:
	// _1.x = index of vertex
	// _1.y = channel mask
	// _1.zw = sprite size in pixels
	
	float4 UserData_1 : TEXCOORD1;

	// Clip window is a local space Rect.  Pixels can only be drawn inside this window.
	// x = minX clip window
	// y = minY clip window
	// z = maxX clip window
	// w = maxY clip window
	float4 ClipWindowLS : TEXCOORD2;

	// x = border texture slot
	// y = interior texture slot
	// z = sprite texture slot
	float4 TextureSlots : TEXCOORD3;

	float4 PositionMS : TEXCOORD4;

	// Sprite user data:
	// xy = Interpolated pixel location on sprite to render.
	float4 SpriteLocation : TEXCOORD7;

}Vertex_OUT, Pixel_IN;

struct Pixel_OUT
{
	float4 FinalColor : COLOR;
};

Vertex_OUT MainVS(Vertex_IN IN)
{
	//////////////////////////////////////////////
	//
	//	Setup return information and pull out some
	//	basic information from what we pulled in
	//
	//////////////////////////////////////////////
	Vertex_OUT returnMe;
	{
		returnMe.InteriorColor = IN.InteriorColor;
		returnMe.BorderColor = IN.BorderColor;
		returnMe.BorderInformation = IN.BorderStyle;
		returnMe.PositionCS = float4(0, 0, 0, 1);
		returnMe.UserData_1 = 0;
		returnMe.TextureSlots = IN.TextureSlots;
		returnMe.SpriteLocation = 0;
	}

	float2 WORLD_TRANSLATION = IN.TranslationAndPosition.xy;
	float2 PRIMITIVE_POSITION = IN.TranslationAndPosition.zw;
	float BORDER_RADIUS_PIXEL = IN.BorderStyle.x;
	float BORDER_RADIUS_PERCENT = IN.BorderStyle.y;
	float SHAPE_TYPE = IN.BorderStyle.w;

	float2x3 World = {
		IN.ScaleAndOrientation.xy, WORLD_TRANSLATION.x,
		IN.ScaleAndOrientation.zw, WORLD_TRANSLATION.y,
	};

	float2x3 WCV = Affine2D_Mul(Affine2D_Mul(View, Camera), World);
	
	// Figure out how thick the border needs to be
	if(SHAPE_TYPE == SHAPE_CIRCLE)
	{
		// For a circle, model space vertices are the same as the UV coords we're going to use.

		// Decompose the current WCV transform in to SVD form U * Scale * V
		float2x2 U, V;
		float2 scale;
		{
			float2x2 rot = {
				WCV[0].xy,
				WCV[1].xy,
			};

			SVD_FullDecompose(rot, U, scale, V);
		}

		// Scale the diagonal scale matrix of the SVD by the percentage		
		float2 scalePerc_InPixels = scale * BORDER_RADIUS_PERCENT;

		// Make sure the actual thickness of the border is at least the proper number of pixels
		float borderThickness = max(BORDER_RADIUS_PIXEL, min(scalePerc_InPixels.x, scalePerc_InPixels.y));
		borderThickness = min(min(scale.x, scale.y), borderThickness);

		// Find the least squares UV coords for the inner ellipse that will best match to a uniform thickness
		{
			int index = (int)IN.UserData_1.x;
			returnMe.UserData_1.xy = UnitSquareVertices[index];

			float2x2 newScale = 
			{
				scale.x / (scale.x - borderThickness), 0,
				0, scale.y / (scale.y - borderThickness),
			};

			float2x2 newXForm = mul(mul(transpose(V), newScale), V);
			returnMe.UserData_1.zw = mul(newXForm, UnitSquareVertices[index]);
		}
	}
	else if(SHAPE_TYPE == SHAPE_BOX)
	{
		float3x2 WCV_Transpose = transpose(WCV);
		
		float2x2 axesPS = {
			WCV_Transpose[0],
			WCV_Transpose[1],
		};
		
		float2 lengths = float2(length(axesPS[0]), length(axesPS[1]));
		float2x2 normalizedAxes = {
			axesPS[0] / lengths.x,
			axesPS[1] / lengths.y,
		};
		
		float borderSizePerc = BORDER_RADIUS_PERCENT * min(lengths[0], lengths[1]);
		
		float r = max(borderSizePerc, BORDER_RADIUS_PIXEL);
		
		float m = sqrt((2 * r * r) / (1 - dot(normalizedAxes[0], normalizedAxes[1])));
		m = min(m, min(lengths.x, lengths.y));
		float2 l_inv = 1 - m / lengths;
		float2 l = 1.0/l_inv;

		{
			int index = IN.UserData_1.x;
			returnMe.UserData_1.xy = l * UnitSquareVertices[index];
		}
	}	
	else if(SHAPE_TYPE == SHAPE_TRIANGLE)
	{
		// Border line in model space.
		float2 A = IN.UserData_1.xy;
		float2 B = IN.UserData_1.zw;

		float2 A_PS = Affine2D_MulPoint(Viewport, Affine2D_MulPoint(Projection, Affine2D_MulPoint(WCV, A)));
		float2 B_PS = Affine2D_MulPoint(Viewport, Affine2D_MulPoint(Projection, Affine2D_MulPoint(WCV, B)));

		returnMe.UserData_1.xy = A_PS;
		returnMe.UserData_1.zw = B_PS;

		// By convention, BORDER_RADIUS_PERCENT is actually distance in model space for triangles.
		float borderRadiusDistance_MS = BORDER_RADIUS_PERCENT;
		
		// Assumes uniform scale (nonuniform scale breaks the straight skeleton decomposition anyway).
		float borderRadiusDistance_PS = length(WCV[0]) * borderRadiusDistance_MS;

		returnMe.BorderInformation.x = max(BORDER_RADIUS_PIXEL, borderRadiusDistance_PS);

		if (returnMe.BorderInformation.x == 0)
		{
			// To prevent numerical noise from causing pixels to get 
			// included in the border when there is no border
			returnMe.BorderInformation.x = 1;
		}
	}
	else if (SHAPE_TYPE == SHAPE_SPRITE)
	{
		returnMe.UserData_1 = IN.UserData_1;
		float2 spriteOrigin = IN.BorderStyle.xy;
		returnMe.BorderInformation.xy = spriteOrigin;
		float2 spriteSize = IN.UserData_1.zw;

		int index = (int)IN.UserData_1.x;

		// 3 -- 2
		// |    |
		// 0 -- 1
		float2 uv = SpriteTextureCoordinates[index];

		// Flip any axes that have a negative scale on the view matrix.
		float2 viewAxisSign = sign(float2(View._m00, View._m11));
		uv = uv * viewAxisSign + saturate(-viewAxisSign);
		
		float2 isFarside = uv;
		returnMe.SpriteLocation.xy = spriteSize * isFarside + 1.0/256.0;
	}
	else if (SHAPE_TYPE == SHAPE_QUADRATIC_BEZIER_CURVE)
	{
		returnMe.UserData_1 = IN.UserData_1;
	}
	
	// Clip space position of vertex
	{
		float2 vertexPositionMS = PRIMITIVE_POSITION;
		if (SHAPE_TYPE < FIRST_NON_BASIC_PRIMITIVE)
		{
			int index = IN.UserData_1.x;
			vertexPositionMS += UnitSquareVertices[index];
		}

		returnMe.PositionMS.xy = vertexPositionMS;
		returnMe.PositionMS.z = 0;
		returnMe.PositionMS.w = 0;

		float2 vertexPositionPS = Affine2D_MulPoint(WCV, vertexPositionMS);
		returnMe.PositionCS.xy = Affine2D_MulPoint(Projection, vertexPositionPS);
	}

	// Pixel space position of clip window
	{
		returnMe.ClipWindowLS = IN.ClipWindowLS;
	}
	
	return returnMe;
}

float cross2D(float2 a, float2 b)
{
	return cross(float3(a, 0), float3(b, 0)).z;
}

float4 SampleTexture(float slot, float2 uv)
{
	if (slot == 0)
	{
		return tex2D(s0, uv);
	}
	else if (slot == 1)
	{
		return tex2D(s1, uv);
	}
	else if (slot == 2)
	{
		return tex2D(s2, uv);
	}
	else if (slot == 3)
	{
		return tex2D(s3, uv);
	}
	else if (slot == 4)
	{
		return tex2D(s4, uv);
	}
	else if (slot == 5)
	{
		return tex2D(s5, uv);
	}
	else if (slot == 6)
	{
		return tex2D(s6, uv);
	}
	else if (slot == 7)
	{
		return tex2D(s7, uv);
	}
	else if (slot == 8)
	{
		return tex2D(s8, uv);
	}
	else if (slot == 9)
	{
		return tex2D(s9, uv);
	}
	else if (slot == 10)
	{
		return tex2D(s10, uv);
	}
	else if (slot == 11)
	{
		return tex2D(s11, uv);
	}
	else if (slot == 12)
	{
		return tex2D(s12, uv);
	}
	else if (slot == 13)
	{
		return tex2D(s13, uv);
	}
	else if (slot == 14)
	{
		return tex2D(s14, uv);
	}
	else if (slot == 15)
	{
		return tex2D(s15, uv);
	}
	else
	{
		return 0;
	}
}

float ChannelSelect(float4 sample, int channel, int defaultChannel)
{
	if (channel == 0) // default channel
	{
		return sample[defaultChannel];
	}
	else if (channel <= 1) // _0
	{
		return 0;
	}
	else if (channel >= 6) // _1
	{
		return 1;
	}
	else // one of rgba
	{
		return sample[channel - 2];
	}
}


Pixel_OUT MainPS(Pixel_IN IN, float2 PIXEL_SPACE_POS : VPOS)
{
	Pixel_OUT returnMe;
	
	float isInsideBorder;
	float isInsideInterior;

	// DX SM3 adds a half offset to pixel position, so we need to undo that.
	PIXEL_SPACE_POS = round(PIXEL_SPACE_POS - 0.5);

	clip(float4(IN.PositionMS.xy - IN.ClipWindowLS.xy + 0.5, IN.ClipWindowLS.zw - IN.PositionMS.xy + 0.5));

	float SHAPE_TYPE = IN.BorderInformation.w;

	if (SHAPE_TYPE == SHAPE_CIRCLE)
	{
		// sqrt(x^2 + y^2) <= 1
		isInsideBorder = step(length(IN.UserData_1.xy), 1);

		// sqrt(x^2 + y^2) <= 1
		isInsideInterior = step(length(IN.UserData_1.zw), 1);
	}
	else if (SHAPE_TYPE == SHAPE_BOX)
	{
		// Always inside the quad
		isInsideBorder = 1;

		float2 subBoxPos = IN.UserData_1.xy;

		// |x| <= 1 && |y| <= 1
		float2 isInsideInteriorXY = step(abs(subBoxPos), 1);
		isInsideInterior = isInsideInteriorXY.x * isInsideInteriorXY.y;
	}
	else if(SHAPE_TYPE == SHAPE_TRIANGLE)
	{
		float2 ptA_PS = IN.UserData_1.xy;
		float2 ptB_PS = IN.UserData_1.zw;
		
		// Always inside the triangle since we're a triangle
		isInsideBorder = 1;

		// scaledDistance / lengthAB = signed distance of point to border line.
		float scaledDistance = cross2D(ptA_PS - PIXEL_SPACE_POS, ptB_PS - PIXEL_SPACE_POS);
		float lengthAB = distance(ptA_PS, ptB_PS);

		float borderOffset = IN.BorderInformation.x;

		// scaledDistance >= borderOffset * lengthAB
		isInsideInterior = step(lengthAB * borderOffset, scaledDistance);
	}
	else if (SHAPE_TYPE == SHAPE_QUADRATIC_BEZIER_CURVE)
	{
		// See Resolution Independent Curve Rendering using Programmable Graphics Hardware
		// Charles Loop and Jim Blinn
		// https://www.microsoft.com/en-us/research/wp-content/uploads/2005/01/p1000-loop.pdf
		float3 uvPosition = IN.UserData_1.xyz;
		float aboveOrBelow = uvPosition.z * (uvPosition.x * uvPosition.x - uvPosition.y);
		isInsideInterior = step(aboveOrBelow, 0);
		isInsideBorder = 1;
	}

	int borderTextureSlot = IN.TextureSlots.x;
	int interiorTextureSlot = IN.TextureSlots.y;
	int spriteSheetSlot = IN.TextureSlots.z;

    if (SHAPE_TYPE == SHAPE_SPRITE)
	{
		float2 spriteLocation = IN.SpriteLocation.xy;
		float2 spriteSheetSize = TextureSizes[spriteSheetSlot].xy;
		int index = (int)IN.UserData_1.x;
		float2 spriteSize = IN.UserData_1.zw;
		float2 spriteOrigin = IN.BorderInformation.xy;
		int colorChannelSwizzle = (int)IN.UserData_1.y;
		int textureIndex = IN.UserData_1.x;

		// Extract 3 bit pattern for each color channel:
		// 0 -> use default channel
		// 1 -> 0.0
		// 2 -> spriteSample.r
		// 3 -> spriteSample.g
		// 4 -> spriteSample.b
		// 5 -> spriteSample.a
		// 6 -> 1.0
		float4 masks = float4(512, 64, 8, 1);
		int4 channels = (int4)fmod(floor(colorChannelSwizzle / masks), 8);

		// [0,1] UV coordinates are mapped to [-0.5, Texture Size - 0.5], and then that is rounded to the nearest
		// integer to find which texel to sample.  Which means 0 and 1 uv coordinates are very likely to bleed to
		// neighboring texels.  See https://learn.microsoft.com/en-us/windows/win32/direct3d9/nearest-point-sampling
		//
		// To avoid sampling near 0 and 1, we want to a value epsilon such that sampling the texture at epsilon
		// produces the sample we expect 0 to produce, and likewise for 1 - epsilon, and then clamp the uv coordinates
		// to [episilon, 1 - epsilon].
		//
		// After math, epsilon <= 0.5 / (Texture Size), or just 0.5 if we modify the pixel coordinates in the texture.
		// Though I've chosen 0.25 here just to be conservative.
		//
		// And actually, we're probably not sampling from the full [0,1] range, but from some subset range [lo, hi].
		// The same arguments apply here, though, we want to clamp the uv range to a narrower range by epsilon, or in
		// other words, clamp to [lo + epsilon, hi - epsilon].
		float2 sampleLocation = spriteOrigin + clamp(spriteLocation, 0.25, spriteSize - 0.25);
		float2 uv = sampleLocation / spriteSheetSize;
		float4 spriteSample = SampleTexture(spriteSheetSlot, uv);
		float4 spriteColor = float4(
			ChannelSelect(spriteSample, channels.r, 0),
			ChannelSelect(spriteSample, channels.g, 1),
			ChannelSelect(spriteSample, channels.b, 2),
			ChannelSelect(spriteSample, channels.a, 3)
		);
		returnMe.FinalColor = IN.InteriorColor * spriteColor;

		if (interiorTextureSlot != INVALID_TEXTURE_SLOT)
		{
			returnMe.FinalColor *= SampleTexture(interiorTextureSlot, PIXEL_SPACE_POS);
		}
	}
	else
	{
		float isInsideBorderButNotInterior = isInsideBorder * (1 - isInsideInterior);

		float4 borderColor = IN.BorderColor;
		if (borderTextureSlot != INVALID_TEXTURE_SLOT)
		{
			float2 size = TextureSizes[borderTextureSlot].xy;
			borderColor *= SampleTexture(borderTextureSlot, PIXEL_SPACE_POS / size);
		}

		float4 interiorColor = IN.InteriorColor;
		if (interiorTextureSlot != INVALID_TEXTURE_SLOT)
		{
			float2 size = TextureSizes[borderTextureSlot].xy;
			interiorColor *= SampleTexture(interiorTextureSlot, PIXEL_SPACE_POS / size);
		}

		// TODO: take in to account smooth border
		returnMe.FinalColor = isInsideBorder * lerp(interiorColor, borderColor, isInsideBorderButNotInterior);
	}

	return returnMe;
}

technique Main
{
	pass Main
	{
		VertexShader = compile vs_3_0 MainVS();
		
		PixelShader = compile ps_3_0 MainPS();
	}
}