#include <iostream>
#include <limits>
#include <numeric>

#include "BarnesHutTree.hpp"
#include "Box.hpp"
#include "Node.hpp"

namespace nbody {
	BarnesHutTree::BarnesHutTree(std::size_t parallelId) : Tree(parallelId) {}

	//determine octree subboxes
	std::array<Box, 8> BarnesHutTree::splitBB(const Node* node) {
		return node->getBB().octreeSplit();
	}

	std::size_t BarnesHutTree::numberOfChildren() const {
		return 8;
	}

	//update upper tree nodes according to moved particles
	void BarnesHutTree::update() {
		//iterate for updating representatives
		for (auto& current : reverse(nodesView)){
			auto siblingNodesView = getSiblingNodesView(&current);
			auto& currRep = current.representative;
			currRep.id = std::numeric_limits<std::size_t>::max();
			currRep.mass = 0.0;
			currRep.position = Vec3{ 0.0 };
			currRep.mass = [&] {
				return current.leaf 
					? std::accumulate(std::begin(current.bodies),   std::end(current.bodies),   0.0, [](double m, const Body& b) { return m + b.mass; })
					: std::accumulate(std::begin(siblingNodesView), std::end(siblingNodesView), 0.0, [](double m, const Node& n) { return m + n.representative.mass; }); }();
			currRep.position = [&] {
				return current.leaf 
					? std::accumulate(std::begin(current.bodies),   std::end(current.bodies),   Vec3{}, [](Vec3 p, const Body& b) { return p + (b.position * b.mass); })
					: std::accumulate(std::begin(siblingNodesView), std::end(siblingNodesView), Vec3{}, [](Vec3 p, const Node& n) { return p + (n.representative.position * n.representative.mass); }); }();
			currRep.position = (currRep.mass > std::numeric_limits<float>::epsilon())
				? currRep.position / currRep.mass
				: Vec3{ 0.0 };
		}
	}

	//split tree node into sub-boxes during tree build
	void BarnesHutTree::split(Node* current) {
		const auto& subboxes = BarnesHutTree::splitBB(current);
		current->leaf = false;
		Node* after = current->next;

		bool first = true;
		for (const auto& subBox : subboxes) {
			auto child = current->tree->createNode(subBox, subBox.copyBodies(current->bodies), current->tree);

			after->insertBefore(child);
			if (!first) {
				child->prev->nextSibling = child;
				child->prevSibling = child->prev;
				child->prev->afterSubtree = child;
			}
			first = false;
		}
		after->prev->afterSubtree = current->afterSubtree;
		current->bodies.clear();
	}

	//initialize tree for build process
	void BarnesHutTree::init(std::vector<Body> bodies, Box domain) {
		Node* current;

		clean();
		if (bodies.empty()) return;
		//insert root node
		nodes->insertAfter(createNode(this));
		current = nodes->next;
		//assign bodies to root node
		current->bodies = std::move(bodies);
		//setup proper bounding box
		current->bb = domain;
		current->extendBBforBodies();
		current->extendBBtoCube();
		current->afterSubtree = current->next;
	}

	//check if split is required and perform it
	bool BarnesHutTree::splitNode(Node* current) {
		bool result = current->isSplitable();

		if (result) {
			split(current);
		}
		return result;
	}

	//build tree with given domain
	void BarnesHutTree::build(std::vector<Body> bodies, Box domain) {
		init(std::move(bodies), domain);
		//iterate over existing boxes and split if it contains too much bodies
		BarnesHutTree::splitSubtree(nodes->next);
		update();
	}

	//build tree
	void BarnesHutTree::build(std::vector<Body> bodies) {
		Box bb{ bodies };
		build(std::move(bodies), bb);
	}

	//merge remote refinement particles into local tree
	//(this are remote particles from other processes needed for force computation on local particles)
	void BarnesHutTree::mergeLET(const std::vector<Body>& bodies) {
		//put all new bodies into fitting leaves, walk through tree and split
		Node* current;

		for (const auto& b : bodies) {
			current = nodes->next;
			while (!current->leaf) {
				Node* child = current->next;

				while (child != nullptr && !child->getBB().contained(b.position)) {
					child = child->nextSibling;
				}
				//TODO(pheinzlr): check for child == nullptr?
				current = child;
			}
			current->bodies.push_back(b);
			current->bodies.back().refinement = true;
		}
		for (auto& node : nodesView) {
			splitNode(&node);
		}
		update();
	}

	//node splitting if required
	void BarnesHutTree::splitSubtree(Node* root) {
		bool toSplitLeft;
		Node* current = root;

		do {
			toSplitLeft = false;
			while (current != root->afterSubtree) {
				if (current->isSplitable()) {
					split(current);
					toSplitLeft = true;
				}
				current = current->next;
			}
		} while (toSplitLeft);
	}
} // namespace nbody
