#include <algorithm>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>

#include "Node.hpp"
#include "Simulation.hpp"
#include "Tree.hpp"

namespace nbody {
	Tree::Tree(std::size_t parallelId):
		nodes(createNode(this)), //insert dummy root node
		maxLeafBodies(16),
		parallelId_(parallelId) {}

	void Tree::clean() {
		//remove all nodes; refresh dummy first node
		nodesMemory.clear();
		nodes = createNode(this);
	}


	bool Tree::isCorrect() const {
		auto nodesView = getNodesView();
		return std::all_of(std::begin(nodesView), std::end(nodesView), [](const Node& n) { return n.isCorrect(); });
	}

	//accumulate force from the whole local tree on parameter body
	void Tree::accumulateForceOnto(Body& body) {
		Node* n = nodes->next;

		body.resetAcceleration();
		while (n != nodes) {
			if (n->sufficientForBody(body)) {
				body.accumulateForceOntoBody(n->representative);
				n = n->afterSubtree;
			} else if (n->leaf) {
				for (const auto& b : n->bodies) {
					body.accumulateForceOntoBody(b);
				}
				n = n->afterSubtree;
			} else {
				n = n->next;
			}
		}
	}

	//accumulate forces for whole tree (local particles)
	void Tree::computeForces() {
		foreach_localParticle([&](Body& b) { accumulateForceOnto(b); });
	}

	//get local bodies for rebuildiong tree after moving particles
	std::vector<Body> Tree::extractLocalBodies() {
		std::vector<Body> localBodies;
		for (auto& node : getNodesView()) {
			if (node.leaf) {
				std::copy_if(std::begin(node.bodies), std::end(node.bodies), std::back_inserter(localBodies), [](const Body& b) {return !b.refinement; });
			}
		}
		clean();
		return localBodies;
	}

	//get refinement particles required to compute forces within remote domains
	std::vector<Body> Tree::copyRefinements(Box domain) const {
		Node* current = nodes->next;
		if (!current->bb.isValid()) { //empty tree means no refinements
			return {};
		}
		std::vector<Body> refinements;
		while (current != nodes) {
			const auto isSufficient = current->sufficientForBox(domain);
			if (isSufficient && current->representative.mass > 0.0) {
					refinements.push_back(current->representative);
			} else if (current->leaf) {
				refinements.insert(std::end(refinements), std::begin(current->bodies), std::end(current->bodies));
			}
			current = isSufficient ? current->afterSubtree : current->next;
		}
		return refinements;
	}

	//rebuild with predefined root node bounding box
	void Tree::rebuild(Box domain) {
		build(extractLocalBodies(), domain);
	}

	void Tree::rebuild() {
		build(extractLocalBodies());
	}

	//rebuild with predefined root node bounding box and bodies
	void Tree::rebuild(Box domain, const std::vector<Body>& bodies) {
		build(bodies, domain);
	}

	//move particles according to accumulated forces
	Box Tree::advance() {
		Box bb;
		foreach_localParticle([&](Body& b) { 
			b.integrate();
			bb.extend(b); 
		});
		return bb;
	}

	void Tree::print(std::size_t parallelId) const {
		for (const auto& n : getNodesView()) {
			n.bb.printBB(parallelId);
			if (n.leaf) {
				for (const auto& b : n.bodies) {
					if (!b.refinement) {
						b.print(this->parallelId_);
					}
				}
			}
		}
	}
} // namespace nbody
