//#define DEBUG_FFT

using System;
using System.Collections.Generic;

using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;
using Random = UnityEngine.Random;


namespace KWS
{
    internal class OceanFftWavesPass : WaterPass
    {
        internal override string PassName => "Water.FftWavesPass";

        static int kernelSpectrumInit;
        static int kernelSpectrumUpdate;
        static int kernelNormal;

        static Dictionary<int, Texture2D> _butterflyTextures         = new Dictionary<int, Texture2D>();
        private static readonly int                        KwsDisplaceXYZ             = Shader.PropertyToID("_displaceXYZ");
        private static readonly int                        KwsKwsNormalFoamTargetRW   = Shader.PropertyToID("KWS_NormalFoamTargetRW");
        private static readonly int                        KwsKwsPrevNormalFoamTarget = Shader.PropertyToID("KWS_PrevNormalFoamTarget");

        private WindZone _lastWindZone;
        private float _lastWindZoneSpeed;
        private float _lastWindZoneTurbulence;
        private Vector3 _lastWindZoneRotation;
        private CommandBuffer _cmd;

        private const float          DefaultTimeScale = 1.5f;
        private const GraphicsFormat fftFormat        = GraphicsFormat.R16G16B16A16_SFloat;
        private const GraphicsFormat spectrumFormat   = GraphicsFormat.R16G16_SFloat;
        private const GraphicsFormat normalFormat     = GraphicsFormat.R16G16B16A16_SFloat;


        static RTHandle[] DisplaceTexture = new RTHandle[2];
        static RTHandle[] NormalTextures  = new RTHandle[2];

        static RTHandle spectrumInit;
        static RTHandle spectrumDisplaceX;
        static RTHandle spectrumDisplaceY;
        static RTHandle spectrumDisplaceZ;

        static RTHandle fftTemp1;
        static RTHandle fftTemp2;
        static RTHandle fftTemp3;

        static ComputeShader spectrumShader;
        static ComputeShader shaderFFT;

        static bool RequireReinitializeSpectrum;

        static int Frame;

        static bool RequireReinitialize(int size, int cascades)
        {
            if (DisplaceTexture[0] == null || DisplaceTexture[0].rt == null || shaderFFT == null) return true;

            var rt = DisplaceTexture[0].rt;
            if (rt.width != size || rt.volumeDepth != cascades) return true;

            return false;
        }

        static void Initialize(int size, int cascades)
        {
       
            spectrumInit      = KWS_CoreUtils.RTHandles.Alloc(size, size, colorFormat: fftFormat,      enableRandomWrite: true, dimension: TextureDimension.Tex2DArray, slices: cascades);
            spectrumDisplaceY = KWS_CoreUtils.RTHandles.Alloc(size, size, colorFormat: spectrumFormat, enableRandomWrite: true, dimension: TextureDimension.Tex2DArray, slices: cascades);
            spectrumDisplaceX = KWS_CoreUtils.RTHandles.Alloc(size, size, colorFormat: spectrumFormat, enableRandomWrite: true, dimension: TextureDimension.Tex2DArray, slices: cascades);
            spectrumDisplaceZ = KWS_CoreUtils.RTHandles.Alloc(size, size, colorFormat: spectrumFormat, enableRandomWrite: true, dimension: TextureDimension.Tex2DArray, slices: cascades);

            fftTemp1 = KWS_CoreUtils.RTHandles.Alloc(size, size, colorFormat: spectrumFormat, enableRandomWrite: true, dimension: TextureDimension.Tex2DArray, slices: cascades);
            fftTemp2 = KWS_CoreUtils.RTHandles.Alloc(size, size, colorFormat: spectrumFormat, enableRandomWrite: true, dimension: TextureDimension.Tex2DArray, slices: cascades);
            fftTemp3 = KWS_CoreUtils.RTHandles.Alloc(size, size, colorFormat: spectrumFormat, enableRandomWrite: true, dimension: TextureDimension.Tex2DArray, slices: cascades);

            DisplaceTexture[0] = KWS_CoreUtils.RTHandles.Alloc(size, size, name: "KWS_FftWavesDisplacement0", colorFormat: fftFormat, enableRandomWrite: true, dimension: TextureDimension.Tex2DArray, slices: cascades);
            DisplaceTexture[1] = KWS_CoreUtils.RTHandles.Alloc(size, size, name: "KWS_FftWavesDisplacement1", colorFormat: fftFormat, enableRandomWrite: true, dimension: TextureDimension.Tex2DArray, slices: cascades);

            NormalTextures[0] = KWS_CoreUtils.RTHandles.Alloc(size, size, name: "KWS_FftWavesNormal1", colorFormat: normalFormat, enableRandomWrite: true,
                                                           autoGenerateMips: false, useMipMap: true, dimension: TextureDimension.Tex2DArray, slices: cascades, filterMode: FilterMode.Trilinear);
            NormalTextures[1] = KWS_CoreUtils.RTHandles.Alloc(size, size, name: "KWS_FftWavesNormal2", colorFormat: normalFormat, enableRandomWrite: true,
                                                           autoGenerateMips: false, useMipMap: true, dimension: TextureDimension.Tex2DArray, slices: cascades, filterMode: FilterMode.Trilinear);


            GetOrCreateButterflyTexture(size);
            //this.WaterLog(DisplaceTexture[0], NormalTextures[0]);
        }

