using UnityEngine;
using System.Collections;

// TBD: if some shader is not supported, turn off!

[AddComponentMenu("Image Effects/Contrast Stretch")]
public class ContrastStretchEffect : MonoBehaviour
{
    /// Adaptation speed - percents per frame, if playing at 30FPS.
    /// Default is 0.02 (2% each 1/30s).
    public float adaptationSpeed = 0.02f;

    /// If our scene is really dark (or really bright), we might not want to
    /// stretch its contrast to the full range.
    /// limitMinimum=0, limitMaximum=1 is the same as not applying the effect at all.
    /// limitMinimum=1, limitMaximum=0 is always stretching colors to full range.

    /// The limit on the minimum luminance (0...1) - we won't go above this.
    public float limitMinimum = 0.2f;

    /// The limit on the maximum luminance (0...1) - we won't go below this.
    public float limitMaximum = 0.6f;


    // To maintain adaptation levels over time, we need two 1x1 render textures
    // and ping-pong between them.
    private RenderTexture[] adaptRenderTex = new RenderTexture[2];
    private int curAdaptIndex = 0;


    // Computes scene luminance (grayscale) image
    public Shader shaderLum;
    private Material m_materialLum;
    protected Material materialLum
    {
        get
        {
            if (m_materialLum == null)
            {
                m_materialLum = new Material(shaderLum);
                m_materialLum.hideFlags = HideFlags.HideAndDontSave;
            }
            return m_materialLum;
        }
    }

    // Reduces size of the image by 2x2, while computing maximum/minimum values.
    // By repeatedly applying this shader, we reduce the initial luminance image
    // to 1x1 image with minimum/maximum luminances found.
    public Shader shaderReduce;
    private Material m_materialReduce;
    protected Material materialReduce
    {
        get
        {
            if (m_materialReduce == null)
            {
                m_materialReduce = new Material(shaderReduce);
                m_materialReduce.hideFlags = HideFlags.HideAndDontSave;
            }
            return m_materialReduce;
        }
    }

    // Adaptation shader - gradually "adapts" minimum/maximum luminances,
    // based on currently adapted 1x1 image and the actual 1x1 image of the current scene.
    public Shader shaderAdapt;
    private Material m_materialAdapt;
    protected Material materialAdapt
    {
        get
        {
            if (m_materialAdapt == null)
            {
                m_materialAdapt = new Material(shaderAdapt);
                m_materialAdapt.hideFlags = HideFlags.HideAndDontSave;
            }
            return m_materialAdapt;
        }
    }

    // Final pass - stretches the color values of the original scene, based on currently
    // adpated minimum/maximum values.
    public Shader shaderApply;
    private Material m_materialApply;
    protected Material materialApply
    {
        get
        {
            if (m_materialApply == null)
            {
                m_materialApply = new Material(shaderApply);
                m_materialApply.hideFlags = HideFlags.HideAndDontSave;
            }
            return m_materialApply;
        }
    }

    void Start()
    {
        // Disable if we don't support image effects
        if (!SystemInfo.supportsImageEffects)
        {
            enabled = false;
            return;
        }
    }

    void OnEnable()
    {
        for (int i = 0; i < 2; ++i)
        {
            if (!adaptRenderTex[i])
            {
                adaptRenderTex[i] = new RenderTexture(1, 1, 32);
                adaptRenderTex[i].hideFlags = HideFlags.HideAndDontSave;
            }
        }
    }

    void OnDisable()
    {
        for (int i = 0; i < 2; ++i)
        {
            DestroyImmediate(adaptRenderTex[i]);
            adaptRenderTex[i] = null;
        }
        if (m_materialLum)
            DestroyImmediate(m_materialLum);
        if (m_materialReduce)
            DestroyImmediate(m_materialReduce);
        if (m_materialAdapt)
            DestroyImmediate(m_materialAdapt);
        if (m_materialApply)
            DestroyImmediate(m_materialApply);
    }


    /// Apply the filter	
    void OnRenderImage(RenderTexture source, RenderTexture destination)
    {
        // Blit to smaller RT and convert to luminance on the way
        const int TEMP_RATIO = 1; // 4x4 smaller
        RenderTexture rtTempSrc = RenderTexture.GetTemporary(source.width / TEMP_RATIO, source.height / TEMP_RATIO);
        ImageEffects.BlitWithMaterial(materialLum, source, rtTempSrc);

        // Repeatedly reduce this image in size, computing min/max luminance values
        // In the end we'll have 1x1 image with min/max luminances found.
        const int FINAL_SIZE = 1;
        //const int FINAL_SIZE = 1;
        while (rtTempSrc.width > FINAL_SIZE || rtTempSrc.height > FINAL_SIZE)
        {
            const int REDUCE_RATIO = 2; // our shader does 2x2 reduction
            int destW = rtTempSrc.width / REDUCE_RATIO;
            if (destW < FINAL_SIZE) destW = FINAL_SIZE;
            int destH = rtTempSrc.height / REDUCE_RATIO;
            if (destH < FINAL_SIZE) destH = FINAL_SIZE;
            RenderTexture rtTempDst = RenderTexture.GetTemporary(destW, destH);
            ImageEffects.BlitWithMaterial(materialReduce, rtTempSrc, rtTempDst);

            // Release old src temporary, and make new temporary the source
            RenderTexture.ReleaseTemporary(rtTempSrc);
            rtTempSrc = rtTempDst;
        }

        // Update viewer's adaptation level
        CalculateAdaptation(rtTempSrc);

        // Apply contrast strech to the original scene, using currently adapted parameters
        materialApply.SetTexture("_AdaptTex", adaptRenderTex[curAdaptIndex]);
        ImageEffects.BlitWithMaterial(materialApply, source, destination);

        RenderTexture.ReleaseTemporary(rtTempSrc);
    }


    /// Helper function to do gradual adaptation to min/max luminances
    private void CalculateAdaptation(Texture curTexture)
    {
        int prevAdaptIndex = curAdaptIndex;
        curAdaptIndex = (curAdaptIndex + 1) % 2;

        // Adaptation speed is expressed in percents/frame, based on 30FPS.
        // Calculate the adaptation lerp, based on current FPS.
        float adaptLerp = 1.0f - Mathf.Pow(1.0f - adaptationSpeed, 30.0f * Time.deltaTime);
        const float kMinAdaptLerp = 0.01f;
        adaptLerp = Mathf.Clamp(adaptLerp, kMinAdaptLerp, 1);

        materialAdapt.SetTexture("_CurTex", curTexture);
        materialAdapt.SetVector("_AdaptParams", new Vector4(
            adaptLerp,
            limitMinimum,
            limitMaximum,
            0.0f
        ));
        ImageEffects.BlitWithMaterial(materialAdapt,
            adaptRenderTex[prevAdaptIndex],
            adaptRenderTex[curAdaptIndex]);
    }
}