using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Microsoft.WindowsAPICodePack.DirectX.Direct3D11;
using Microsoft.WindowsAPICodePack.DirectX.Direct3D;
using System.IO;
using System.Runtime.InteropServices;
using Microsoft.WindowsAPICodePack.DirectX.DXGI;

namespace TestSimple
{
    struct BuffType
    {
        public int I;
        public float F;

        /// <summary>
        /// Initializes a new instance of the BuffType structure.
        /// </summary>
        /// <param name="i"></param>
        /// <param name="f"></param>
        public BuffType(int i, float f)
        {
            I = i;
            F = f;
        }
    }

    class Program
    {
        static int ELEMENT_SIZE = Marshal.SizeOf(typeof(BuffType));

        static void Main(string[] args)
        {
            const int NUM_ELEMENTS = 1024;

            FeatureLevel[] level = new FeatureLevel[]{ FeatureLevel.FeatureLevel_10_0 };
            using (D3DDevice device = D3DDevice.CreateDevice(null, DriverType.Hardware, null, CreateDeviceFlag.SingleThreaded | CreateDeviceFlag.Debug, level))
            {
                DeviceContext context = device.GetImmediateContext();
                Console.WriteLine(device.DeviceFeatureLevel);
                FeatureDataD3D10XHardwareOptions options;
                if (device.CheckFeatureDataD3D10XHardwareOptions(out options))
                    Console.WriteLine(options.ComputeShadersPlusRawAndStructuredBuffersViaShader4x);
                else
                    Console.WriteLine("Failed to check hw options");

                ComputeShader shader = null;
                try
                {
                    
                    using (Stream hlsl = File.OpenRead("unit.fx"))
                        shader = device.CreateComputeShader(hlsl);
                    BuffType[] buffer;

                    buffer = new BuffType[NUM_ELEMENTS];
                    for (int i = 0; i < NUM_ELEMENTS; i++)
                    {
                    	buffer[i] = new BuffType { I = i, F = i };
                    }

                    using (D3DBuffer buf0 = CreateBuffer(device, buffer, true))
                    using (D3DBuffer buf1 = CreateBuffer(device, buffer, true))
                    using (D3DBuffer buf2 = CreateBuffer(device, buffer, false))
                    {
                        using (ShaderResourceView view0 = CreateView(device, buf0))
                        using (ShaderResourceView view1 = CreateView(device, buf1))
                        using (UnorderedAccessView view2 = CreateUnorderedAccessView(device, buf2))
                        {
                            RunComputerShader(context, shader, new ShaderResourceView[] { view0, view1 }, new UnorderedAccessView[] { view2 }, NUM_ELEMENTS, 1, 1);
                            D3DBuffer debugbuf = CreateAndCopyToDebugBuf(device, context, buf2);
                            MappedSubresource mappedResource = (MappedSubresource)context.Map(debugbuf, 0, Map.Read, MapFlag.Unspecified);
                            IntPtr p = mappedResource.Data;
                            int cursor = p.ToInt32();
                            for (int i = 0; i < NUM_ELEMENTS; i++)
                            {
                                p = new IntPtr(cursor);
                                BuffType bt = (BuffType)Marshal.PtrToStructure(p, typeof(BuffType));
                                cursor += ELEMENT_SIZE;
                                if (bt.I != buffer[i].I * 2 || bt.F != buffer[i].F * 2)
                                {
                                    Console.WriteLine("Failed");
                                }
                            }
                        }
                    }
                }
                finally
                {
                    if (shader != null)
                        shader.Dispose();
                }
            }
            Console.WriteLine("End");
            Console.ReadLine();
        }

        static D3DBuffer CreateAndCopyToDebugBuf(D3DDevice device, DeviceContext context, D3DBuffer buffer)
        {
            BufferDescription desc = buffer.Description;
            desc.CpuAccessFlags = CpuAccessFlag.Read;
            desc.Usage = Usage.Staging;
            desc.BindFlags = BindFlag.Undefined;
            desc.MiscFlags = ResourceMiscFlag.Undefined;
            D3DBuffer result = device.CreateBuffer(desc);
            context.CopyResource(result, buffer);
            return result;
        }