        static void InitializeShaders(int size, int cascades)
        {

            if (spectrumShader == null) spectrumShader = KWS_CoreUtils.LoadComputeShader("Common/CommandPass/KWS_WavesSpectrum");
            if (shaderFFT == null) shaderFFT = KWS_CoreUtils.LoadComputeShader("Common/CommandPass/KWS_WavesFFT");


            if (spectrumShader != null)
            {
                spectrumShader.name = "WavesSpectrum";
                kernelSpectrumInit = spectrumShader.FindKernel("SpectrumInitalize");
                kernelSpectrumUpdate = spectrumShader.FindKernel("SpectrumUpdate");

                spectrumShader.SetTexture(kernelSpectrumUpdate, "SpectrumInit", spectrumInit);
                spectrumShader.SetTexture(kernelSpectrumUpdate, "SpectrumDisplaceX", spectrumDisplaceX);
                spectrumShader.SetTexture(kernelSpectrumUpdate, "SpectrumDisplaceY", spectrumDisplaceY);
                spectrumShader.SetTexture(kernelSpectrumUpdate, "SpectrumDisplaceZ", spectrumDisplaceZ);
            }

            if (shaderFFT != null)
            {
                shaderFFT.name = "WavesFFT";
                kernelNormal = shaderFFT.FindKernel("ComputeNormal");

                var fftKernel = GetKernelBySize(size);

                shaderFFT.SetTexture(fftKernel, "SpectrumDisplaceX", spectrumDisplaceX);
                shaderFFT.SetTexture(fftKernel, "SpectrumDisplaceY", spectrumDisplaceY);
                shaderFFT.SetTexture(fftKernel, "SpectrumDisplaceZ", spectrumDisplaceZ);
                shaderFFT.SetTexture(fftKernel, "inputButterfly", GetOrCreateButterflyTexture(size));
                shaderFFT.SetTexture(fftKernel, "_displaceX", fftTemp1);
                shaderFFT.SetTexture(fftKernel, "_displaceY", fftTemp2);
                shaderFFT.SetTexture(fftKernel, "_displaceZ", fftTemp3);

                shaderFFT.SetTexture(fftKernel + 1, "SpectrumDisplaceX", fftTemp1);
                shaderFFT.SetTexture(fftKernel + 1, "SpectrumDisplaceY", fftTemp2);
                shaderFFT.SetTexture(fftKernel + 1, "SpectrumDisplaceZ", fftTemp3);
                shaderFFT.SetTexture(fftKernel + 1, "inputButterfly", GetOrCreateButterflyTexture(size));
                //shaderFFT.SetTexture(fftKernel + 1, "_displaceXYZ",      DisplaceTexture);

                shaderFFT.SetVector("KWS_FFT_TexelSize", new Vector4(1f / NormalTextures[0].rt.width, 1f / NormalTextures[0].rt.height, NormalTextures[0].rt.width, NormalTextures[0].rt.height));
                //shaderFFT.SetTexture(kernelNormal, "_displaceXYZ", DisplaceTexture);
            }
        }

