using System;
using System.Runtime.InteropServices;
using UnityEngine;

namespace FidelityFX.FrameGen
{
    public static class OpticalFlow
    {
        public static readonly int MinBlockSize = 8;
        internal const int OpticalFlowMaxPyramidLevels = 7;
        internal const int HistogramBins = 256;
        internal const int HistogramsPerDim = 3;
        internal const int HistogramShifts = 3;
        
        public static OpticalFlowContext CreateContext(Vector2Int resolution, OpticalFlowShaders shaders)
        {
            Debug.Log($"Setting up Optical Flow with resolution: {resolution.x}x{resolution.y}");
            
            var contextDescription = new ContextDescription
            {
                resolution = resolution,
                shaders = shaders,
            };
            
            var context = new OpticalFlowContext();
            context.Create(contextDescription);
            return context;
        }
        
        public struct ContextDescription
        {
            public Vector2Int resolution;
            public OpticalFlowShaders shaders;
        }

        public class DispatchDescription
        {
            public ResourceView Color;
            public ResourceView OpticalFlowVector;
            public ResourceView OpticalFlowSCD;
            public bool Reset;
            
            public BackbufferTransferFunction BackbufferTransferFunction;
            public Vector2 MinMaxLuminance;
        }

        public enum BackbufferTransferFunction
        {
            LDR_sRGB,
            HDR_PQ,
            HDR_scRGB,
        }

        public static Vector2Int GetOpticalFlowTextureSize(Vector2Int displaySize)
        {
            return GetOpticalFlowTextureSize(displaySize, MinBlockSize);
        }
        
        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);
        }

        internal 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);
        }

        internal 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);
        }

        internal static int GetSCDHistogramTextureWidth()
        {
            return HistogramBins * (HistogramsPerDim * HistogramsPerDim);
        }
        
        [Serializable, StructLayout(LayoutKind.Sequential)]
        internal struct OpticalFlowConstants
        {
            public Vector2Int inputLumaResolution;
            public uint opticalFlowPyramidLevel;
            public uint opticalFlowPyramidLevelCount;

            public int frameIndex;
            public uint backbufferTransferFunction;
            public Vector2 minMaxLuminance;
        }
        
        [Serializable, StructLayout(LayoutKind.Sequential)]
        internal struct SpdConstants
        {
            public FfxSpd.SpdConstants spd;

            public uint numWorkGroupsOpticalFlowInputPyramid;
            public uint pad0_;
            public uint pad1_;
            public uint pad2_;
        }
    }
}