using System;
using UnityEngine;

namespace UnityStandardAssets.ImageEffects
{
    [ExecuteInEditMode]
    [RequireComponent (typeof(Camera))]
    [AddComponentMenu ("Image Effects/Camera/Depth of Field (Lens Blur, Scatter, DX11)") ]
    public class DepthOfField : PostEffectsBase {

        public bool  visualizeFocus = false;
        public float focalLength = 10.0f;
        public float focalSize = 0.05f;
        public float aperture = 0.5f;
        public Transform focalTransform = null;
        public float maxBlurSize = 2.0f;
        public bool  highResolution = false;

        public enum BlurType {
            DiscBlur = 0,
            DX11 = 1,
        }

        public enum BlurSampleCount {
            Low = 0,
            Medium = 1,
            High = 2,
        }

        public BlurType blurType = BlurType.DiscBlur;
        public BlurSampleCount blurSampleCount = BlurSampleCount.High;

        public bool  nearBlur = false;
        public float foregroundOverlap = 1.0f;

        public Shader dofHdrShader;
        private Material dofHdrMaterial = null;

        public Shader dx11BokehShader;
        private Material dx11bokehMaterial;

        public float dx11BokehThreshold = 0.5f;
        public float dx11SpawnHeuristic = 0.0875f;
        public Texture2D dx11BokehTexture = null;
        public float dx11BokehScale = 1.2f;
        public float dx11BokehIntensity = 2.5f;

        private float focalDistance01 = 10.0f;
        private ComputeBuffer cbDrawArgs;
        private ComputeBuffer cbPoints;
        private float internalBlurWidth = 1.0f;

        private Camera cachedCamera;

        public override bool CheckResources () {
            CheckSupport (true); // only requires depth, not HDR

            dofHdrMaterial = CheckShaderAndCreateMaterial (dofHdrShader, dofHdrMaterial);
            if (supportDX11 && blurType == BlurType.DX11) {
                dx11bokehMaterial = CheckShaderAndCreateMaterial(dx11BokehShader, dx11bokehMaterial);
                CreateComputeResources ();
            }

            if (!isSupported)
                ReportAutoDisable ();

            return isSupported;
        }

        void OnEnable () {
            cachedCamera = GetComponent<Camera>();
            cachedCamera.depthTextureMode |= DepthTextureMode.Depth;
        }

        void OnDisable () {
            ReleaseComputeResources ();

            if (dofHdrMaterial) DestroyImmediate(dofHdrMaterial);
            dofHdrMaterial = null;
            if (dx11bokehMaterial) DestroyImmediate(dx11bokehMaterial);
            dx11bokehMaterial = null;
        }

        void ReleaseComputeResources () {
            if (cbDrawArgs != null) cbDrawArgs.Release();
            cbDrawArgs = null;
            if (cbPoints != null) cbPoints.Release();
            cbPoints = null;
        }

        void CreateComputeResources () {
            if (cbDrawArgs == null)
            {
                cbDrawArgs = new ComputeBuffer (1, 16, ComputeBufferType.IndirectArguments);
                var args= new int[4];
                args[0] = 0; args[1] = 1; args[2] = 0; args[3] = 0;
                cbDrawArgs.SetData (args);
            }
            if (cbPoints == null)
            {
                cbPoints = new ComputeBuffer (90000, 12+16, ComputeBufferType.Append);
            }
        }

        float FocalDistance01 ( float worldDist) {
            return cachedCamera.WorldToViewportPoint((worldDist-cachedCamera.nearClipPlane) * cachedCamera.transform.forward + cachedCamera.transform.position).z / (cachedCamera.farClipPlane-cachedCamera.nearClipPlane);
        }

        private void WriteCoc ( RenderTexture fromTo, bool fgDilate) {
            dofHdrMaterial.SetTexture("_FgOverlap", null);

            if (nearBlur && fgDilate) {

                int rtW = fromTo.width/2;
                int rtH = fromTo.height/2;

                // capture fg coc
                RenderTexture temp2 = RenderTexture.GetTemporary (rtW, rtH, 0, fromTo.format);
                Graphics.Blit (fromTo, temp2, dofHdrMaterial, 4);

                // special blur
                float fgAdjustment = internalBlurWidth * foregroundOverlap;

                dofHdrMaterial.SetVector ("_Offsets", new Vector4 (0.0f, fgAdjustment , 0.0f, fgAdjustment));
                RenderTexture temp1 = RenderTexture.GetTemporary (rtW, rtH, 0, fromTo.format);
                Graphics.Blit (temp2, temp1, dofHdrMaterial, 2);
                RenderTexture.ReleaseTemporary(temp2);

                dofHdrMaterial.SetVector ("_Offsets", new Vector4 (fgAdjustment, 0.0f, 0.0f, fgAdjustment));
                temp2 = RenderTexture.GetTemporary (rtW, rtH, 0, fromTo.format);
                Graphics.Blit (temp1, temp2, dofHdrMaterial, 2);
                RenderTexture.ReleaseTemporary(temp1);

                // "merge up" with background COC
                dofHdrMaterial.SetTexture("_FgOverlap", temp2);
                fromTo.MarkRestoreExpected(); // only touching alpha channel, RT restore expected
                Graphics.Blit (fromTo, fromTo, dofHdrMaterial,  13);
                RenderTexture.ReleaseTemporary(temp2);
            }
            else {
                // capture full coc in alpha channel (fromTo is not read, but bound to detect screen flip)
				fromTo.MarkRestoreExpected(); // only touching alpha channel, RT restore expected
                Graphics.Blit (fromTo, fromTo, dofHdrMaterial,  0);
            }
        }

