using UnityEngine;
using UnityEngine.Experimental.Rendering;

namespace FidelityFX.OpticalFlow
{
    internal class OpticalFlowResources
    {
        internal const int OpticalFlowMaxPyramidLevels = 7;
        internal const int HistogramBins = 256;
        internal const int HistogramsPerDim = 3;
        internal const int HistogramShifts = 3;

        public readonly RenderTexture[] OpticalFlowInput = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowInputLevel1 = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowInputLevel2 = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowInputLevel3 = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowInputLevel4 = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowInputLevel5 = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowInputLevel6 = new RenderTexture[2];
        
        public readonly RenderTexture[] OpticalFlow = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowLevel1 = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowLevel2 = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowLevel3 = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowLevel4 = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowLevel5 = new RenderTexture[2];
        public readonly RenderTexture[] OpticalFlowLevel6 = new RenderTexture[2];

        public RenderTexture OpticalFlowSCDHistogram;
        public RenderTexture OpticalFlowSCDPreviousHistogram;
        public RenderTexture OpticalFlowSCDTemp;
        
        public void Create(OpticalFlow.ContextDescription contextDescription)
        {
            Vector2Int opticalFlowInputTextureSize = contextDescription.resolution;
            
            const int minBlockSize = 8;
            Vector2Int opticalFlowTextureSize = GetOpticalFlowTextureSize(contextDescription.resolution, minBlockSize);
            Vector2Int opticalFlowLevel1TextureSize = NextLevelSize(opticalFlowTextureSize);
            Vector2Int opticalFlowLevel2TextureSize = NextLevelSize(opticalFlowLevel1TextureSize);
            Vector2Int opticalFlowLevel3TextureSize = NextLevelSize(opticalFlowLevel2TextureSize);
            Vector2Int opticalFlowLevel4TextureSize = NextLevelSize(opticalFlowLevel3TextureSize);
            Vector2Int opticalFlowLevel5TextureSize = NextLevelSize(opticalFlowLevel4TextureSize);
            Vector2Int opticalFlowLevel6TextureSize = NextLevelSize(opticalFlowLevel5TextureSize);

            Vector2Int opticalFlowHistogramTextureSize = GetOpticalFlowHistogramSize(0);
            Vector2Int globalMotionSearchMaxDispatchSize = GetGlobalMotionSearchDispatchSize(0);
            int globalMotionSearchTextureWidth = 4 + (globalMotionSearchMaxDispatchSize.x * globalMotionSearchMaxDispatchSize.y);
            
            CreateDoubleBufferedResource(OpticalFlowInput, "OPTICALFLOW_OpticalFlowInput", opticalFlowInputTextureSize, GraphicsFormat.R8_UInt);
            opticalFlowInputTextureSize = new Vector2Int(opticalFlowInputTextureSize.x >> 1, opticalFlowInputTextureSize.y >> 1);
            CreateDoubleBufferedResource(OpticalFlowInputLevel1, "OPTICALFLOW_OpticalFlowInputLevel1-", opticalFlowInputTextureSize, GraphicsFormat.R8_UInt);
            opticalFlowInputTextureSize = new Vector2Int(opticalFlowInputTextureSize.x >> 1, opticalFlowInputTextureSize.y >> 1);
            CreateDoubleBufferedResource(OpticalFlowInputLevel2, "OPTICALFLOW_OpticalFlowInputLevel2-", opticalFlowInputTextureSize, GraphicsFormat.R8_UInt);
            opticalFlowInputTextureSize = new Vector2Int(opticalFlowInputTextureSize.x >> 1, opticalFlowInputTextureSize.y >> 1);
            CreateDoubleBufferedResource(OpticalFlowInputLevel3, "OPTICALFLOW_OpticalFlowInputLevel3-", opticalFlowInputTextureSize, GraphicsFormat.R8_UInt);
            opticalFlowInputTextureSize = new Vector2Int(opticalFlowInputTextureSize.x >> 1, opticalFlowInputTextureSize.y >> 1);
            CreateDoubleBufferedResource(OpticalFlowInputLevel4, "OPTICALFLOW_OpticalFlowInputLevel4-", opticalFlowInputTextureSize, GraphicsFormat.R8_UInt);
            opticalFlowInputTextureSize = new Vector2Int(opticalFlowInputTextureSize.x >> 1, opticalFlowInputTextureSize.y >> 1);
            CreateDoubleBufferedResource(OpticalFlowInputLevel5, "OPTICALFLOW_OpticalFlowInputLevel5-", opticalFlowInputTextureSize, GraphicsFormat.R8_UInt);
            opticalFlowInputTextureSize = new Vector2Int(opticalFlowInputTextureSize.x >> 1, opticalFlowInputTextureSize.y >> 1);
            CreateDoubleBufferedResource(OpticalFlowInputLevel6, "OPTICALFLOW_OpticalFlowInputLevel6-", opticalFlowInputTextureSize, GraphicsFormat.R8_UInt);
            
            CreateDoubleBufferedResource(OpticalFlow, "OPTICALFLOW_OpticalFlow", opticalFlowTextureSize, GraphicsFormat.R16G16_SInt);
            CreateDoubleBufferedResource(OpticalFlowLevel1, "OPTICALFLOW_OpticalFlowLevel1-", opticalFlowLevel1TextureSize, GraphicsFormat.R16G16_SInt);
            CreateDoubleBufferedResource(OpticalFlowLevel2, "OPTICALFLOW_OpticalFlowLevel2-", opticalFlowLevel2TextureSize, GraphicsFormat.R16G16_SInt);
            CreateDoubleBufferedResource(OpticalFlowLevel3, "OPTICALFLOW_OpticalFlowLevel3-", opticalFlowLevel3TextureSize, GraphicsFormat.R16G16_SInt);
            CreateDoubleBufferedResource(OpticalFlowLevel4, "OPTICALFLOW_OpticalFlowLevel4-", opticalFlowLevel4TextureSize, GraphicsFormat.R16G16_SInt);
            CreateDoubleBufferedResource(OpticalFlowLevel5, "OPTICALFLOW_OpticalFlowLevel5-", opticalFlowLevel5TextureSize, GraphicsFormat.R16G16_SInt);
            CreateDoubleBufferedResource(OpticalFlowLevel6, "OPTICALFLOW_OpticalFlowLevel6-", opticalFlowLevel6TextureSize, GraphicsFormat.R16G16_SInt);

            OpticalFlowSCDHistogram = CreateResource("OPTICALFLOW_OpticalFlowSCDHistogram", new Vector2Int(GetSCDHistogramTextureWidth(), 1), GraphicsFormat.R32_UInt);
            OpticalFlowSCDPreviousHistogram = CreateResource("OPTICALFLOW_OpticalFlowSCDPreviousHistogram", new Vector2Int(GetSCDHistogramTextureWidth(), 1), GraphicsFormat.R32_SFloat);
            OpticalFlowSCDTemp = CreateResource("OPTICALFLOW_OpticalFlowSCDTemp", new Vector2Int(3, 1), GraphicsFormat.R32_UInt);
        }

