Iterative quadtree

using System.Collections.Generic;
using UnityEngine;
using Particle = Simulation.Particle;

struct Node
{
    public Quad boundary;
    public float mass;
    public Vector2 centerOfMass;
    public int childIndex; //NW index
}
public struct Quad
{
    //Bottom left of rectangle
    public Vector2 position;
    public float size;

    public bool Intersects(Vector2 pos)
    {
        return pos.x > position.x && pos.x < position.x + size &&
               pos.y > position.y && pos.y < position.y + size;
    }
    public int FindNode(Vector2 pos)
    {
        bool wX = pos.x - position.x < size / 2;
        bool sY = pos.y - position.y < size / 2;

        if (wX && !sY) return 0;
        if (!wX && !sY) return 1;
        if (wX && sY) return 2;
        if(!wX && sY) return 3;

        return 0;
    }
}

public class QuadTree
{
    List<Node> quadTree;

    public QuadTree()
    {
        quadTree = new();
        Quad boundary = new() { position = new Vector2(-5, -5), size = 10 };
        quadTree.Add(new Node
        {
            boundary = boundary,
        });
    }

    public void DrawTree()
    {
        foreach (var node in quadTree)
        {
            DrawSquare(node.boundary.position, node.boundary.size);
        }
    }
    static void DrawSquare(Vector2 position, float size)
    {
        Debug.DrawLine(position, new Vector2(position.x, position.y + size), Color.red);
        Debug.DrawLine(new Vector2(position.x, position.y + size), new Vector2(position.x + size, position.y + size), Color.red);
        Debug.DrawLine(new Vector2(position.x + size, position.y + size), new Vector2(position.x + size, position.y), Color.red);
        Debug.DrawLine(new Vector2(position.x + size, position.y), new Vector2(position.x, position.y), Color.red);
    }

    public Vector2 CalculateAcceleration(Particle particle, float theta, float epsilon, float timeStep)
    {
        Vector2 acceleration = new();

        float thetaSqr = theta * theta;
        float epsilonSqr = epsilon * epsilon;

        Stack<int> stack = new();
        stack.Push(0);
        while (stack.Count > 0)
        {
            Node node = quadTree[stack.Pop()];
            Vector2 diff = node.centerOfMass - particle.position;
            float sqrDst = Mathf.Max(diff.sqrMagnitude, 0.05f);

            if (node.childIndex == 0 || node.boundary.size * node.boundary.size < sqrDst * thetaSqr)
            {
                float denom = (sqrDst + epsilonSqr) * Mathf.Sqrt(sqrDst);
                acceleration += diff * node.mass / denom * timeStep;
                continue;
            }

            for (int i = 3; i >= 0; i--)
            {
                stack.Push(node.childIndex + i);
            }
        }

        return acceleration;
    }

    public void Propagate()
    {
        for (int i = quadTree.Count - 1; i > 0; i--)
        {
            if(quadTree[i].childIndex == 0) continue;

            int child = quadTree[i].childIndex;
            Node node = quadTree[i];

            node.centerOfMass = quadTree[child].centerOfMass * quadTree[child].mass
                                + quadTree[child + 1].centerOfMass * quadTree[child + 1].mass
                                + quadTree[child + 2].centerOfMass * quadTree[child + 2].mass
                                + quadTree[child + 3].centerOfMass * quadTree[child + 3].mass;
            node.mass = quadTree[child].mass
                        + quadTree[child + 1].mass
                        + quadTree[child + 2].mass
                        + quadTree[child + 3].mass;
            node.centerOfMass /= node.mass;
            quadTree[i] = node;

        }
    }

    public void Insert(Particle particle)
    {
        int maxIterations = 100;
        int iteration = 0;
        while (true && iteration < maxIterations)
        {
            iteration++;
            int bestNode = 0;
            for (int i = 0; i < quadTree.Count; i++)
            {
                Node node = quadTree[i];

                //Continue if node isn't leaf node
                if(node.childIndex != 0) continue;
                //Continue if particle does not belong in this node
                if (!node.boundary.Intersects(particle.position)) continue;

                bestNode = i;
                //If node is empty, set the mass and position
                if (node.mass == 0)
                {
                    node.mass = particle.mass;
                    node.centerOfMass = particle.position;
                    quadTree[i] = node;
                    return;
                }
                //If the node's position is the same as the particle's, add the particle's mass
                if (node.centerOfMass == particle.position)
                {
                    node.mass += particle.mass;
                    quadTree[i] = node;
                    return;
                }
            }

            Vector2 nodePos = quadTree[bestNode].centerOfMass;
            float nodeMass = quadTree[bestNode].mass;

            int iterations = 0;
            while (true && iterations < 100)
            {
                iterations++;
                //Check what node particle and bestNode ends up in
                //If they are the same, continue and set bestNode to FindNode(particle)
                //Otherwise, set the mass and pos of FindNode(particle)
                Subdivide(bestNode);
                int childIndex = quadTree[bestNode].childIndex;
                int q1 = quadTree[bestNode].boundary.FindNode(nodePos);
                int q2 = quadTree[bestNode].boundary.FindNode(particle.position);

                if (q1 == q2)
                {
                    bestNode = q1 + childIndex;
                }
                else
                {
                    int n1 = q1 + childIndex;
                    int n2 = q2 + childIndex;

                    quadTree[n1] = new Node
                    {
                        centerOfMass = nodePos,
                        mass = nodeMass,
                        boundary = quadTree[n1].boundary,
                    };
                    quadTree[n2] = new Node
                    {
                        centerOfMass = particle.position,
                        mass = particle.mass,
                        boundary = quadTree[n2].boundary,
                    };
                    return;
                }
            }
        }
    }

    void Subdivide(int parentIndex)
    {
        Node parent = quadTree[parentIndex];
        parent.childIndex = quadTree.Count;
        quadTree[parentIndex] = parent;
        Vector2 parentPos = parent.boundary.position;
        Vector2 nwPosition = new(parentPos.x, parentPos.y + parent.boundary.size / 2);
        Vector2 nePosition = new(parentPos.x + parent.boundary.size / 2, parentPos.y + parent.boundary.size / 2);
        Vector2 swPosition = new(parentPos.x, parentPos.y);
        Vector2 sePosition = new(parentPos.x + parent.boundary.size / 2, parentPos.y);

        Node nw = new()
        {
            boundary = new Quad { position = nwPosition, size = parent.boundary.size / 2 }
        };
        Node ne = new()
        {
            boundary = new Quad { position = nePosition, size = parent.boundary.size / 2 }
        };
        Node sw = new()
        {
            boundary = new Quad { position = swPosition, size = parent.boundary.size / 2 }
        };
        Node se = new()
        {
            boundary = new Quad { position = sePosition, size = parent.boundary.size / 2 }
        };

        quadTree.Add(nw);
        quadTree.Add(ne);
        quadTree.Add(sw);
        quadTree.Add(se);
    }
}