        void OnRenderImage (RenderTexture source, RenderTexture destination) {
            if (!CheckResources ()) {
                Graphics.Blit (source, destination);
                return;
            }

            // clamp & prepare values so they make sense

            if (aperture < 0.0f) aperture = 0.0f;
            if (maxBlurSize < 0.1f) maxBlurSize = 0.1f;
            focalSize = Mathf.Clamp(focalSize, 0.0f, 2.0f);
            internalBlurWidth = Mathf.Max(maxBlurSize, 0.0f);

            // focal & coc calculations

            focalDistance01 = (focalTransform) ? (cachedCamera.WorldToViewportPoint (focalTransform.position)).z / (cachedCamera.farClipPlane) : FocalDistance01 (focalLength);
            dofHdrMaterial.SetVector("_CurveParams", new Vector4(1.0f, focalSize, (1.0f / (1.0f - aperture) - 1.0f), focalDistance01));

            // possible render texture helpers

            RenderTexture rtLow = null;
            RenderTexture rtLow2 = null;
            RenderTexture rtSuperLow1 = null;
            RenderTexture rtSuperLow2 = null;
            float fgBlurDist = internalBlurWidth * foregroundOverlap;

            if (visualizeFocus)
            {

                //
                // 2.
                // visualize coc
                //
                //

                WriteCoc (source, true);
                Graphics.Blit (source, destination, dofHdrMaterial, 16);
            }
            else if ((blurType == BlurType.DX11) && dx11bokehMaterial)
            {

                //
                // 1.
                // optimized dx11 bokeh scatter
                //
                //


                if (highResolution) {

                    internalBlurWidth = internalBlurWidth < 0.1f ? 0.1f : internalBlurWidth;
                    fgBlurDist = internalBlurWidth * foregroundOverlap;

                    rtLow = RenderTexture.GetTemporary (source.width, source.height, 0, source.format);

                    var dest2= RenderTexture.GetTemporary (source.width, source.height, 0, source.format);

                    // capture COC
                    WriteCoc (source, false);

                    // blur a bit so we can do a frequency check
                    rtSuperLow1 = RenderTexture.GetTemporary(source.width>>1, source.height>>1, 0, source.format);
                    rtSuperLow2 = RenderTexture.GetTemporary(source.width>>1, source.height>>1, 0, source.format);

                    Graphics.Blit(source, rtSuperLow1, dofHdrMaterial, 15);
                    dofHdrMaterial.SetVector ("_Offsets", new Vector4 (0.0f, 1.5f , 0.0f, 1.5f));
                    Graphics.Blit (rtSuperLow1, rtSuperLow2, dofHdrMaterial, 19);
                    dofHdrMaterial.SetVector ("_Offsets", new Vector4 (1.5f, 0.0f, 0.0f, 1.5f));
                    Graphics.Blit (rtSuperLow2, rtSuperLow1, dofHdrMaterial, 19);

                    // capture fg coc
                    if (nearBlur)
                        Graphics.Blit (source, rtSuperLow2, dofHdrMaterial, 4);

                    dx11bokehMaterial.SetTexture ("_BlurredColor", rtSuperLow1);
                    dx11bokehMaterial.SetFloat ("_SpawnHeuristic", dx11SpawnHeuristic);
                    dx11bokehMaterial.SetVector ("_BokehParams", new Vector4(dx11BokehScale, dx11BokehIntensity, Mathf.Clamp(dx11BokehThreshold, 0.005f, 4.0f), internalBlurWidth));
                    dx11bokehMaterial.SetTexture ("_FgCocMask", nearBlur ? rtSuperLow2 : null);

                    // collect bokeh candidates and replace with a darker pixel
                    Graphics.SetRandomWriteTarget (1, cbPoints);
                    Graphics.Blit (source, rtLow, dx11bokehMaterial, 0);
                    Graphics.ClearRandomWriteTargets ();

                    // fg coc blur happens here (after collect!)
                    if (nearBlur) {
                        dofHdrMaterial.SetVector ("_Offsets", new Vector4 (0.0f, fgBlurDist , 0.0f, fgBlurDist));
                        Graphics.Blit (rtSuperLow2, rtSuperLow1, dofHdrMaterial, 2);
                        dofHdrMaterial.SetVector ("_Offsets", new Vector4 (fgBlurDist, 0.0f, 0.0f, fgBlurDist));
                        Graphics.Blit (rtSuperLow1, rtSuperLow2, dofHdrMaterial, 2);

                        // merge fg coc with bg coc
                        Graphics.Blit (rtSuperLow2, rtLow, dofHdrMaterial, 3);
                    }

                    // NEW: LAY OUT ALPHA on destination target so we get nicer outlines for the high rez version
                    Graphics.Blit (rtLow, dest2, dofHdrMaterial, 20);

                    // box blur (easier to merge with bokeh buffer)
                    dofHdrMaterial.SetVector ("_Offsets", new Vector4 (internalBlurWidth, 0.0f , 0.0f, internalBlurWidth));
                    Graphics.Blit (rtLow, source, dofHdrMaterial, 5);
                    dofHdrMaterial.SetVector ("_Offsets", new Vector4 (0.0f, internalBlurWidth, 0.0f, internalBlurWidth));
                    Graphics.Blit (source, dest2, dofHdrMaterial, 21);

                    // apply bokeh candidates
                    Graphics.SetRenderTarget (dest2);
                    ComputeBuffer.CopyCount (cbPoints, cbDrawArgs, 0);
                    dx11bokehMaterial.SetBuffer ("pointBuffer", cbPoints);
                    dx11bokehMaterial.SetTexture ("_MainTex", dx11BokehTexture);
                    dx11bokehMaterial.SetVector ("_Screen", new Vector3(1.0f/(1.0f*source.width), 1.0f/(1.0f*source.height), internalBlurWidth));
                    dx11bokehMaterial.SetPass (2);

                    Graphics.DrawProceduralIndirect (MeshTopology.Points, cbDrawArgs, 0);

                    Graphics.Blit (dest2, destination);	// hackaround for DX11 high resolution flipfun (OPTIMIZEME)

                    RenderTexture.ReleaseTemporary(dest2);
                    RenderTexture.ReleaseTemporary(rtSuperLow1);
                    RenderTexture.ReleaseTemporary(rtSuperLow2);
                }
                else {
                    rtLow = RenderTexture.GetTemporary (source.width>>1, source.height>>1, 0, source.format);
                    rtLow2 = RenderTexture.GetTemporary (source.width>>1, source.height>>1, 0, source.format);

                    fgBlurDist = internalBlurWidth * foregroundOverlap;

                    // capture COC & color in low resolution
                    WriteCoc (source, false);
                    source.filterMode = FilterMode.Bilinear;
                    Graphics.Blit (source, rtLow, dofHdrMaterial, 6);

                    // blur a bit so we can do a frequency check
                    rtSuperLow1 = RenderTexture.GetTemporary(rtLow.width>>1, rtLow.height>>1, 0, rtLow.format);
                    rtSuperLow2 = RenderTexture.GetTemporary(rtLow.width>>1, rtLow.height>>1, 0, rtLow.format);

                    Graphics.Blit(rtLow, rtSuperLow1, dofHdrMaterial, 15);
                    dofHdrMaterial.SetVector ("_Offsets", new Vector4 (0.0f, 1.5f , 0.0f, 1.5f));
                    Graphics.Blit (rtSuperLow1, rtSuperLow2, dofHdrMaterial, 19);
                    dofHdrMaterial.SetVector ("_Offsets", new Vector4 (1.5f, 0.0f, 0.0f, 1.5f));
                    Graphics.Blit (rtSuperLow2, rtSuperLow1, dofHdrMaterial, 19);

                    RenderTexture rtLow3 = null;

                    if (nearBlur) {
                        // capture fg coc
                        rtLow3 = RenderTexture.GetTemporary (source.width>>1, source.height>>1, 0, source.format);
                        Graphics.Blit (source, rtLow3, dofHdrMaterial, 4);
                    }

                    dx11bokehMaterial.SetTexture ("_BlurredColor", rtSuperLow1);
                    dx11bokehMaterial.SetFloat ("_SpawnHeuristic", dx11SpawnHeuristic);
                    dx11bokehMaterial.SetVector ("_BokehParams", new Vector4(dx11BokehScale, dx11BokehIntensity, Mathf.Clamp(dx11BokehThreshold, 0.005f, 4.0f), internalBlurWidth));
                    dx11bokehMaterial.SetTexture ("_FgCocMask", rtLow3);

                    // collect bokeh candidates and replace with a darker pixel
                    Graphics.SetRandomWriteTarget (1, cbPoints);
                    Graphics.Blit (rtLow, rtLow2, dx11bokehMaterial, 0);
                    Graphics.ClearRandomWriteTargets ();

                    RenderTexture.ReleaseTemporary(rtSuperLow1);
                    RenderTexture.ReleaseTemporary(rtSuperLow2);

                    // fg coc blur happens here (after collect!)
                    if (nearBlur) {
                        dofHdrMaterial.SetVector ("_Offsets", new Vector4 (0.0f, fgBlurDist , 0.0f, fgBlurDist));
                        Graphics.Blit (rtLow3, rtLow, dofHdrMaterial, 2);
                        dofHdrMaterial.SetVector ("_Offsets", new Vector4 (fgBlurDist, 0.0f, 0.0f, fgBlurDist));
                        Graphics.Blit (rtLow, rtLow3, dofHdrMaterial, 2);

                        // merge fg coc with bg coc
                        Graphics.Blit (rtLow3, rtLow2, dofHdrMaterial, 3);
                    }

                    // box blur (easier to merge with bokeh buffer)
                    dofHdrMaterial.SetVector ("_Offsets", new Vector4 (internalBlurWidth, 0.0f , 0.0f, internalBlurWidth));
                    Graphics.Blit (rtLow2, rtLow, dofHdrMaterial, 5);
                    dofHdrMaterial.SetVector ("_Offsets", new Vector4 (0.0f, internalBlurWidth, 0.0f, internalBlurWidth));
                    Graphics.Blit (rtLow, rtLow2, dofHdrMaterial, 5);

                    // apply bokeh candidates
                    Graphics.SetRenderTarget (rtLow2);
                    ComputeBuffer.CopyCount (cbPoints, cbDrawArgs, 0);
                    dx11bokehMaterial.SetBuffer ("pointBuffer", cbPoints);
                    dx11bokehMaterial.SetTexture ("_MainTex", dx11BokehTexture);
                    dx11bokehMaterial.SetVector ("_Screen", new Vector3(1.0f/(1.0f*rtLow2.width), 1.0f/(1.0f*rtLow2.height), internalBlurWidth));
                    dx11bokehMaterial.SetPass (1);
                    Graphics.DrawProceduralIndirect (MeshTopology.Points, cbDrawArgs, 0);

                    // upsample & combine
                    dofHdrMaterial.SetTexture ("_LowRez", rtLow2);
                    dofHdrMaterial.SetTexture ("_FgOverlap", rtLow3);
                    dofHdrMaterial.SetVector ("_Offsets",  ((1.0f*source.width)/(1.0f*rtLow2.width)) * internalBlurWidth * Vector4.one);
                    Graphics.Blit (source, destination, dofHdrMaterial, 9);

                    if (rtLow3) RenderTexture.ReleaseTemporary(rtLow3);
                }
            }
            else
            {

                //
                // 2.
                // poisson disc style blur in low resolution
                //
                //

                source.filterMode = FilterMode.Bilinear;

                if (highResolution) internalBlurWidth *= 2.0f;

                WriteCoc (source, true);

                rtLow = RenderTexture.GetTemporary (source.width >> 1, source.height >> 1, 0, source.format);
                rtLow2 = RenderTexture.GetTemporary (source.width >> 1, source.height >> 1, 0, source.format);

                int blurPass = (blurSampleCount == BlurSampleCount.High || blurSampleCount == BlurSampleCount.Medium) ? 17 : 11;

                if (highResolution) {
                    dofHdrMaterial.SetVector ("_Offsets", new Vector4 (0.0f, internalBlurWidth, 0.025f, internalBlurWidth));
                    Graphics.Blit (source, destination, dofHdrMaterial, blurPass);
                }
                else {
                    dofHdrMaterial.SetVector ("_Offsets", new Vector4 (0.0f, internalBlurWidth, 0.1f, internalBlurWidth));

                    // blur
                    Graphics.Blit (source, rtLow, dofHdrMaterial, 6);
                    Graphics.Blit (rtLow, rtLow2, dofHdrMaterial, blurPass);

                    // cheaper blur in high resolution, upsample and combine
                    dofHdrMaterial.SetTexture("_LowRez", rtLow2);
                    dofHdrMaterial.SetTexture("_FgOverlap", null);
                    dofHdrMaterial.SetVector ("_Offsets",  Vector4.one * ((1.0f*source.width)/(1.0f*rtLow2.width)) * internalBlurWidth);
                    Graphics.Blit (source, destination, dofHdrMaterial, blurSampleCount == BlurSampleCount.High ? 18 : 12);
                }
            }

            if (rtLow) RenderTexture.ReleaseTemporary(rtLow);
            if (rtLow2) RenderTexture.ReleaseTemporary(rtLow2);
        }
    }
}