#include "BarnesHutTree.hpp"
#include "Node.hpp"
#include <iostream>

namespace nbody {
	using namespace std;

	BarnesHutTree::BarnesHutTree() {
		this->control.maxLeafBodies = 16;
		this->control.tree = this;
		this->control.processedNodes = 0;
		pthread_rwlock_init(&this->control.lock, NULL);
	}

	BarnesHutTree::~BarnesHutTree() {
		pthread_rwlock_destroy(&this->control.lock);
	}

	vector<Box> BarnesHutTree::splitBB(Node* node) {
		return node->getBB().octreeSplit();
	}

	int BarnesHutTree::numberOfChildren() {
		return 8;
	}

	void BarnesHutTree::update() {
		//iterate for updating representatives
		Node* current = this->nodes->prev;
		while (current != this->nodes) {
			current->representative.mass = 0.0;
			current->representative.position[0] = 0.0;
			current->representative.position[1] = 0.0;
			current->representative.position[2] = 0.0;
			if (current->leaf) {
				for (unsigned int i = 0; i < current->bodies.size(); i++) {
					current->representative.mass += current->bodies[i].mass;
					for (int j = 0; j < 3; j++) {
						current->representative.position[j] += current->bodies[i].position[j] * current->bodies[i].mass;
					}
				}
			} else {
				Node* child = current->next;

				do {
					current->representative.mass += child->representative.mass;
					for (int j = 0; j < 3; j++) {
						current->representative.position[j] += child->representative.position[j] * child->representative.mass;
					}
					child = child->nextSibling;
				} while (child != NULL);
			}
			for (int j = 0; j < 3; j++) {
				current->representative.position[j] /= current->representative.mass;
			}
			current = current->prev;
		}
	}

	void BarnesHutTree::splitNode(Node* current) {
		if (current->isSplitable(&this->control)) {
			vector<Box> subboxes = this->splitBB(current);
			current->leaf = false;
			Node* after = current->next;

			for (vector<Box>::iterator it = subboxes.begin(); it != subboxes.end(); it++) {
				Node* child = new Node(this);

				child->parent = current;
				child->bb = *it;
				child->bodies = it->copyBodies(current->bodies);
				child->nextSibling = NULL;
				child->prevSibling = NULL;
				after->insertBefore(child);
				if (it != subboxes.begin()) {
					child->prev->nextSibling = child;
					child->prevSibling = child->prev;
					child->prev->afterSubtree = child;
				}
			}
			after->prev->afterSubtree = current->afterSubtree;
			current->bodies.clear();
		}
	}

	void BarnesHutTree::build(vector<Body> bodies, Box domain) {
		Node* current;

		this->clean();
		if (bodies.empty()) return;
		//insert root node
		this->nodes->insertAfter(new Node(this));
		current = this->nodes->next;
		//assign bodies to root node
		current->bodies = bodies;
		//setup proper bounding box
		current->bb = domain;
		current->extendBBforBodies();
		current->extendBBtoCube();
		current->afterSubtree = current->next;
		//iterate over existing boxes and split if it contains too much bodies
		while (current != this->nodes) {
			this->splitNode(current);
			current = current->next;
		}
		this->update();
	}

	void BarnesHutTree::build(vector<Body> bodies) {
		Box bb;

		bb.extendForBodies(bodies);
		this->build(bodies, bb);
	}

	void BarnesHutTree::mergeLET(vector<Body> bodies) {
		for (vector<Body>::iterator it = bodies.begin(); it != bodies.end(); it++) {
			Node* current = this->nodes->next;

			while (!current->leaf) {
				Node* child = current->next;

				while (child != NULL && !child->getBB().contained(it->position)) {
					child = child->nextSibling;
				}
				if (child == NULL) {
					cout << it - bodies.begin() << " Error: not contained: " << it->position[0] << " " << it->position[1] << " " << it->position[2] << endl;
				} else {
					current = child;
				}
			}
			current->bodies.push_back(*it);
			current->bodies.back().refinement = true;
			this->splitNode(current);
		}
		this->update();
	}
}