Verlet gpu

using UnityEngine;
using System.Collections;
using System.Diagnostics;
using TMPro;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using Debug = UnityEngine.Debug;
using Random = UnityEngine.Random;


public class VerletIntegration : MonoBehaviour
{
    public int numParticles;
    public float radius;
    public float gravity;
    public float constraintRadius;
    public int collisionSubSteps;
    public Material objectMaterial;
    public ComputeShader verletShader;

    public int maxCount;
    public float cellSize;

    Particle[] particles;
    int[] grid;
    ComputeBuffer particlesBuffer;
    ComputeBuffer gridBuffer;

    new ParticleRenderer renderer;
    TextMeshProUGUI latencyText;
    TextMeshProUGUI numParticlesText;

    int currentNumParticles;
    float radius2;
    int gridSize;

    void Start()
    {
        Application.targetFrameRate = 50;

        particles = new Particle[numParticles];

        float cellSize = radius * (2f + this.cellSize);
        gridSize = Mathf.CeilToInt(10 / cellSize);
        grid = new int[gridSize * gridSize * maxCount]; // The maximum number of circles in one cell is 6
        radius2 = radius * 2;

        latencyText = GameObject.Find("Latency").GetComponent<TextMeshProUGUI>();
        numParticlesText = GameObject.Find("NumParticles").GetComponent<TextMeshProUGUI>();
        GameObject.Find("ConstraintCircle").transform.localScale = new Vector3(constraintRadius * 2 + radius, constraintRadius * 2 + radius, 0);

        renderer = new(numParticles, objectMaterial, radius);

        verletShader.SetInt("gridSize", gridSize);
        verletShader.SetFloat("radius2", radius2);
        particlesBuffer = new ComputeBuffer(numParticles, sizeof(float) * 4);
        gridBuffer = new ComputeBuffer(grid.Length, sizeof(int));

        StartCoroutine(SpawnObjects());
    }

    void FixedUpdate()
    {
        Stopwatch watch = new();
        watch.Start();

        for (int i = 0; i < collisionSubSteps; i++)
        {
            UpdatePositions();
            RebuildGrid();

            particlesBuffer.SetData(particles);
            gridBuffer.SetData(grid);

            verletShader.SetBuffer(0, "particles", particlesBuffer);
            verletShader.SetBuffer(0, "grid", gridBuffer);
            verletShader.Dispatch(0, (gridSize - 2) / 8, (gridSize - 2) / 8, 1);

            particlesBuffer.GetData(particles);
            gridBuffer.GetData(grid);
        }

        watch.Stop();
        latencyText.text = watch.ElapsedMilliseconds + " ms";

        if (watch.ElapsedMilliseconds > 16)
        {
            Debug.Log(numParticlesText);
            return;
        }
    }

    void Update()
    {
        renderer.Render(currentNumParticles, particles);
        if (!Application.isPlaying)
        {
            renderer.ReleaseBuffers();
        }
    }

    void UpdatePositions()
    {
        float dt = Time.fixedDeltaTime / collisionSubSteps;

        for (int i = 0; i < currentNumParticles; i++)
        {
            float2 position = particles[i].position;

            //Perform verlet integration
            float2 positionBeforeUpdate = position;
            position += position - particles[i].prevPosition;
            position += gravity * dt * dt * new float2(0, -1);

            //Constraint object inside a circle
            float2 constraintPosition = Vector2.zero;
            float2 constraintDiff = position - constraintPosition;
            float len = math.length(constraintDiff);

            if (len > constraintRadius - radius)
            {
                float2 n = constraintDiff / len;
                position = constraintPosition + n * (constraintRadius - radius);
            }

            particles[i] = new Particle
            {
                position = position,
                prevPosition = positionBeforeUpdate
            };
        }
    }

    void SolveCollisions()
    {
        for (int y = 1; y < gridSize - 1; y++)
        {
            for (int x = 1; x < gridSize - 1; x++)
            {
                int index = (x + y * gridSize) * 3;

                for (int yOffset = -1; yOffset <= 1; yOffset++)
                {
                    for (int xOffset = -1; xOffset <= 1; xOffset++)
                    {
                        int otherIndex = (x + xOffset + (y + yOffset) * gridSize) * 3;
                        SolveCellCollisions(index, otherIndex);
                    }
                }
            }
        }
    }

