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

public class VoxelRaycastGpuManager
{
    ComputeShader raycastShader;

//---------------- Octree ----------------------------------------------------
    OctreeNode root;                 // assign your built octree root
    ComputeBuffer nodeBuffer;

    public LinearTree linearTree;
    
//--------------- RayCasts ------------------------------------------------------
    int kernel;

    ComputeBuffer datasBuffer;
    ComputeBuffer hitCounterBuffer = null;
    ComputeBuffer rayBuffer = null;
    ComputeBuffer hitBuffer;

    int raysPerBatch;

    int batchDataClassSize = Marshal.SizeOf(typeof(VoxelRaycastGPU.BatchData));

    int groupsX;

    int maxRaycastPerIteration;
    

    //--------------- Clustering ------------------------------------------------------

    ComputeShader clusteringShader;

    int gridSize;
    int threadCount = 64;
    int groupCountHits;
    int groupCountGrid;

    ComputeBuffer cellClusteringSums;
    ComputeBuffer cellClusteringCounts;
    ComputeBuffer clustered;


    //---------------------------------------------------------------------------------

    public VoxelRaycastGpuManager(ComputeShader computeShader, ComputeShader clusteringShader, OctreeNode octreeRoot)
    {
        raycastShader = computeShader;
        this.clusteringShader = clusteringShader;
        root = octreeRoot;
    }

    public VoxelRaycastGPU.BatchData[] Raycast(in VoxelRaycastGPU.BatchData[] batchData, int datasLenght)
    {
        ComputeBuffer countBuffer = new ComputeBuffer(1, sizeof(int), ComputeBufferType.Raw);

        int iteration = 0;
        int currentCount = datasLenght;
        int previousCount = datasLenght;

        datasBuffer.SetData(batchData, 0, 0, currentCount);

        while (iteration < 5 && currentCount > 0)
        {
            previousCount = currentCount;

            hitBuffer.SetCounterValue(0);

            raycastShader.SetBuffer(kernel, "batchDatas", datasBuffer );
            raycastShader.SetBuffer(kernel, "hits", hitBuffer);

            int threadsY = 8;

            int groupsY = Mathf.CeilToInt((float)currentCount / threadsY);

            Stopwatch sw = Stopwatch.StartNew();
            raycastShader.Dispatch(kernel, groupsX, groupsY, 1);

            ComputeBuffer.CopyCount(hitBuffer, countBuffer, 0);
            int[] countArr = new int[1];
            countBuffer.GetData(countArr);
            currentCount = countArr[0];

            sw.Stop();
            UnityEngine.Debug.Log($"Dispatch done in {sw.Elapsed.TotalMilliseconds}ms for {previousCount*raysPerBatch} casts retrieving {currentCount} hits");

            if (currentCount > 0)
            {
                if (currentCount * raysPerBatch > maxRaycastPerIteration && iteration < 5)
                {
                    sw = Stopwatch.StartNew();
                    currentCount = Clustering(currentCount);
                    sw.Stop();

                    UnityEngine.Debug.Log($"Clustering done in {sw.Elapsed.TotalMilliseconds}ms for {currentCount} casts");

                    VoxelRaycastGPU.BatchData[] hits = new VoxelRaycastGPU.BatchData[currentCount];
                    datasBuffer.GetData(hits, 0, 0, currentCount);

                    for (int i = 0; i < currentCount; i++)
                    {
                        GameObject sphere = GameObject.CreatePrimitive(PrimitiveType.Sphere);
                        sphere.transform.position = hits[i].origin;
                        sphere.transform.localScale = Vector3.one * 0.5f;
                    }
            
                }
                else
                {

                    datasBuffer = hitBuffer;
                }
            }

            iteration++;
        }

        VoxelRaycastGPU.BatchData[] result = new VoxelRaycastGPU.BatchData[previousCount];
        if (currentCount == 0)
            datasBuffer.GetData(result, 0, 0, previousCount);
        else
            hitBuffer.GetData(result, 0, 0, previousCount);


        countBuffer.Release();

        return result;
    }

