using RayTracer;
using System.Collections;
using System.Collections.Generic;
using Unity.Collections;
using UnityEngine;
using UnityEngine.Profiling;

public class RaytraceGeneration : MonoBehaviour
{
    Generator generator;

    public int currentElement = 0;
    public int updateElement = 0;
    public int totalElements = 0;

    public int resolution = 8;

    public int rtSteps = 256;
    public int rtUpdate = 4096;

    public int bvhDebug = 0;

    RayTracer.Triangle[] triangles;
    RayTracer.BVH accel;

    bool prevBaking = false;

    ComputeBuffer trianglesBuffer;
    ComputeBuffer bvhNodesBuffer;
    ComputeBuffer bvhTrianglesBuffer;

    [SerializeField]
    ComputeShader rayTracer;

    void CalculateGPU()
    {
        rayTracer.SetBuffer(0, "_Triangles", trianglesBuffer);
        rayTracer.SetBuffer(0, "_BVHNodes", bvhNodesBuffer);
        rayTracer.SetBuffer(0, "_BVHIndices", bvhTrianglesBuffer);
        rayTracer.SetTexture(0, "_LightMap", generator.LightMap);
        rayTracer.SetTexture(0, "_RayOrigin", generator.PositionTexture);
        rayTracer.SetTexture(0, "_RayDirection", generator.NormalsTexture);
        rayTracer.SetInt("_Resolution", generator.settings.Resolution);
        rayTracer.SetFloat("_Time", Time.time);

        int workgroup = 8;
        int groups = Mathf.CeilToInt(generator.settings.Resolution / workgroup);

        rayTracer.Dispatch(0, groups, groups, 1);
    }

    // Start is called before the first frame update
    void Start()
    {
        generator = Object.FindObjectOfType<Generator>();
    }

    // Update is called once per frame
    void Update()
    {
        if (generator.IsBaking)
        {
            if (prevBaking == false)
            {
                totalElements = generator.settings.Resolution * generator.settings.Resolution;
                Prepare();
            }

            //Calculate(rtSteps);
            CalculateGPU();
        }

        prevBaking = generator.IsBaking;
    }

    public class StackElement
    {
        public Transform transform;
        public Matrix4x4 worldMatrix;

        public StackElement(Transform transform, Matrix4x4 worldMatrix)
        {
            this.transform = transform;
            this.worldMatrix = worldMatrix;
        }
    }