        static RTHandle GetTargetNormal()
        {
            return NormalTextures[Frame];
        }

        static RTHandle GetPreviousTargetNormal()
        {
            return NormalTextures[(Frame + 1) % 2];
        }

        static RTHandle GetDisplacement()
        {
            return DisplaceTexture[Frame];
        }

        static RTHandle GetPreviousDisplacement()
        {
            return DisplaceTexture[(Frame + 1) % 2];
        }

        static void SwapTargetNormal()
        {
            Frame = (Frame + 1) % 2;
        }


        static void ReleaseTextures()
        {
            spectrumInit?.Release();
            spectrumDisplaceY?.Release();
            spectrumDisplaceX?.Release();
            spectrumDisplaceZ?.Release();

            fftTemp1?.Release();
            fftTemp2?.Release();
            fftTemp3?.Release();

            DisplaceTexture[0]?.Release();
            DisplaceTexture[1]?.Release();

            NormalTextures[0]?.Release();
            NormalTextures[1]?.Release();

            DisplaceTexture[0] = DisplaceTexture[1] = NormalTextures[0] = NormalTextures[1] = null;
            
            spectrumInit = spectrumDisplaceX = spectrumDisplaceY = spectrumDisplaceZ = spectrumDisplaceZ = null;
            
            fftTemp1 = fftTemp2 = fftTemp3 = null;

           // this.WaterLog(String.Empty, KW_Extensions.WaterLogMessageType.ReleaseRT);
        }


        public OceanFftWavesPass()
        {
            WaterSystem.OnAnyWaterSettingsChanged += OnAnyWaterSettingsChanged;
        }


        public override void Release()
        {
            WaterSystem.OnAnyWaterSettingsChanged -= OnAnyWaterSettingsChanged;

            ReleaseFFT();
            
            this.WaterLog(String.Empty, KW_Extensions.WaterLogMessageType.ReleaseRT);
            this.WaterLog(String.Empty, KW_Extensions.WaterLogMessageType.Release);
        }

        internal static void ReleaseFFT()
        {
            foreach (var butterflyTexture in _butterflyTextures) KW_Extensions.SafeDestroy(butterflyTexture.Value);
            _butterflyTextures.Clear();

            ReleaseTextures();
            KW_Extensions.SafeDestroy(spectrumShader, shaderFFT);
            spectrumShader              = null;
            shaderFFT                   = null;
            RequireReinitializeSpectrum = true;

        }

        private void OnAnyWaterSettingsChanged(WaterSystem.WaterSettingsCategory changedTabs)
        {   
            if (KWS_Ocean.Instance == false) return;
            
            if (!changedTabs.HasTab(WaterSystem.WaterSettingsCategory.Ocean)) return;
            
            var size     = (int)KWS_Ocean.Instance.FftWavesQuality;
            var cascades = KWS_Ocean.Instance.FftWavesCascades;

            InitializeFftWavesData(size, cascades);
        }

        static void InitializeFftWavesData(int size, int cascades)
        {  
            if (RequireReinitialize(size, cascades))
            {
                ReleaseTextures();
                Initialize(size, cascades);
                InitializeShaders(size, cascades);
            }

            RequireReinitializeSpectrum = true;
        }

        public override void ExecutePerFrame(HashSet<Camera> cameras, CustomFixedUpdates fixedUpdates)
        {
            #if DEBUG_FFT
                return;
            #endif
            
            if (KWS_Ocean.Instance == false)
            {
                Shader.SetGlobalTexture(KWS_ShaderConstants.FFT.KWS_FftWavesDisplace, KWS_CoreUtils.DefaultBlackTexture);
                Shader.SetGlobalTexture(KWS_ShaderConstants.FFT.KWS_FftWavesNormal,   KWS_CoreUtils.DefaultBlackTexture);
                
                return;
            }
            
            if (WaterSystem.UseNetworkBuoyancy == false && fixedUpdates.TimeScaledFramesCount_60fps == 0) return;

            if (_cmd == null) _cmd = new CommandBuffer() { name = PassName };
            _cmd.Clear();
            
            if (KWS_Ocean.Instance.WindZone != null && IsWindZoneChanged()) RequireReinitializeSpectrum = true;

            ExecuteInstance(_cmd);
            WaterSharedResources.FftWavesDisplacement = GetDisplacement();
            WaterSharedResources.FftWavesNormal       = GetTargetNormal();

            Shader.SetGlobalTexture(KWS_ShaderConstants.FFT.KWS_FftWavesDisplace, WaterSharedResources.FftWavesDisplacement);
            Shader.SetGlobalTexture(KWS_ShaderConstants.FFT.KWS_FftWavesNormal,   WaterSharedResources.FftWavesNormal);
        
            Graphics.ExecuteCommandBuffer(_cmd);
           
        }