    public void Init(int nbRaysPerBatch, in VoxelRaycastGPU.Ray[] rays)
    {
        maxRaycastPerIteration = 1000000;
        raysPerBatch = nbRaysPerBatch;

        // Flatten octree
        linearTree = OctreeGpuHelpers.FlattenOctree(root);
        int nodeStride = Marshal.SizeOf(typeof(LinearNode)); // should be 64

        hitBuffer = new ComputeBuffer(maxRaycastPerIteration * raysPerBatch, batchDataClassSize, ComputeBufferType.Append);
        datasBuffer = new ComputeBuffer(maxRaycastPerIteration, batchDataClassSize, ComputeBufferType.Default);

        rayBuffer = new ComputeBuffer(rays.Length, Marshal.SizeOf(typeof(VoxelRaycastGPU.Ray)), ComputeBufferType.Default);
        rayBuffer.SetData(rays, 0, 0, rays.Length);

        // Create GPU buffer for nodes
        nodeBuffer = new ComputeBuffer(linearTree.nodes.Length, nodeStride, ComputeBufferType.Default);
        nodeBuffer.SetData(linearTree.nodes);

        hitCounterBuffer = new ComputeBuffer(1, sizeof(int), ComputeBufferType.Raw);
        uint[] counterInit = { 0 };
        counterInit[0] = 0;
        hitCounterBuffer.SetData(counterInit);

        kernel = raycastShader.FindKernel("CSMain");

        raycastShader.SetBuffer(kernel, "nodes", nodeBuffer);
        raycastShader.SetBuffer(kernel, "hitCount", hitCounterBuffer);
        raycastShader.SetBuffer(kernel, "rays", rayBuffer);

        raycastShader.SetInt("raysPerBatch", nbRaysPerBatch);
        raycastShader.SetInt("rootIndex", linearTree.rootIndex);
        raycastShader.SetInt("nodeCount", linearTree.nodes.Length);

        raycastShader.SetFloat("rootHalfSize", root.bounds.size.x / 2f);
        raycastShader.SetFloats("rootCenter", new float[3] { root.bounds.center.x, root.bounds.center.y, root.bounds.center.z });

        groupsX = Mathf.CeilToInt((float)raysPerBatch / 8);

        gridSize = maxRaycastPerIteration / raysPerBatch;
        groupCountGrid = Mathf.CeilToInt((float)gridSize / threadCount);

        cellClusteringSums = new ComputeBuffer(gridSize, sizeof(float) * 4);
        cellClusteringCounts = new ComputeBuffer(gridSize, sizeof(uint));
        clustered = new ComputeBuffer(gridSize, Marshal.SizeOf(typeof(VoxelRaycastGPU.BatchData)), ComputeBufferType.Append);
    }

    public int Clustering( int inputCount )
    {
        groupCountHits = Mathf.CeilToInt((float)inputCount / threadCount);

        clustered.SetCounterValue(0);
        float[] zero4 = new float[gridSize * 4];
        uint[] zeroU = new uint[gridSize];
        cellClusteringSums.SetData(zero4);
        cellClusteringCounts.SetData(zeroU);

        // Shader
        int kernelAcc = clusteringShader.FindKernel("AccumulateHits");
        int kernelMean = clusteringShader.FindKernel("ComputeMeans");

        clusteringShader.SetInt("inputCount", inputCount);
        clusteringShader.SetFloat("cellSize", 1.0f);  // taille des cellules
        clusteringShader.SetVector("gridOrigin", Vector3.zero);

        clusteringShader.SetBuffer(kernelAcc, "inputHits", hitBuffer);
        clusteringShader.SetBuffer(kernelAcc, "cellSums", cellClusteringSums);
        clusteringShader.SetBuffer(kernelAcc, "cellCounts", cellClusteringCounts);
        clusteringShader.Dispatch(kernelAcc, groupCountHits, 1, 1);

        clusteringShader.SetBuffer(kernelMean, "cellSums", cellClusteringSums);
        clusteringShader.SetBuffer(kernelMean, "cellCounts", cellClusteringCounts);
        clusteringShader.SetBuffer(kernelMean, "clusteredHits", clustered);
        clusteringShader.Dispatch(kernelMean, groupCountGrid, 1, 1);

        // Lecture du résultat
        ComputeBuffer countBuffer = new ComputeBuffer(1, sizeof(int), ComputeBufferType.Raw);
        ComputeBuffer.CopyCount(clustered, countBuffer, 0);
        int[] countArr = new int[1];
        countBuffer.GetData(countArr);
        int clusterCount = countArr[0];

        datasBuffer = clustered;

        return clusterCount;
    }

    ~VoxelRaycastGpuManager()
    {
        if (hitCounterBuffer != null)
            hitCounterBuffer.Release();

        if (rayBuffer != null)
            rayBuffer.Release();

        hitBuffer.Release();
        datasBuffer.Release();
    }

}