using System;
using System.Reflection;
using System.Runtime.InteropServices;
using System.Threading;
using System.Windows.Forms;
using SharpDX;
using SharpDX.Direct3D9;

namespace Blacklight.Bezier.D3D9
{
    /// <summary>
    /// Runs one pixel per sample through a variant on the default adapter, for readback or for timing.
    /// </summary>
    public sealed class D3D9Bench : IDisposable
    {
        private static readonly Vector4[] FullScreenTriangle =
        {
            new Vector4(-1, -1, 0, 1),
            new Vector4(3, -1, 0, 1),
            new Vector4(-1, 3, 0, 1),
        };

        // D3DGETDATA_FLUSH, so a poll can't spin on a query the driver never submitted.
        private const int GetDataFlush = 1;

        // Query.GetData<T> can't tell "not ready yet" from a result, this can.
        private static readonly Func<Query, IntPtr, int, int, Result> RawGetData =
            (Func<Query, IntPtr, int, int, Result>)Delegate.CreateDelegate(
                typeof(Func<Query, IntPtr, int, int, Result>),
                typeof(Query).GetMethod(nameof(Query.GetData),
                    BindingFlags.Instance | BindingFlags.NonPublic | BindingFlags.DeclaredOnly,
                    binder: null, types: new[] { typeof(IntPtr), typeof(int), typeof(int) }, modifiers: null));

        private readonly Form window;
        private readonly Direct3D direct3D;
        private readonly Device device;
        private readonly VertexDeclaration declaration;
        private readonly VertexShader vertexShader;

        public readonly string AdapterDescription;

        public D3D9Bench(string vertexShaderSource)
        {
            // Never shown: D3D9 needs a window to own the device, and nothing here is presented.
            window = new Form();
            direct3D = new Direct3D();
            AdapterDescription = direct3D.Adapters[0].Details.Description.Trim();

            if (!direct3D.CheckDeviceFormat(0, DeviceType.Hardware, Format.X8R8G8B8, Usage.RenderTarget,
                ResourceType.Texture, Format.R32F))
            {
                throw new NotSupportedException(AdapterDescription + " can't render to R32F.");
            }

            device = new Device(direct3D, 0, DeviceType.Hardware, window.Handle,
                // FpuPreserve: otherwise D3D9 drops the x87 FPU to single precision under the double reference.
                CreateFlags.HardwareVertexProcessing | CreateFlags.FpuPreserve, new PresentParameters(1, 1));

            declaration = new VertexDeclaration(device, new[]
            {
                new VertexElement(0, 0, DeclarationType.Float4, DeclarationMethod.Default,
                    DeclarationUsage.Position, 0),
                VertexElement.VertexDeclarationEnd,
            });
            
            // D3DX9's compiler, the one production D3D9 shaders go through.
            vertexShader = new VertexShader(device, ShaderBytecode.Compile(vertexShaderSource, "FullScreenVS", "vs_3_0",
                ShaderFlags.OptimizationLevel3));
        }

        /// <summary>
        /// Distance evaluations TimingPS runs per pixel; matches EvaluationsPerTimedPixel in Harness.hlsl.
        /// </summary>
        public const int EvaluationsPerTimedPixel = 16;

        /// <param name="entryPoint">DistancePS to read distances back, TimingPS to time them.</param>
        public PixelShader CompilePixelShader(Variant variant, string entryPoint)
        {
            return new PixelShader(device, ShaderBytecode.Compile(variant.Source, entryPoint, "ps_3_0",
                ShaderFlags.OptimizationLevel3));
        }

        /// <summary>
        /// A side x side texture of (x, y, phi, end), each rounded to float32 on its own; samples repeat to fill it.
        /// </summary>
        /// <param name="fourth">Replaces end in the last channel when given, as g does for shell tests.</param>
        public Texture UploadSamples(ShellSample[] samples, int side, double[] fourth = null)
        {
            var returnMe = new Texture(device, side, side, 1, Usage.None, Format.A32B32G32R32F, Pool.Managed);
            DataRectangle rectangle = returnMe.LockRectangle(0, LockFlags.None);
            var row = new float[side * 4];
            for (int rowIndex = 0; rowIndex < side; rowIndex++)
            {
                for (int column = 0; column < side; column++)
                {
                    int index = (rowIndex * side + column) % samples.Length;
                    ShellSample sample = samples[index];
                    row[column * 4] = (float)sample.X;
                    row[column * 4 + 1] = (float)sample.Y;
                    row[column * 4 + 2] = (float)sample.Phi;
                    row[column * 4 + 3] = (float)(fourth == null ? sample.End : fourth[index]);
                }
                
                Marshal.Copy(row, 0, IntPtr.Add(rectangle.DataPointer, rowIndex * rectangle.Pitch), row.Length);
            }

            returnMe.UnlockRectangle(0);
            return returnMe;
        }

