#include <vector>
#include <iostream>
#include <math.h> 
#include <chrono>

using namespace std;
using namespace std::chrono;
#include "Node.h"


/*
Calculates the self gravitational acceleration for a set of particles located
at r_vec (n, 3) with masses m_vec (n,) using a Barnes-Hut tree with theta = 0.7

Parameters:
    r_vec: the position of the particles.
    m_vec: the masses of the particles.
    n: the number of particles. This cannot be infered by C++ and must be
        passed directly.
    size: size of the top node in the octtree.
    px, py, pz: coordinates of the lowest corner of the top node.
    size: characteristic softening scale.

Returns:
    The accelerations computed for each mass (n, 3).
*/
extern "C" double* barnesHutCPP(double** r_vec, double* m_vec, int n, 
    double size, double px, double py, double pz, double soft){
    // Create nodes
    std::vector<Node*> nodes;
    for (int i = 0; i < n; i++){
        nodes.push_back(new Node(r_vec[i], m_vec[i]));
    }

    // Create the tree
    Node* tree = new Node(nodes, size, px, py, pz);

    // Calculate the accelerations for each node. We want to return the
    // result as an array and use a 1D array for simplicity since this will
    // be allocated continously in the heap and can be reshaped in Numpy.
    double* accel = new double[3*n];
    for (int i = 0; i < nodes.size(); i++){
        nodes[i]->treeWalk(*tree, 0.7, soft); // thetamax = 0.7
        accel[3*i+0] = nodes[i]->g[0];
        accel[3*i+1] = nodes[i]->g[1];
        accel[3*i+2] = nodes[i]->g[2];
    }
    
    // Return as an (n,) array
    return accel;
}

/*
Calculates the self gravitational acceleration for a set of particles located
at r_vec (n, 3) with masses m_vec (n,) using Brute Force pairwise summation.
Massive particles must appear first.

Parameters:
    r_vec: the position of the particles.
    m_vec: the masses of the particles.
    n: the number of particles. This cannot be infered by C++ and must be
        passed directly.
    size: characteristic softening scale.

Returns:
    The accelerations computed for each mass (n, 3).
*/
extern "C" double* bruteForceCPP(double** r_vec, double* m_vec, 
    int n, double soft){

    // Initialize result and fill with 0s
    // Use a 1D array so as not to have to convert back in Numpy
    double* accel = new double[3*n];
    for (int i=0; i<3*n; i++){
        accel[i] = 0;
    }

    // Compute the acceleration
    for (int i=0; i<n; i++){
        // Only consider pairs with at least one massive particle i
        if (m_vec[i] == 0.) break;
        for (int j=i+1; j<n; j++){
            // Distance between particles
            double delta[3];
            for (int k = 0; k < 3; k++) delta[k] = r_vec[j][k] - r_vec[i][k];
            double r = sqrt(delta[0]*delta[0] 
                + delta[1]*delta[1] 
                + delta[2]*delta[2]);
            double tripler = (r+soft) * (r+soft) * r;

            // Compute acceleration on first particle
            double mr3inv = m_vec[i]/tripler;
            for (int k = 0; k < 3; k++) accel[3*j+k] -= mr3inv * delta[k];

            // Compute acceleration on second particle
            // For pairs with one massless particle, no reaction force
            if (m_vec[j] == 0.) break;
            // Otherwise, opposite direction (+)
            mr3inv = m_vec[j]/tripler;
            for (int k = 0; k < 3; k++) accel[3*i+k] += mr3inv * delta[k];
        }
    }

    // Return as an (n,) array
    return accel;
}