#include "Node.h"
#include <vector>
#include <iostream>
#include <math.h>
using namespace std;

/*
Node constructor. Used recursively to build the Barnes-Hut tree. px, py, pz 
denote the corner (lowest value for each dimension) of the box of size pSize.
pBodies is a vector containing all the nodes that must be placed in this box.
*/
Node::Node(const vector<Node*> &pBodies, const double pSize, 
    const double px, const double py, const double pz){
    size = pSize; // Required later for treeWalk

    // Divide into subnodes (octants)
    vector<Node*> subBodies[2][2][2];

    for (int i = 0; i < pBodies.size(); i++){
        int xIndex, yIndex, zIndex;

        if (pBodies[i]->COM[0] < (px + (size / 2))) xIndex = 0;
        else xIndex = 1;

        if (pBodies[i]->COM[1] < (py + (size / 2))) yIndex = 0;
        else yIndex = 1;

        if (pBodies[i]->COM[2] < (pz + (size / 2))) zIndex = 0;
        else zIndex = 1;

        subBodies[xIndex][yIndex][zIndex].push_back(pBodies[i]);
    }


    // Recursively place the nodes
    // g++ will unroll these loops
    for (int i = 0; i < 2; i++){
        for (int j = 0; j < 2; j++){
            for (int k = 0; k < 2; k++){
                switch(subBodies[i][j][k].size()){
                    case 0: continue;
                    case 1:
                        subBodies[i][j][k][0]->size = size/2;
                        children.push_back(subBodies[i][j][k][0]);
                        break;
                    default:
                        children.push_back(new Node(subBodies[i][j][k], size/2,
                            px + size/2*i, py + size/2*j, pz + size/2*k));
                }
            }
        }
    }

    // Recursively calculate the COM
    memset(COM, 0, sizeof(COM)); // Set COM to 0s
    m = 0.; // mass
    for (int i = 0; i < children.size(); i++){
        m += children[i]->m;
        for (int j = 0; j < 3; j++) 
            COM[j] += children[i]->m * children[i]->COM[j];
    }
    // COM only relevant if there is mass in the octant
    if (m > 0) for (int i = 0; i < 3; i++) COM[i] /= m;

}

/*
Node constructor. Used to build the leaf nodes directly from the data passed
from Python using ctypes.
*/
Node::Node(const double* pr_vec, const double pm){
    // Initialize a node (leaf)
    for (int i = 0; i < 3; i++) COM[i] = pr_vec[i];
    memset(g, 0, sizeof(g)); // Set g to 0s
    m = pm; // mass
}

/*
Calculate the acceleration at the this node. Used recursively calling
treeWalk(topNode, thetamax). This is O(log n) and will be called for
each node in the tree: O(n log n). Soft defines the characteristic
plummer softening scale.
*/
void Node::treeWalk(const Node &node, 
    const double thetamax, const double soft){
    // Calculate distance to node
    double delta[3];
    for (int i = 0; i < 3; i++) delta[i] = node.COM[i] - COM[i];
    double r = sqrt(delta[0]*delta[0] + delta[1]*delta[1] + delta[2]*delta[2]);

    if (r==0) return; // Do not interact with self

    // If it satisfies the size/r < thetamax criterion, add the g contribution
    if (node.children.size() == 0 || node.size/r < thetamax){
        double tripler = (r+soft) * (r+soft) * r;
        for(int i = 0; i < 3; i++) g[i] += node.m * delta[i] / tripler;
    }
    else{ // Otherwise recurse into its children
        for (int i = 0; i < node.children.size(); i++){
            treeWalk((*node.children[i]), thetamax, soft);
        }
    }
}