        /// <summary>
        /// A side x side texture with one value per texel in .x, and optionally a second in .y, repeating to fill it.
        /// </summary>
        public Texture UploadValues(double[] values, int side, double[] secondValues = null)
        {
            var returnMe = new Texture(device, side, side, 1, Usage.None, Format.A32B32G32R32F, Pool.Managed);
            DataRectangle rectangle = returnMe.LockRectangle(0, LockFlags.None);
            var row = new float[side * 4];
            for (int rowIndex = 0; rowIndex < side; rowIndex++)
            {
                for (int column = 0; column < side; column++)
                {
                    int index = (rowIndex * side + column) % values.Length;
                    row[column * 4] = (float)values[index];
                    
                    if (secondValues != null)
                    {
                        row[column * 4 + 1] = (float)secondValues[index];
                    }
                }

                Marshal.Copy(row, 0, IntPtr.Add(rectangle.DataPointer, rowIndex * rectangle.Pitch), row.Length);
            }

            returnMe.UnlockRectangle(0);
            return returnMe;
        }

        private void PrepareDraw(PixelShader shader, Texture[] samples, int side, Surface target)
        {
            device.SetRenderTarget(0, target);
            // The device's own depth buffer is 1x1; one smaller than the target makes every draw invalid.
            device.DepthStencilSurface = null;
            device.SetRenderState(RenderState.ZEnable, false);
            device.SetRenderState(RenderState.AlphaBlendEnable, false);
            device.SetRenderState(RenderState.CullMode, Cull.None);

            device.VertexDeclaration = declaration;
            device.VertexShader = vertexShader;
            device.PixelShader = shader;

            for (int stage = 0; stage < samples.Length; stage++)
            {
                device.SetTexture(stage, samples[stage]);
                device.SetSamplerState(stage, SamplerState.MinFilter, TextureFilter.Point);
                device.SetSamplerState(stage, SamplerState.MagFilter, TextureFilter.Point);
                device.SetSamplerState(stage, SamplerState.MipFilter, TextureFilter.None);
                device.SetSamplerState(stage, SamplerState.AddressU, TextureAddress.Clamp);
                device.SetSamplerState(stage, SamplerState.AddressV, TextureAddress.Clamp);
            }

            device.SetPixelShaderConstant(0, new[] { new Vector4(side, side, 1.0f / side, 1.0f / side) });
        }

        private void DrawPass()
        {
            device.DrawUserPrimitives(PrimitiveType.TriangleList, 1, FullScreenTriangle);
        }

        /// <summary>
        /// Every sample's distance, as the variant computed it, in sample order.
        /// </summary>
        public float[] Measure(PixelShader shader, Texture samples, int side)
        {
            return Measure(shader, new[] { samples }, side);
        }

        /// <param name="samples">Bound to samplers s0, s1, ... in order.</param>
        public float[] Measure(PixelShader shader, Texture[] samples, int side)
        {
            using (var target = new Texture(device, side, side, 1, Usage.RenderTarget, Format.R32F, Pool.Default))
            using (Surface targetSurface = target.GetSurfaceLevel(0))
            using (Surface readback = Surface.CreateOffscreenPlain(device, side, side, Format.R32F, Pool.SystemMemory))
            {
                PrepareDraw(shader, samples, side, targetSurface);
                device.BeginScene();
                DrawPass();
                device.EndScene();
                device.GetRenderTargetData(targetSurface, readback);

                var returnMe = new float[side * side];
                DataRectangle rectangle = readback.LockRectangle(LockFlags.ReadOnly);
                
                for (int row = 0; row < side; row++)
                {
                    Marshal.Copy(IntPtr.Add(rectangle.DataPointer, row * rectangle.Pitch), returnMe, row * side, side);
                }
                
                readback.UnlockRectangle();
                return returnMe;
            }
        }

        private static long WaitForQuery(Query query, IntPtr buffer, int size)
        {
            Result result;
            
            while ((result = RawGetData(query, buffer, size, GetDataFlush)) != Result.Ok)
            {
                result.CheckError();
                Thread.SpinWait(64);
            }
            
            if (size == sizeof(int))
            {
                return Marshal.ReadInt32(buffer);
            }
            
            return Marshal.ReadInt64(buffer);
        }