        static void RunComputerShader(DeviceContext context, ComputeShader shader, ShaderResourceView[] views, UnorderedAccessView[] unordered, uint x, uint y, uint z)
        {
            ComputeShaderPipelineStage cs = context.CS;
            cs.SetShader(shader);
            cs.SetShaderResources(0, views);
            // 4 size of handle?
            cs.SetUnorderedAccessViews(0, unordered, new uint[] { 0 });
            context.Dispatch(x, y, z);
            // cs.SetUnorderedAccessViews(0, null, new uint[] { 0 });
            // cs.SetShaderResources(0, 1, null);
        }

        static ShaderResourceView CreateView(D3DDevice device, D3DBuffer buffer)
        {
            ShaderResourceViewDescription desc = new ShaderResourceViewDescription
            {
                ViewDimension =  ShaderResourceViewDimension.BufferEx
            };
            desc.BufferEx.FirstElement = 0;
            if ((buffer.Description.MiscFlags & ResourceMiscFlag.BufferAllowRawViews) == ResourceMiscFlag.BufferAllowRawViews)
            {
                desc.Format = Format.R32_TYPELESS;
                desc.BufferEx.Flags = BufferExShaderResourceViewFlag.Raw;
                desc.BufferEx.NumElements = buffer.Description.ByteWidth / 4;
            }
            else if ((buffer.Description.MiscFlags & ResourceMiscFlag.BufferStructured) == ResourceMiscFlag.BufferStructured)
            {
                desc.Format = Format.Unknown;
                desc.BufferEx.NumElements = buffer.Description.ByteWidth / buffer.Description.StructureByteStride;
            }
            else
                throw new Exception("Unsupported buffer format");
            return device.CreateShaderResourceView(buffer, desc);
        }

        static UnorderedAccessView CreateUnorderedAccessView(D3DDevice device, D3DBuffer buffer)
        {
            UnorderedAccessViewDescription desc = new UnorderedAccessViewDescription
            {
                ViewDimension =  UnorderedAccessViewDimension.Buffer
            };
            desc.Buffer.FirstElement = 0;

            if ((buffer.Description.MiscFlags & ResourceMiscFlag.BufferAllowRawViews) == ResourceMiscFlag.BufferAllowRawViews)
            {
                desc.Format = Format.R32_TYPELESS;
                desc.Buffer.Flags = BufferUnorderedAccessViewFlag.Raw;
                desc.Buffer.NumElements = buffer.Description.ByteWidth / 4;
            }
            else if ((buffer.Description.MiscFlags & ResourceMiscFlag.BufferStructured) == ResourceMiscFlag.BufferStructured)
            {
                desc.Format = Format.Unknown;
                desc.Buffer.NumElements = buffer.Description.ByteWidth / buffer.Description.StructureByteStride;
            }
            else
                throw new Exception("Unsupported buffer format");
            return device.CreateUnorderedAccessView(buffer, desc);
        }

        static D3DBuffer CreateBuffer(D3DDevice device, BuffType[] buffer, bool init)
        {
            BufferDescription desc = new BufferDescription
                {
                    BindFlags = BindFlag.UnorderedAccess | BindFlag.ShaderResource,
                    ByteWidth = (uint)(ELEMENT_SIZE * buffer.Length),
                    MiscFlags = ResourceMiscFlag.BufferStructured,
                    StructureByteStride = (uint)ELEMENT_SIZE                    
                };
            IntPtr ptr = Marshal.AllocHGlobal((int)desc.ByteWidth);
            int current = ptr.ToInt32();
            foreach (BuffType bt in buffer)
            {
                IntPtr offset = new IntPtr(current);
                Marshal.StructureToPtr(bt, offset, false);
                current += ELEMENT_SIZE;
            }
            D3DBuffer result;
            SubresourceData initData;
            if (init)
            {
                initData = new SubresourceData
                {
                    SysMem = ptr
                };
                result = device.CreateBuffer(desc, initData);
            }
            else
            {
                initData = new SubresourceData();
                result = device.CreateBuffer(desc);
            }
            return result;
        }
    }
}