    void Prepare()
    {
        Profiler.BeginSample("LightMap - Prepare LightMap Objects Transform");

        // Get all lightmap objects and add them into DFS stack
        LightMap[] lightMaps = Object.FindObjectsOfType<LightMap>();
        Stack<StackElement> transforms = new Stack<StackElement>();
        for (int i = 0; i < lightMaps.Length; i++)
        {
            transforms.Push(new StackElement(lightMaps[i].transform, lightMaps[i].transform.localToWorldMatrix));
        }

        Profiler.EndSample();

        Profiler.BeginSample("LightMap - Count Triangles");

        // Loop through stack and count number of triangles
        int trianglesCount = 0;
        while (transforms.Count > 0)
        {
            // Pop Transform
            StackElement se = transforms.Pop();
            Transform t = se.transform;
            Matrix4x4 curr = se.worldMatrix * t.localToWorldMatrix;

            // Enqueue all children for search
            for (int i = 0; i < t.childCount; i++)
            {
                transforms.Push(new StackElement(t.GetChild(i), curr));
            }

            // If transform has game object, game object has mesh filter and mesh filter contains mesh - process that mesh into triangle list
            if (t.gameObject != null)
            {
                MeshFilter mf = t.GetComponent<MeshFilter>();

                if (mf != null)
                {
                    if (mf.sharedMesh != null)
                    {
                        Vector3[] vertices = mf.sharedMesh.vertices;

                        for (int j = 0; j < mf.sharedMesh.subMeshCount; j++)
                        {
                            int[] indices = mf.sharedMesh.GetTriangles(j);

                            trianglesCount += indices.Length / 3;
                        }
                    }
                }
            }
        }

        Profiler.EndSample();

        Profiler.BeginSample("LightMap - Restart Stack");

        // Clean tri list
        triangles = new RayTracer.Triangle[trianglesCount];

        // Restart stack
        for (int i = 0; i < lightMaps.Length; i++)
        {
            transforms.Push(new StackElement(lightMaps[i].transform, lightMaps[i].transform.localToWorldMatrix));
        }

        Profiler.EndSample();

        Profiler.BeginSample("LightMap - Add Triangles");

        // Add triangles
        int counter = 0;
        while (transforms.Count > 0)
        {
            // Pop Transform
            StackElement se = transforms.Pop();
            Transform t = se.transform;
            Matrix4x4 curr = se.worldMatrix * t.localToWorldMatrix;

            // Enqueue all children for search
            for (int i = 0; i < t.childCount; i++)
            {
                transforms.Push(new StackElement(t.GetChild(i), curr));
            }

            // If transform has game object, game object has mesh filter and mesh filter contains mesh - process that mesh into triangle list
            if (t.gameObject != null)
            {
                MeshFilter mf = t.GetComponent<MeshFilter>();

                if (mf != null)
                {
                    if (mf.sharedMesh != null)
                    {
                        Vector3[] vertices = mf.sharedMesh.vertices;

                        for (int j = 0; j < mf.sharedMesh.subMeshCount; j++)
                        {
                            int[] indices = mf.sharedMesh.GetTriangles(j);

                            for (int k = 0; k < indices.Length; k += 3)
                            {
                                int a = indices[k + 0];
                                int b = indices[k + 1];
                                int c = indices[k + 2];

                                Vector3 va = curr * new Vector4(vertices[a].x, vertices[a].y, vertices[a].z, 1.0f);
                                Vector3 vb = curr * new Vector4(vertices[b].x, vertices[b].y, vertices[b].z, 1.0f);
                                Vector3 vc = curr * new Vector4(vertices[c].x, vertices[c].y, vertices[c].z, 1.0f);

                                triangles[counter] = new RayTracer.Triangle(va, vb, vc);

                                counter++;
                            }
                        }
                    }
                }
            }
        }

        if (trianglesBuffer != null)
        {
            trianglesBuffer.Release();
            trianglesBuffer = null;
        }

        trianglesBuffer = new ComputeBuffer(counter, 36, ComputeBufferType.Structured, ComputeBufferMode.Immutable);
        trianglesBuffer.SetData(triangles);

        Profiler.EndSample();

        // Build BVH
        accel = new RayTracer.BVH(triangles);

        if (bvhNodesBuffer != null)
        {
            bvhNodesBuffer.Release();
            bvhNodesBuffer = null;
        }

        bvhNodesBuffer = new ComputeBuffer(accel.GetNodes().Count, 64, ComputeBufferType.Structured, ComputeBufferMode.Immutable);
        bvhNodesBuffer.SetData(accel.GetNodes().ToArray());

        if (bvhTrianglesBuffer != null)
        {
            bvhTrianglesBuffer.Release();
            bvhTrianglesBuffer = null;
        }

        bvhTrianglesBuffer = new ComputeBuffer(accel.GetPrimitiveIndexes().Count, 4, ComputeBufferType.Structured, ComputeBufferMode.Immutable);
        bvhTrianglesBuffer.SetData(accel.GetPrimitiveIndexes().ToArray());
    }

    private void OnDisable()
    {
        if (trianglesBuffer != null)
        {
            trianglesBuffer.Release();
            trianglesBuffer = null;
        }
    }

    void Calculate(int steps)
    {
        Profiler.BeginSample("Raytrace - Loop Samples");

        NativeArray<float> positions = generator.PositionTexture.GetRawTextureData<float>();
        NativeArray<float> normals = generator.NormalsTexture.GetRawTextureData<float>();

        for (int i = 0; i < steps; i++)
        {
            Profiler.BeginSample("Raytrace - Gather Position/Normals");

            Vector3 pos = new Vector3(positions[(currentElement + i) * 4 + 0], positions[(currentElement + i) * 4 + 1], positions[(currentElement + i) * 4 + 2]);
            Vector3 dir = new Vector3(normals[(currentElement + i) * 4 + 0], normals[(currentElement + i) * 4 + 1], normals[(currentElement + i) * 4 + 2]);

            if (Vector3.Dot(pos, pos) == 0.0f && Vector3.Dot(dir, dir) == 0.0f)
            {
                Profiler.EndSample();

                continue;
            }

            Profiler.EndSample();

            Profiler.BeginSample("Raytrace - Trace Ray");

            RayTracer.Ray r = new RayTracer.Ray(pos + dir * 0.01f, new Vector3(0.0f, 1.0f, 0.0f));

            /*float closestDistance = float.MaxValue;
            float closestU = 0.0f;
            float closestV = 0.0f;
            int closestIndex;
            bool closestHit = false;

            for (int j = 0; j < triangles.Count; j++)
            {
                RayTracer.IntersectResult result;
                result.baryU = 0.0f;
                result.baryV = 0.0f;
                result.distance = 0.0f;
                bool hit = RayTracer.Intersect.RayTriangle(r, triangles[j], ref result);

                if (hit == true)
                {
                    if (result.distance > 0.0f && result.distance < closestDistance)
                    {
                        closestDistance = result.distance;
                        closestU = result.baryU;
                        closestV = result.baryV;
                        closestHit = true;
                        closestIndex = j;
                    }
                }
            }

            float resultRed = 1.0f;
            float resultGreen = 1.0f;
            float resultBlue = 1.0f;

            if (closestHit)
            {
                resultRed = resultGreen = resultBlue = 0.0f;
            }*/

            RayTracer.IntersectResult result = new RayTracer.IntersectResult();

            int index = accel.IntersectRay(r, ref result);

            Profiler.EndSample();

            Profiler.BeginSample("Raytrace - Write Pixel");

            float resultRed = 1.0f;
            float resultGreen = 1.0f;
            float resultBlue = 1.0f;

            if (index >= 0)
            {
                resultRed = resultGreen = resultBlue = 0.0f;
            }

            /*generator.LightMap.SetPixel((currentElement + i) % generator.settings.Resolution,
                (currentElement + i) / generator.settings.Resolution,
                new Color(resultRed, resultGreen, resultBlue));*/

            Profiler.EndSample();
        }

        currentElement += steps;
        updateElement += steps;
        if (currentElement >= totalElements)
        {
            currentElement = 0;
            updateElement = 0;
            //generator.LightMap.Apply();
        }

        if (updateElement > rtUpdate)
        {
            updateElement -= rtUpdate;
            //generator.LightMap.Apply();
        }

        Profiler.EndSample();
    }