    void SolveCellCollisions(int cell1Index, int cell2Index)
    {
        for (int i = 0; i < 3; i++)
        {
            int gridIndex1 = cell1Index + i;
            int particleIndex1 = grid[gridIndex1];

            if(particleIndex1 == -1) continue;

            for (int j = 0; j < 3; j++)
            {
                int gridIndex2 = cell2Index + j;
                int particleIndex2 = grid[gridIndex2];

                if (particleIndex2 == -1) continue;
                if(particleIndex1 == particleIndex2) continue;

                var object1 = particles[particleIndex1];
                var object2 = particles[particleIndex2];

                Vector2 collisionAxis = object1.position - object2.position;
                float distSqr = collisionAxis.x * collisionAxis.x + collisionAxis.y * collisionAxis.y;
                float minDist = radius2;

                if(distSqr < minDist * minDist)
                {
                    float dist = Mathf.Sqrt(distSqr);
                    float2 n = collisionAxis / dist; // Normalize collision axis
                    float delta = minDist - dist;

                    // Apply corrections to both objects
                    object1.position += 0.5f * delta * n;
                    object2.position -= 0.5f * delta * n;

                    particles[particleIndex1] = object1;
                    particles[particleIndex2] = object2;
                }
            }
        }
    }

    void RebuildGrid()
    {
        for (int y = 0; y < gridSize; y++)
        {
            for (int x = 0; x < gridSize; x++)
            {
                int baseIndex = (x + y * gridSize) * 3;
                // grid[baseIndex] = -1;
                // grid[baseIndex + 1] = -1;
                // grid[baseIndex + 2] = -1;

                for (int i = 0; i < maxCount; i++)
                {
                    grid[baseIndex + i] = -1;
                }
            }
        }

        for (int i = 0; i < currentNumParticles; i++)
        {
            float xPos01 = (particles[i].position.x + 5 - radius) / 10f;
            float yPos01 = (particles[i].position.y + 5 - radius) / 10f;
            if (xPos01 < 0 || xPos01 > 1 || yPos01 < 0 || yPos01 > 1)
            {
                Debug.LogWarning("X: " + xPos01 + ", Y: " + yPos01);
            }

            int x = Mathf.FloorToInt(xPos01 * gridSize);
            int y =  Mathf.FloorToInt(yPos01 * gridSize);

            int baseIndex = (x + y * gridSize) * 3;
            // if(grid[baseIndex] == -1) grid[baseIndex] = i;
            // else if(grid[baseIndex + 1] == -1) grid[baseIndex + 1] = i;
            // else if (grid[baseIndex + 2] == -1) grid[baseIndex + 2] = i;
            // else if (grid[baseIndex + 3] == -1) grid[baseIndex + 3] = i;
            // else if (grid[baseIndex + 4] == -1) grid[baseIndex + 4] = i;
            // else if (grid[baseIndex + 5] == -1) grid[baseIndex + 5] = i;

            bool b = false;
            for (int j = 0; j < maxCount; j++)
            {
                if (grid[baseIndex + j] == -1)
                {
                    grid[baseIndex + j] = i;
                    b = true;
                    break;
                }

            }
            if(!b) Debug.LogWarning("Overflow");
        }
    }


    IEnumerator SpawnObjects()
    {
        while (true)
        {
            // Vector2 mousePositionWorld = Camera.main.ScreenToWorldPoint(new Vector3(Input.mousePosition.x, Input.mousePosition.y, 0));
            Vector2 pos = new(Random.Range(-1f, 1f), 3);
            Particle particle = new()
            {
                position = pos,
                prevPosition = pos,
            };

            particles[currentNumParticles] = particle;
            currentNumParticles++;
            numParticlesText.text = currentNumParticles + " particles";

            if (currentNumParticles == numParticles) break;

            yield return new WaitForSeconds(0.07f);
        }
    }
}

public struct Particle
{
    public float2 position, prevPosition;
}