        /// <summary>
        /// Seconds the GPU spent on passes full-screen draws, or null if the clock jumped mid-measurement.
        /// </summary>
        private double? TimeOnce(PixelShader shader, Texture[] samples, int side, Surface target, int passes,
            Query disjoint, Query frequency, Query start, Query finish, IntPtr buffer)
        {
            PrepareDraw(shader, samples, side, target);
            disjoint.Issue(Issue.Begin);
            start.Issue(Issue.End);
            device.BeginScene();
            
            for (int pass = 0; pass < passes; pass++)
            {
                DrawPass();
            }
            
            device.EndScene();
            finish.Issue(Issue.End);
            disjoint.Issue(Issue.End);
            frequency.Issue(Issue.End);

            long startTicks = WaitForQuery(start, buffer, sizeof(long));
            long finishTicks = WaitForQuery(finish, buffer, sizeof(long));
            long ticksPerSecond = WaitForQuery(frequency, buffer, sizeof(long));
            
            if (WaitForQuery(disjoint, buffer, sizeof(int)) != 0)
            {
                return null;
            }
            
            return (double)(finishTicks - startTicks) / ticksPerSecond;
        }

        /// <summary>
        /// Median GPU nanoseconds per pixel for each shader, or null if the driver can't time.
        /// </summary>
        /// <param name="progress">Called after each measurement with how many are done and how many there are.</param>
        public double[] NanosecondsPerPixel(PixelShader[] shaders, Texture samples, int side, int passes, int repeats,
            Action<int, int> progress)
        {
            return NanosecondsPerPixel(shaders, new[] { samples }, side, passes, repeats, progress);
        }

        /// <param name="samples">Bound to samplers s0, s1, ... in order.</param>
        public double[] NanosecondsPerPixel(PixelShader[] shaders, Texture[] samples, int side, int passes,
            int repeats, Action<int, int> progress)
        {
            Query disjoint;
            Query frequency;
            Query start;
            Query finish;
            try
            {
                disjoint = new Query(device, QueryType.TimestampDisjoint);
                frequency = new Query(device, QueryType.TimestampFreq);
                start = new Query(device, QueryType.Timestamp);
                finish = new Query(device, QueryType.Timestamp);
            }
            catch (SharpDXException)
            {
                return null;
            }

            IntPtr buffer = Marshal.AllocHGlobal(sizeof(long));
            try
            {
                using (disjoint)
                using (frequency)
                using (start)
                using (finish)
                using (var target = new Texture(device, side, side, 1, Usage.RenderTarget, Format.R32F, Pool.Default))
                using (Surface targetSurface = target.GetSurfaceLevel(0))
                {
                    var perPixel = new double[shaders.Length][];
                    for (int shaderIndex = 0; shaderIndex < shaders.Length; shaderIndex++)
                    {
                        perPixel[shaderIndex] = new double[repeats];
                    }

                    // Shaders take turns within each repeat, to make comparisons apples to apples.
                    int measurements = (repeats + 1) * shaders.Length;
                    int measured = 0;
                    for (int repeat = -1; repeat < repeats; repeat++)
                    {
                        for (int shaderIndex = 0; shaderIndex < shaders.Length; shaderIndex++)
                        {
                            double? seconds = null;
                            for (int attempt = 0; attempt < 10 && seconds == null; attempt++)
                            {
                                seconds = TimeOnce(shaders[shaderIndex], samples, side, targetSurface, passes,
                                    disjoint, frequency, start, finish, buffer);
                            }

                            if (repeat >= 0)
                            {
                                perPixel[shaderIndex][repeat] = (seconds ?? double.NaN) * 1e9
                                    / ((double)passes * side * side);
                            }

                            measured++;
                            progress(measured, measurements);
                        }
                    }

                    var returnMe = new double[shaders.Length];
                    for (int shaderIndex = 0; shaderIndex < shaders.Length; shaderIndex++)
                    {
                        Array.Sort(perPixel[shaderIndex]);
                        returnMe[shaderIndex] = perPixel[shaderIndex][repeats / 2];
                    }
                    
                    return returnMe;
                }
            }
            finally
            {
                Marshal.FreeHGlobal(buffer);
            }
        }

        public void Dispose()
        {
            vertexShader.Dispose();
            declaration.Dispose();
            device.Dispose();
            direct3D.Dispose();
            window.Dispose();
        }
    }
}