Shader "LeviathanVideo/YUV420P_3D_ChromaKey"
{
    Properties
    {
        _MainTex ("YTexture", 2D) = "white" {}
        _UTex("UTexture",2D)="white"{}
        _VTex("VTexture",2D)="white"{}
        [Toggle]_ReverseY("Reverse Y",Float)=1
        
        _Color ("Tint", Color) = (1,1,1,1)
        
        _KeyColor("KeyColor", Color) = (0,1,0,0)
        _ColorCutoff("Cutoff", Range(0, 1)) = 0.2
        _ColorFeathering("ColorFeathering", Range(0, 1)) = 0.33
        _MaskFeathering("MaskFeathering", Range(0, 1)) = 1
        _Sharpening("Sharpening", Range(0, 1)) = 0.5
        
        [Enum(UnityEngine.Rendering.CullMode)] _Cull ("Cull Mode", Float) = 2
        [Enum(Off, 0, On, 1)] _ZWrite ("ZWrite", Float) = 1
        [Enum(UnityEngine.Rendering.CompareFunction)] _ZTest ("ZTest", Float) = 4
    }

    SubShader
    {
        Tags
        {
            "Queue"="Transparent"
            "IgnoreProjector"="True"
            "RenderType"="Transparent"
        }

        Cull [_Cull]
        Lighting Off
        ZWrite [_ZWrite]
        ZTest [_ZTest]
        Blend SrcAlpha OneMinusSrcAlpha

        Pass
        {
            Name "Default"
        CGPROGRAM
            #pragma vertex vert
            #pragma fragment frag
            #pragma target 2.0

            #include "UnityCG.cginc"

            //Rec.709
            static const half3x3 YUV_TO_RGB = half3x3(
                1.16438,  1.16438,  1.16438,
                0.0,      -0.21325, 2.11240,
                1.79274,  -0.53291,     0.0
            );
            static const float FLOAT_16_255 = 16.0 / 255.0;
            static const float FLOAT_128_255 = 128.0 / 255.0;

            struct appdata_t
            {
                float4 vertex   : POSITION;
                float4 color    : COLOR;
                float2 texcoord : TEXCOORD0;
                UNITY_VERTEX_INPUT_INSTANCE_ID
            };

            struct v2f
            {
                float4 vertex   : SV_POSITION;
                fixed4 color    : COLOR;
                float2 texcoord  : TEXCOORD0;
                UNITY_VERTEX_OUTPUT_STEREO
            };

            sampler2D _MainTex,_UTex,_VTex;
            fixed4 _Color;
            float4 _MainTex_ST;
            half _ReverseY;

            float4 _MainTex_TexelSize;
            float4 _KeyColor;
            float _ColorCutoff;
            float _ColorFeathering;
            float _MaskFeathering;
            float _Sharpening;

            float rgb2y(float3 c) 
            {
                return (0.299*c.r + 0.587*c.g + 0.114*c.b);
            }

            float rgb2cb(float3 c) 
            {
                return (0.5 + -0.168736*c.r - 0.331264*c.g + 0.5*c.b);
            }

            float rgb2cr(float3 c) 
            {
                return (0.5 + 0.5*c.r - 0.418688*c.g - 0.081312*c.b);
            }

            float colorclose(float Cb_p, float Cr_p, float Cb_key, float Cr_key, float tola, float tolb)
            {
                float temp = (Cb_key-Cb_p)*(Cb_key-Cb_p)+(Cr_key-Cr_p)*(Cr_key-Cr_p);
                float tola2 = tola*tola;
                float tolb2 = tolb*tolb;
                if (temp < tola2) return (0);
                if (temp < tolb2) return (temp-tola2)/(tolb2-tola2);
                return (1);
            }

            half3 getRGB(sampler2D tex1, sampler2D tex2, sampler2D tex3, float2 uv)
            {
                half3 c = 0;
                c.r = tex2D(tex1, uv).a - FLOAT_16_255;
                c.g = tex2D(tex2, uv).a - FLOAT_128_255;
                c.b = tex2D(tex3, uv).a - FLOAT_128_255;
                c = mul(c, YUV_TO_RGB);
                return c;
            }
            
            float maskedTex2D(sampler2D tex1, sampler2D tex2, sampler2D tex3, float2 uv)
            {
                float4 color = float4(getRGB(tex1, tex2, tex3, uv), 1.0);
                
                // Chroma key to CYK conversion
                float key_cb = rgb2cb(_KeyColor.rgb);
                float key_cr = rgb2cr(_KeyColor.rgb);
                float pix_cb = rgb2cb(color.rgb);
                float pix_cr = rgb2cr(color.rgb);

                return colorclose(pix_cb, pix_cr, key_cb, key_cr, _ColorCutoff, _ColorFeathering);
            }

            v2f vert(appdata_t v)
            {
                v2f OUT;
                UNITY_SETUP_INSTANCE_ID(v);
                UNITY_INITIALIZE_VERTEX_OUTPUT_STEREO(OUT);
                OUT.vertex = UnityObjectToClipPos(v.vertex);
                OUT.texcoord = TRANSFORM_TEX(v.texcoord.xy, _MainTex);
                if (_ReverseY > 0)
                {
                    OUT.texcoord.y = 1 - OUT.texcoord.y;
                }
                OUT.color = v.color * _Color;
                return OUT;
            }

            fixed4 frag(v2f IN) : SV_Target
            {
                half4 color = half4(getRGB(_MainTex, _UTex, _VTex, IN.texcoord), IN.color.a);
                color.rgb = saturate(color * _Color);
#ifdef UNITY_COLORSPACE_GAMMA
#else
                color.rgb = pow(color.rgb, 2.2);
#endif

                // Get pixel width
                float2 pixelWidth = float2(1.0 / _MainTex_TexelSize.z, 0);
                float2 pixelHeight = float2(0, 1.0 / _MainTex_TexelSize.w);

                // Unfeathered mask
                float mask = maskedTex2D(_MainTex, _UTex, _VTex, IN.texcoord);

                // Feathering & smoothing
                float c = mask;
                float r = maskedTex2D(_MainTex, _UTex, _VTex, IN.texcoord + pixelWidth);
                float l = maskedTex2D(_MainTex, _UTex, _VTex, IN.texcoord - pixelWidth);
                float d = maskedTex2D(_MainTex, _UTex, _VTex, IN.texcoord + pixelHeight); 
                float u = maskedTex2D(_MainTex, _UTex, _VTex, IN.texcoord - pixelHeight);
                float rd = maskedTex2D(_MainTex, _UTex, _VTex, IN.texcoord + pixelWidth + pixelHeight) * .707;
                float dl = maskedTex2D(_MainTex, _UTex, _VTex, IN.texcoord - pixelWidth + pixelHeight) * .707;
                float lu = maskedTex2D(_MainTex, _UTex, _VTex, IN.texcoord - pixelHeight - pixelWidth) * .707;
                float ur = maskedTex2D(_MainTex, _UTex, _VTex, IN.texcoord + pixelWidth - pixelHeight) * .707;
                float blurContribution = (r + l + d + u + rd + dl + lu + ur + c) * 0.12774655;
                float smoothedMask = smoothstep(_Sharpening, 1, lerp(c, blurContribution, _MaskFeathering));
                float4 result = color * smoothedMask;
                
                return float4(result.xyz, smoothedMask) * IN.color;
            }
        ENDCG
        }
    }
}