using System.Collections.Generic;
using UnityEngine;
using Vector2 = UnityEngine.Vector2;
using Particle = Simulation.Particle;
public class QuadTree
{
    // Arbitrary constant to indicate how many elements can be stored in this quad
    const int Capacity = 1;
    const float MinSize = .1f;

    Vector2 centerOfMass;
    float mass;

    // Axis-aligned bounding box stored as a center with half-dimensions
    // to represent the boundaries of this quad tree
    Rectangle boundary;

    //Objects in this quad tree node
    readonly List<Particle> particles;

    // Children
    QuadTree northWest;
    QuadTree northEast;
    QuadTree southWest;
    QuadTree southEast;

    //True if the quad subdivides
    bool divided;

    readonly float particleRadius;

    public QuadTree(Rectangle boundary, float particleRadius)
    {
        this.boundary = boundary;
        this.particleRadius = particleRadius;
        particles = new();
        divided = false;
    }

    public void DrawTree()
    {
        DrawRectangle(boundary.position, boundary.size);
        Debug.DrawLine(centerOfMass, centerOfMass + new Vector2(0.1f, 0), Color.blue);

        if (divided)
        {
            northWest.DrawTree();
            northEast.DrawTree();
            southWest.DrawTree();
            southEast.DrawTree();
        }
    }
    static void DrawRectangle(Vector2 position, Vector2 size)
    {
        Debug.DrawLine(position, new Vector2(position.x, position.y + size.y), Color.red);
        Debug.DrawLine(new Vector2(position.x, position.y + size.y), new Vector2(position.x + size.x, position.y + size.y), Color.red);
        Debug.DrawLine(new Vector2(position.x + size.x, position.y + size.y), new Vector2(position.x + size.x, position.y), Color.red);
        Debug.DrawLine(new Vector2(position.x + size.x, position.y), new Vector2(position.x, position.y), Color.red);
    }

    public Particle UpdateVelocity(Particle particle, float timeStep)
    {
        //If the node hasn't been divided, calculate the force exerted by all particles in this node
        if (!divided)
        {
            for (int i = 0; i < particles.Count; i++)
            {
                if(particles[i].position == particle.position) continue;
                particle.velocity += CalculateAcceleration(particle, particles[i]) * timeStep;
            }

            return particle;
        }
        //Otherwise, calculate the ratio s/d. If node is sufficiently far away, treat it as a single particle
        float sd = boundary.size.x / (particle.position - centerOfMass).magnitude;
        if (sd < 0.5f)
        {
            Particle particle2 = new(){ position = centerOfMass, mass = mass };
            particle.velocity += CalculateAcceleration(particle, particle2) * timeStep;
            return particle;
        }
        //Otherwise, check children if the node has been divided
        particle = northWest.UpdateVelocity(particle, timeStep);
        particle = northEast.UpdateVelocity(particle, timeStep);
        particle = southWest.UpdateVelocity(particle, timeStep);
        particle = southEast.UpdateVelocity(particle, timeStep);
        return particle;
    }

    Vector2 CalculateAcceleration(Particle p1, Particle p2)
    {
        Vector2 diff = p2.position - p1.position;
        Vector2 forceDir = diff.normalized;
        float sqrDst = Mathf.Max(diff.sqrMagnitude, 0.05f);
        Vector2 force = forceDir * (p1.mass * p2.mass) / sqrDst;
        Vector2 acceleration = force / p1.mass;

        return acceleration;
    }

    public void CalculateMassAndCenterOfMass()
    {
        centerOfMass = new Vector2();

        for (int i = 0; i < particles.Count; i++)
        {
            mass += particles[i].mass;
        }

        float x = 0;
        float y = 0;

        for (int i = 0; i < particles.Count; i++)
        {
            x += particles[i].position.x * particles[i].mass;
            y += particles[i].position.y * particles[i].mass;
        }

        x /= mass;
        y /= mass;
        centerOfMass = new Vector2(x, y);

        if (divided)
        {
            northWest.CalculateMassAndCenterOfMass();
            northEast.CalculateMassAndCenterOfMass();
            southWest.CalculateMassAndCenterOfMass();
            southEast.CalculateMassAndCenterOfMass();
        }
    }

    public void Insert(Particle particle)
    {
        //If the quad has already divided, insert it into on of the children
        if (divided)
        {
            InsertIntoChildren(particle);
        }
        else
        {
            // Ignore objects that do not belong in this quad tree
            if (!boundary.Intersects(particle)) return;

            // If there is space in this quad tree and if doesn't have subdivisions, add the object here
            if (particles.Count < Capacity)
            {
                particles.Add(particle);
            }
            // Otherwise, subdivide and then add the point to whichever node will accept it
            else
            {
                Subdivide(particle);
            }

        }
    }

    void Subdivide(Simulation.Particle particle)
    {
        float size = Mathf.Min(boundary.size.x, boundary.size.y) / 2;

        if (size > MinSize)
        {
            northWest = new(new Rectangle
            {
                position = new Vector2(boundary.position.x, boundary.position.y + boundary.size.y / 2f),
                size = boundary.size / 2,
                pRadius = particleRadius
            }, particleRadius);
            northEast = new(new Rectangle
            {
                position = new Vector2(boundary.position.x + boundary.size.x / 2, boundary.position.y + boundary.size.y / 2f),
                size = boundary.size / 2,
                pRadius = particleRadius
            }, particleRadius);
            southWest = new(new Rectangle
            {
                position = new Vector2(boundary.position.x, boundary.position.y),
                size = boundary.size / 2,
                pRadius = particleRadius
            }, particleRadius);
            southEast = new(new Rectangle
            {
                position = new Vector2(boundary.position.x + boundary.size.x / 2, boundary.position.y),
                size = boundary.size / 2,
                pRadius = particleRadius
            }, particleRadius);

            //Move all objects from this quad into children
            InsertIntoChildren(particle);
            for (int i = 0; i < particles.Count; i++)
            {
                InsertIntoChildren(particles[i]);
            }

            divided = true;
        }
    }

    void InsertIntoChildren(Simulation.Particle particle)
    {
        northWest.Insert(particle);
        northEast.Insert(particle);
        southWest.Insert(particle);
        southEast.Insert(particle);
    }
}

public struct Rectangle
{
    //Bottom left of rectangle
    public Vector2 position;
    public Vector2 size;
    public float pRadius;

    public bool Intersects(Simulation.Particle p)
    {
        return p.position.x > position.x - pRadius && p.position.x < position.x + size.x + pRadius &&
               p.position.y > position.y - pRadius && p.position.y < position.y + size.y + pRadius;
    }
}