        public override void ExecuteBeforeCameraRendering(Camera cam, ScriptableRenderContext context)
        {
            #if DEBUG_FFT
                if (KWS_Ocean.Instance == false)
                {
                    Shader.SetGlobalTexture(KWS_ShaderConstants.FFT.KWS_FftWavesDisplace, KWS_CoreUtils.DefaultBlackTexture);
                    Shader.SetGlobalTexture(KWS_ShaderConstants.FFT.KWS_FftWavesNormal,   KWS_CoreUtils.DefaultBlackTexture);
                    
                    return;
                }
              
                if (_cmd == null) _cmd = new CommandBuffer() { name = PassName };
                _cmd.Clear();
                
                RequireReinitializeSpectrum = true;

                ExecuteInstance(_cmd);
                WaterSharedResources.FftWavesDisplacement = GetDisplacement();
                WaterSharedResources.FftWavesNormal       = GetTargetNormal();

                Shader.SetGlobalTexture(KWS_ShaderConstants.FFT.KWS_FftWavesDisplace, WaterSharedResources.FftWavesDisplacement);
                Shader.SetGlobalTexture(KWS_ShaderConstants.FFT.KWS_FftWavesNormal,   WaterSharedResources.FftWavesNormal);
            
                Graphics.ExecuteCommandBuffer(_cmd);
            #endif

        }

 

        void ExecuteInstance(CommandBuffer cmd)
        {
            var size     = (int)KWS_Ocean.Instance.FftWavesQuality;
            var cascades = KWS_Ocean.Instance.FftWavesCascades;
            
            if (RequireReinitialize(size, cascades))
            {  
                InitializeFftWavesData(size, cascades);
                return; //todo one frame delay to avoid nan init. Why? 
            }

            cmd.SetGlobalFloat(KWS_ShaderConstants.ConstantWaterParams.KWS_WavesAreaScale, KWS_Ocean.Instance.WavesAreaScale);
            var time = KW_Extensions.TotalTime() * WaterSystem.Instance.SimulationFPSMultiplier * DefaultTimeScale;
            
            if (RequireReinitializeSpectrum) InitializeSpectrum(cmd, size, cascades);
            UpdateSpectrum(cmd, size, cascades, time);
            DispatchFFT(cmd, size, cascades);
        }


        static void InitializeSpectrum(CommandBuffer cmd, int size, int cascades)
        {
            cmd.SetComputeFloatParam(spectrumShader, "KWS_WindSpeed",    KWS_Ocean.Instance.WindSpeed);
            cmd.SetComputeFloatParam(spectrumShader, "KWS_Turbulence",   KWS_Ocean.Instance.WindTurbulence);
            cmd.SetComputeFloatParam(spectrumShader, "KWS_WindRotation", KWS_Ocean.Instance.WindRotation);

            cmd.SetComputeIntParam(spectrumShader, "KWS_Size", size);

            //cmd.SetComputeFloatParams(spectrumShader, KWS_ShaderConstants.ConstantWaterParams.KWS_WavesDomainSizes, KWS_Settings.FFT.FftDomainSize);
            cmd.SetComputeFloatParam(spectrumShader, KWS_ShaderConstants.ConstantWaterParams.KWS_WavesAreaScale, KWS_Ocean.Instance.WavesAreaScale);

            cmd.SetComputeTextureParam(spectrumShader, kernelSpectrumInit, "RW_SpectrumInit", spectrumInit);

            cmd.DispatchCompute(spectrumShader, kernelSpectrumInit, size / 8, size / 8, cascades);
            RequireReinitializeSpectrum = false;

            //this.WaterLog($"InitializeSpectrum");
        }

