Node.cpp

//
// Created by lucy on 04/03/25.
//

#include "Node.h"

#include <cassert>
#include <iostream>
#include <numeric>
#define G 0.0043009172706 // Gravitational constant in pc/M_sun/(km/s)^2
#define SOLAR_MASS 1.989e30 // Solar mass in kg


void Node::makeChildren() {
    // Make new children nodes
    // particles: vector of particles in the node
    this->calculateCOM();
    if (this->particles.size() <= 1) {
        if (this->particles.size() == 1) {
            this->particles[0]->node = this;
        }
        return;
    }
    // Pos refers to the centre of the node
    long double x = this->posX;
    long double y = this->posY;
    long double z = this->posZ;
    long double quarterSize = this->size / 4.0;

    // Create 8 children nodes (octants)
    for (int i = 0 ; i < 8; i++) {
        children[i] = new Node((i % 2 == 0) ? x - quarterSize : x + quarterSize,
                                (i % 4 > 1) ? y - quarterSize : y + quarterSize,
                                (i % 8 > 3) ? z - quarterSize : z + quarterSize,
                                this->halfSize);
    }
    for (Particle* p : this->particles) {
        for (auto & i : children) {
            if (i->containsParticle(*p)) {
                i->addParticle(p);
                break;
            }
        }
    }
    // For each child, if it has more than 0 particles, make its children
    for (auto & i : children) {
        if (i->particles.size() > 0) {
            i->makeChildren();
        }
    }

}

void Node::addParticle(Particle* particle) {
    // Add a particle to the node
    // particle: particle to add
    this->particles.push_back(particle);
}

bool Node::containsParticle(const Particle& particle) const {
    // Check if the node contains a particle
    // particle: particle to check
    return (this->posX + this->halfSize >= particle.posX && this->posX - this->halfSize <= particle.posX &&
            this->posY + this->halfSize >= particle.posY && this->posY - this->halfSize <= particle.posY &&
            this->posZ + this->halfSize >= particle.posZ && this->posZ - this->halfSize <= particle.posZ);
}

void Node::calculateCOM() {
    // Calculate the center of mass of the node
    // particles: vector of particles in the node
    this->totalMass = 0.0;
    this->comX = 0.0;
    this->comY = 0.0;
    this->comZ = 0.0;
    if (this->particles.empty()) {
        return;
    }
    // Calculate the total mass and center of mass
    this->totalMass = std::accumulate(this->particles.begin(), this->particles.end(), 0.0,
        [](double sum, const Particle* p) { return sum + p->mass; });

    this->comX = std::accumulate(this->particles.begin(), this->particles.end(), 0.0,
        [](double sum, const Particle* p) { return sum + (p->posX * p->mass); }) / this->totalMass;
    this->comY = std::accumulate(this->particles.begin(), this->particles.end(), 0.0,
        [](double sum, const Particle* p) { return sum + (p->posY * p->mass); }) / this->totalMass;
    this->comZ = std::accumulate(this->particles.begin(), this->particles.end(), 0.0,
        [](double sum, const Particle* p) { return sum + (p->posZ * p->mass); }) / this->totalMass;
    if (this->comX == 0 | this->comY == 0 | this->comZ == 0) {
        std::cout << "COM is 0" << std::endl;
    }
    if (this->totalMass == 0) {
        std::cout << "Total mass is 0" << std::endl;
    }
}

void Node::printNode() const {
    // Print the node
    std::cout << "\nNODE INFO" << std::endl;
    std::cout << "Node position: (" << this->posX << ", " << this->posY << ", " << this->posZ << ")" << std::endl;
    std::cout << "Node size: " << this->size << std::endl;
    std::cout << "Number of particles: " << this->particles.size() << std::endl;
    std::cout << "Center of mass: (" << this->comX << ", " << this->comY << ", " << this->comZ << ")" << std::endl;
    std::cout << "Total mass: " << this->totalMass << std::endl;
    for (auto & i : children) {
        if (i != nullptr) {
            i->printNode();
        }
    }
}