    private void OnDrawGizmosSelected()
    {
        if (Application.isPlaying == true && generator != null && generator.IsBaking)
        {
            /*NativeArray<float> positions = generator.PositionTexture.GetRawTextureData<float>();
            NativeArray<float> normals = generator.NormalsTexture.GetRawTextureData<float>();

            Vector3 pos = new Vector3(positions[currentElement * 4 + 0], positions[currentElement * 4 + 1], positions[currentElement * 4 + 2]);
            Vector3 dir = new Vector3(normals[currentElement * 4 + 0], normals[currentElement * 4 + 1], normals[currentElement * 4 + 2]);

            RayTracer.Ray r = new RayTracer.Ray(pos + dir * 0.01f, new Vector3(0.0f, 1.0f, 0.0f));

            float closestDistance = float.MaxValue;
            float closestU = 0.0f;
            float closestV = 0.0f;
            bool closestHit = false;

            Color prev = Gizmos.color;

            for (int j = 0; j < triangles.Count; j++)
            {
                RayTracer.IntersectResult result = new RayTracer.IntersectResult();
                bool hit = RayTracer.Intersect.RayTriangle(r, triangles[j], ref result);

                if (hit == true)
                {
                    if (result.distance > 0.0f && result.distance < closestDistance)
                    {
                        closestDistance = result.distance;
                        closestU = result.baryU;
                        closestV = result.baryV;
                        closestHit = true;
                    }
                }

                Gizmos.color = Color.white;
                Gizmos.DrawLine(triangles[j].a, triangles[j].b);
                Gizmos.DrawLine(triangles[j].b, triangles[j].c);
                Gizmos.DrawLine(triangles[j].c, triangles[j].a);
            }

            Gizmos.color = Color.green;
            Gizmos.DrawLine(r.origin + new Vector3(1.0f, 0.0f, 0.0f), r.origin - new Vector3(1.0f, 0.0f, 0.0f));
            Gizmos.DrawLine(r.origin + new Vector3(0.0f, 1.0f, 0.0f), r.origin - new Vector3(0.0f, 1.0f, 0.0f));
            Gizmos.DrawLine(r.origin + new Vector3(0.0f, 0.0f, 1.0f), r.origin - new Vector3(0.0f, 0.0f, 1.0f));

            if (closestHit)
            {
                Gizmos.color = Color.red;
                Gizmos.DrawLine(r.origin, r.origin + r.direction * closestDistance);
            }
            else
            {
                Gizmos.color = Color.blue;
                Gizmos.DrawLine(r.origin, r.origin + r.direction * 100.0f);
            }*/

            NativeArray<float> positions = generator.PositionTexture.GetRawTextureData<float>();
            NativeArray<float> normals = generator.NormalsTexture.GetRawTextureData<float>();

            Vector3 pos = new Vector3(positions[currentElement * 4 + 0], positions[currentElement * 4 + 1], positions[currentElement * 4 + 2]);
            Vector3 dir = new Vector3(normals[currentElement * 4 + 0], normals[currentElement * 4 + 1], normals[currentElement * 4 + 2]);

            RayTracer.Ray r = new RayTracer.Ray(pos + dir * 0.01f, new Vector3(0.0f, 1.0f, 0.0f));

            RayTracer.IntersectResult result = new RayTracer.IntersectResult();

            int index = accel.IntersectRay(r, ref result);

            Color prev = Gizmos.color;

            Gizmos.color = Color.green;
            Gizmos.DrawLine(r.origin + new Vector3(1.0f, 0.0f, 0.0f), r.origin - new Vector3(1.0f, 0.0f, 0.0f));
            Gizmos.DrawLine(r.origin + new Vector3(0.0f, 1.0f, 0.0f), r.origin - new Vector3(0.0f, 1.0f, 0.0f));
            Gizmos.DrawLine(r.origin + new Vector3(0.0f, 0.0f, 1.0f), r.origin - new Vector3(0.0f, 0.0f, 1.0f));

            if (index >= 0)
            {
                Gizmos.color = Color.red;
                Gizmos.DrawLine(r.origin, r.origin + r.direction * result.distance);
            }
            else
            {
                Gizmos.color = Color.blue;
                Gizmos.DrawLine(r.origin, r.origin + r.direction * 100.0f);
            }

            RayTracer.BVH.BuildNode node = accel.GetHelperBuildNode(bvhDebug);
            if (node != null)
            {
                if (node.primitiveCount == 0)
                {
                    Gizmos.color = Color.yellow;
                }
                else
                {
                    for (int i = 0; i < node.primitiveCount; i++)
                    {
                        int baseindex = node.primitiveOffset + i;
                        int triIndex = accel.GetIndices()[baseindex];

                        Gizmos.color = Color.white;
                        Gizmos.DrawLine(triangles[triIndex].a, triangles[triIndex].b);
                        Gizmos.DrawLine(triangles[triIndex].b, triangles[triIndex].c);
                        Gizmos.DrawLine(triangles[triIndex].c, triangles[triIndex].a);
                    }

                    Gizmos.color = Color.red;
                }

                Gizmos.DrawLine(new Vector3(node.bounds.min.x, node.bounds.min.y, node.bounds.min.z),
                    new Vector3(node.bounds.max.x, node.bounds.min.y, node.bounds.min.z));

                Gizmos.DrawLine(new Vector3(node.bounds.max.x, node.bounds.min.y, node.bounds.min.z),
                    new Vector3(node.bounds.max.x, node.bounds.max.y, node.bounds.min.z));

                Gizmos.DrawLine(new Vector3(node.bounds.max.x, node.bounds.max.y, node.bounds.min.z),
                    new Vector3(node.bounds.min.x, node.bounds.max.y, node.bounds.min.z));

                Gizmos.DrawLine(new Vector3(node.bounds.min.x, node.bounds.max.y, node.bounds.min.z),
                    new Vector3(node.bounds.min.x, node.bounds.min.y, node.bounds.min.z));

                Gizmos.DrawLine(new Vector3(node.bounds.min.x, node.bounds.min.y, node.bounds.max.z),
                    new Vector3(node.bounds.max.x, node.bounds.min.y, node.bounds.max.z));

                Gizmos.DrawLine(new Vector3(node.bounds.max.x, node.bounds.min.y, node.bounds.max.z),
                    new Vector3(node.bounds.max.x, node.bounds.max.y, node.bounds.max.z));

                Gizmos.DrawLine(new Vector3(node.bounds.max.x, node.bounds.max.y, node.bounds.max.z),
                    new Vector3(node.bounds.min.x, node.bounds.max.y, node.bounds.max.z));

                Gizmos.DrawLine(new Vector3(node.bounds.min.x, node.bounds.max.y, node.bounds.max.z),
                    new Vector3(node.bounds.min.x, node.bounds.min.y, node.bounds.max.z));

                Gizmos.DrawLine(new Vector3(node.bounds.min.x, node.bounds.max.y, node.bounds.min.z),
                    new Vector3(node.bounds.min.x, node.bounds.min.y, node.bounds.min.z));

                Gizmos.DrawLine(new Vector3(node.bounds.min.x, node.bounds.min.y, node.bounds.min.z),
                    new Vector3(node.bounds.min.x, node.bounds.min.y, node.bounds.max.z));

                Gizmos.DrawLine(new Vector3(node.bounds.max.x, node.bounds.min.y, node.bounds.min.z),
                    new Vector3(node.bounds.max.x, node.bounds.min.y, node.bounds.max.z));

                Gizmos.DrawLine(new Vector3(node.bounds.max.x, node.bounds.max.y, node.bounds.min.z),
                    new Vector3(node.bounds.max.x, node.bounds.max.y, node.bounds.max.z));

                Gizmos.DrawLine(new Vector3(node.bounds.min.x, node.bounds.max.y, node.bounds.min.z),
                    new Vector3(node.bounds.min.x, node.bounds.max.y, node.bounds.max.z));
            }

            Gizmos.color = prev;
        }
    }
}