        static void UpdateSpectrum(CommandBuffer cmd, int size, int cascades, float time)
        {
           
            if (spectrumShader == null)
            {
                Debug.LogError($"Water UpdateSpectrum error: {spectrumShader}");
                return;
            }
            cmd.SetComputeFloatParam(spectrumShader, "time", time);
            cmd.DispatchCompute(spectrumShader, kernelSpectrumUpdate, size / 8, size / 8, cascades);
        }

        public static RenderTexture BakeFFT(int frames, float loopTime, int size, int cascades, int cascadeIndexToBake)
        {   
            var cmd = new CommandBuffer();

            ReleaseFFT();
            
            InitializeFftWavesData(size, cascades);
            InitializeSpectrum(cmd, size, cascades);
            
            var rtArray = new RenderTexture(size, size, 0) { dimension = TextureDimension.Tex2DArray, volumeDepth = frames, graphicsFormat = fftFormat, useMipMap = false };
            rtArray.Create();
            
            var renderTexDesc = new RenderTextureDescriptor(size, size, fftFormat, 0);
            renderTexDesc.volumeDepth = frames;
            renderTexDesc.useMipMap   = false;
            renderTexDesc.dimension   = TextureDimension.Tex2DArray;
            
            var tempPackedRT = KWS_CoreUtils.RTHandles.Alloc(size, size, name: "KWS_FftWavesDisplacementPacked", colorFormat: fftFormat, enableRandomWrite: true, dimension: TextureDimension.Tex2D);
            var kernelBake = shaderFFT.FindKernel("InterpolationPacking");
            
            UpdateSpectrum(cmd, size, cascades, 0); //not sure why first frame doesnt work correctly 
            DispatchFFT(cmd, size, cascades);
            
            for (int i = 0; i < frames; i++)
            {
                cmd.SetComputeFloatParam(spectrumShader, "KWS_BakeLoopTime", loopTime);
                cmd.SetKeyword("FFT_BAKE_MODE", true);
                cmd.SetComputeFloatParam(shaderFFT, "KWS_BakedPackedRange", 3);

                var currentTime = loopTime * (i / (float)frames);
                var nextTime = loopTime * (((i + 1f) % frames) / (float)frames);
               
                UpdateSpectrum(cmd, size, cascades, currentTime);
                DispatchFFT(cmd, size, cascades);
             
                var currentFftFrame  = GetDisplacement();
                var previousFftFrame = GetPreviousDisplacement();
                
                cmd.SetComputeTextureParam(shaderFFT, kernelBake, "KWS_CurrentFFTFrame", currentFftFrame);
                cmd.SetComputeTextureParam(shaderFFT, kernelBake, "KWS_PrevFFTFrame",    previousFftFrame);
                cmd.SetComputeTextureParam(shaderFFT, kernelBake, "KWS_PackedResult",    tempPackedRT);
                cmd.DispatchCompute(shaderFFT, kernelBake, size / 8, size / 8, 1);
                
                
                cmd.Blit(tempPackedRT, rtArray, cascadeIndexToBake, i);
                
                cmd.SetKeyword("FFT_BAKE_MODE", false);
               
            }
            
            tempPackedRT?.Release();
            
            Graphics.ExecuteCommandBuffer(cmd);
            
            return rtArray;
        }

        static void DispatchFFT(CommandBuffer cmd, int size, int cascades)
        {
           // var instance  = WaterSystem.Instance;
            
            var fftKernel = GetKernelBySize(size);

            if (shaderFFT == null)
            {
                Debug.LogError($"Water DispatchFFT error: {shaderFFT}");
                return;
            }

            cmd.SetComputeTextureParam(shaderFFT, fftKernel + 1, KwsDisplaceXYZ, GetDisplacement());
            cmd.DispatchCompute(shaderFFT, fftKernel,     1,    size, cascades);
            cmd.DispatchCompute(shaderFFT, fftKernel + 1, size, 1,    cascades);

            cmd.SetComputeTextureParam(shaderFFT, kernelNormal, KwsDisplaceXYZ, GetDisplacement());
            cmd.SetComputeTextureParam(shaderFFT, kernelNormal, KwsKwsNormalFoamTargetRW, GetTargetNormal());
            cmd.SetComputeTextureParam(shaderFFT, kernelNormal, KwsKwsPrevNormalFoamTarget, GetPreviousTargetNormal());

            cmd.DispatchCompute(shaderFFT, kernelNormal, size / 8, size / 8, cascades);
            cmd.GenerateMips(GetTargetNormal());

            SwapTargetNormal();
        }