        public void Destroy()
        {
            DestroyResource(ref OpticalFlowSCDTemp);
            DestroyResource(ref OpticalFlowSCDPreviousHistogram);
            DestroyResource(ref OpticalFlowSCDHistogram);
            
            DestroyResource(OpticalFlowLevel6);
            DestroyResource(OpticalFlowLevel5);
            DestroyResource(OpticalFlowLevel4);
            DestroyResource(OpticalFlowLevel3);
            DestroyResource(OpticalFlowLevel2);
            DestroyResource(OpticalFlowLevel1);
            DestroyResource(OpticalFlow);
            
            DestroyResource(OpticalFlowInputLevel6);
            DestroyResource(OpticalFlowInputLevel5);
            DestroyResource(OpticalFlowInputLevel4);
            DestroyResource(OpticalFlowInputLevel3);
            DestroyResource(OpticalFlowInputLevel2);
            DestroyResource(OpticalFlowInputLevel1);
            DestroyResource(OpticalFlowInput);
        }

        private static Vector2Int NextLevelSize(Vector2Int size)
        {
            return new Vector2Int(AlignUp(size.x, 2) / 2, AlignUp(size.y, 2) / 2);
        }

        private static int AlignUp(int x, int y)
        {
            return (x + (y - 1)) & ~(y - 1);
        }