void Node::updateParticlesRoot(double deltaTime, float pheta) {
    // Apply forces for all particles in the root node
    for (Particle* particle : this->particles) {
        if (particle->node == this) {
            continue;
        }
        /* If the set pheta is smaller than the size to distance ratio then we search
         * one of the children nodes the particle is in*/
        if (pheta < this->getSDRatio(particle)) {
            for (auto & i : children) {
                if (i != nullptr) {
                    i->updateParticleNode(pheta, particle);
                }
            }
        }
        /* If the set pheta is larger than the size to distance ratio then we calculate
         * the acceleration on that particle based on the mass of the node */
        else {
            double dist = this->getDistToNode(particle);
            //debug
            if (dist == 0) {
                std::cout << "Distance is 0" << std::endl;
                continue;
            }
            std::vector<long double> normalisedVector = this->getNormalisedVector(particle->node, dist);
            auto acceleration = getAcceleration(this->totalMass, dist);
            particle->updateAcceleration(acceleration * normalisedVector[0], acceleration * normalisedVector[1], acceleration * normalisedVector[2]);
        }
    }
    // Update all particle positions and velocity in the node using a leapfrog integrator
    bool needsUpdate = false;
    for (Particle* particle : this->particles) {
        if (particle->node == nullptr) {
            continue;
        }
        particle->update(deltaTime); // The particle update handles the leapfrog integration
        // particle.calculateKineticEnergy();
        // particle.updatePotentialEnergy(0); // Placeholder for potential energy calculation
        // particle.printParticle();

        // Check if the particle is still in the node
        if (!particle->node->containsParticle(*particle)) {
            needsUpdate = true;
            particle->node = nullptr;
        }
        // Check if the particle is still in the simulation area
        if (!this->containsParticle(*particle)) {
            needsUpdate = true;
            std::cout << "Particle out of bounds" << std::endl;
            std::cout << "Erasing" << std::endl;
            particle->node = nullptr;
            std::erase(this->particles, particle);
        }
    }
    // If any particles have moved out of their nodes, we need to update the tree
     if (needsUpdate) {
         std::cout << "Update particles" << std::endl;
         // delete all children
         this->deleteChildren();
         // remake the children
         this->makeChildren();
     }
}

void Node::deleteChildren() {
    // Delete the children of the node
    for (auto & i : children) {
        if (i != nullptr) {
            i->deleteChildren();
            delete i;
            i = nullptr;
        }
    }
}

void Node::updateParticleNode(float pheta, Particle* particle) {
    // Update the particles in the node

    if (particle->node == this) {
        return;
    }
    // Check if the node has any particles
    // If not then we can skip the calculation and return as the node is empty
    if (this->particles.empty()) {
        return;
    }
    // Pheta is smaller than the size to distance ratio then we search the children
    if (pheta < this->getSDRatio(particle)) {
        for (auto & i : children) {
            if (i != nullptr) {
                i->updateParticleNode(pheta, particle);
            }
        }
    }
    /* If pheta is larger than the size to distance ratio then we calculate the
    * acceleration on that particle based on the mass of the node */
    else {
        double dist = this->getDistToNode(particle);
        if (dist == 0) {
            std::cout << "Distance is 0" << std::endl;
            return;
        }
        std::vector<long double> normalisedVector = this->getNormalisedVector(particle->node, dist);
        double acceleration = getAcceleration(this->totalMass, dist);
        particle->updateAcceleration(acceleration * normalisedVector[0], acceleration * normalisedVector[1], acceleration * normalisedVector[2]);
    }
}

long double Node::getDistToNode(Particle* p) {
    // Calculate the distance to another node
    return sqrt(
        pow(this->comX - p->posX, 2) +
        pow(this->comY - p->posY, 2) +
        pow(this->comZ - p->posZ, 2));
}

std::vector<long double> Node::getNormalisedVector(Node *other, double distance) const {
    // Calculate the normalised vector to another node
    if (distance == 0) {
        return {0, 0, 0};
    }
    return {
        (this->comX - other->comX) / distance,
        (this->comY - other->comY) / distance,
        (this->comZ - other->comZ) / distance};
}

double Node::getSDRatio(Particle* p) {
    // Calculate the size to distance ratio
    double dist = this->getDistToNode(p);
    if (dist == 0) {
        return 0;
    }
    return this->size / dist;
}

long double Node::getAcceleration(double mass, long double distance) {
    // Calculate the gravitational acceleration given the mass and distance
    // mass: mass of the object
    // distance: distance from the object
    return (G * mass) / (distance * distance);
}