        static Texture2D GetOrCreateButterflyTexture(int size)
        {
            if (!_butterflyTextures.ContainsKey(size)) _butterflyTextures.Add(size, InitializeButterfly(size));

            return _butterflyTextures[size];
        }

        static Texture2D InitializeButterfly(int size)
        {
            var log2Size = Mathf.RoundToInt(Mathf.Log(size, 2));
            var butterflyColors = new Color[size * log2Size];

            int offset = 1, numIterations = size >> 1;
            for (int rowIndex = 0; rowIndex < log2Size; rowIndex++)
            {
                int rowOffset = rowIndex * size;
                {
                    int start = 0, end = 2 * offset;
                    for (int iteration = 0; iteration < numIterations; iteration++)
                    {
                        var bigK = 0.0f;
                        for (int K = start; K < end; K += 2)
                        {
                            var phase = 2.0f * Mathf.PI * bigK * numIterations / size;
                            var cos = Mathf.Cos(phase);
                            var sin = Mathf.Sin(phase);
                            butterflyColors[rowOffset + K / 2] = new Color(cos, -sin, 0, 1);
                            butterflyColors[rowOffset + K / 2 + offset] = new Color(-cos, sin, 0, 1);

                            bigK += 1.0f;
                        }
                        start += 4 * offset;
                        end = start + 2 * offset;
                    }
                }
                numIterations >>= 1;
                offset <<= 1;
            }
            var texButterfly = new Texture2D(size, log2Size, GraphicsFormat.R32G32B32A32_SFloat, TextureCreationFlags.None);
            texButterfly.SetPixels(butterflyColors);
            texButterfly.Apply();
            return texButterfly;
        }


        bool IsWindZoneChanged()
        {
            var windZone = KWS_Ocean.Instance.WindZone;
            if (KWS_Ocean.Instance.WindZone != _lastWindZone)
            {
                _lastWindZone = KWS_Ocean.Instance.WindZone;
                return true;
            }

            if (Math.Abs(_lastWindZoneSpeed - windZone.windMain * KWS_Ocean.Instance.WindZoneSpeedMultiplier) > 0.001f)
            {
                _lastWindZoneSpeed = windZone.windMain * KWS_Ocean.Instance.WindZoneSpeedMultiplier;
                return true;
            }

            if (Math.Abs(_lastWindZoneTurbulence - windZone.windTurbulence * KWS_Ocean.Instance.WindZoneTurbulenceMultiplier) > 0.001f)
            {
                _lastWindZoneTurbulence = windZone.windTurbulence * KWS_Ocean.Instance.WindZoneTurbulenceMultiplier;
                return true;
            }

            var forward = windZone.transform.forward;
            if (Math.Abs(_lastWindZoneRotation.x - forward.x) > 0.001f || Math.Abs(_lastWindZoneRotation.z - forward.z) > 0.001f)
            {
                _lastWindZoneRotation = forward;
                return true;
            }

            return false;
        }



        static int GetKernelBySize(int size)
        {
            var kernelOffset = 0;
            kernelOffset = size switch
            {
                (int)WaterQualityLevelSettings.FftWavesQualityEnum.Low    => 0,
                (int)WaterQualityLevelSettings.FftWavesQualityEnum.Medium => 2,
                (int)WaterQualityLevelSettings.FftWavesQualityEnum.High   => 4,
                (int)WaterQualityLevelSettings.FftWavesQualityEnum.Ultra  => 6,
                //(int)WaterQualityLevelSettings.FftWavesQualityEnum.Extreme  => 8,
                _                                                         => kernelOffset
            };
            return kernelOffset;
        }

    }
}