        // TODO: move these to OpticalFlow class
        internal static Vector2Int GetOpticalFlowTextureSize(Vector2Int displaySize, int opticalFlowBlockSize)
        {
            int width = (displaySize.x + opticalFlowBlockSize - 1) / opticalFlowBlockSize;
            int height = (displaySize.y + opticalFlowBlockSize - 1) / opticalFlowBlockSize;
            return new Vector2Int(width, height);
        }

        private static Vector2Int GetOpticalFlowHistogramSize(int level)
        {
            const int searchRadius = 8;
            int maxVelocity = searchRadius * (1 << (OpticalFlowMaxPyramidLevels - 1 - level));
            int binsPerDimension = 2 * maxVelocity + 1;
            return new Vector2Int(binsPerDimension, binsPerDimension);
        }

        private static Vector2Int GetGlobalMotionSearchDispatchSize(int level)
        {
            const int threadGroupSizeX = 16;
            const int threadGroupSizeY = 16;
            Vector2Int opticalFlowHistogramSize = GetOpticalFlowHistogramSize(level);
            int additionalElementsDueToShiftsX = opticalFlowHistogramSize.x / threadGroupSizeX;
            int additionalElementsDueToShiftsY = opticalFlowHistogramSize.y / threadGroupSizeY;
            int dispatchX = (opticalFlowHistogramSize.x + additionalElementsDueToShiftsX + threadGroupSizeX - 1) / threadGroupSizeX;
            int dispatchY = (opticalFlowHistogramSize.y + additionalElementsDueToShiftsY + threadGroupSizeY - 1) / threadGroupSizeY;
            return new Vector2Int(dispatchX, dispatchY);
        }

        private static int GetSCDHistogramTextureWidth()
        {
            return HistogramBins * (HistogramsPerDim * HistogramsPerDim);
        }

        private static RenderTexture CreateResource(string name, Vector2Int size, GraphicsFormat format)
        {
            var rt = new RenderTexture(size.x, size.y, 0, format) { name = name, enableRandomWrite = true };
            rt.Create();
            return rt;
        }
        
        private static void CreateDoubleBufferedResource(RenderTexture[] resource, string name, Vector2Int size, GraphicsFormat format)
        {
            for (int i = 0; i < 2; ++i)
            {
                resource[i] = new RenderTexture(size.x, size.y, 0, format) { name = name + (i + 1), enableRandomWrite = true };
                resource[i].Create();
            }
        }
        
        private static void DestroyResource(ref Texture2D resource)
        {
            if (resource == null)
                return;
            
#if UNITY_EDITOR
            if (Application.isPlaying && !UnityEditor.EditorApplication.isPaused)
                UnityEngine.Object.Destroy(resource);
            else
                UnityEngine.Object.DestroyImmediate(resource);
#else
            UnityEngine.Object.Destroy(resource);
#endif
            resource = null;
        }

        private static void DestroyResource(ref RenderTexture resource)
        {
            if (resource == null)
                return;
            
            resource.Release();
            resource = null;
        }

        private static void DestroyResource(RenderTexture[] resource)
        {
            for (int i = 0; i < resource.Length; ++i)
                DestroyResource(ref resource[i]);
        }
    }
}