#ifndef PHYCPP_MODELS_BROWNIANMOTION_HPP
#define PHYCPP_MODELS_BROWNIANMOTION_HPP

#define _USE_MATH_DEFINES
#include <cmath>
#include <map>
#include <utility>

namespace PhyCpp {

struct BrownianMotion {

	struct ReconstructionResult {
		double logLikelihood;
		double rate;
	};

	// ML reconstruction of ancestral state for the given taxon character into the given node character
	// taxa with value equal to missingTaxonValue are treated as missing data
	template <typename Tree, typename TaxonValueAccessor, typename NodeValueAccessor>
	static ReconstructionResult reconstruct(Tree * tree, const TaxonValueAccessor & tval, NodeValueAccessor & nval, double missingTaxonValue) {

		// TODO: ASSERT THAT ALL TIPS HAVE TAXON VALUES

		ReconstructionResult result;

		std::map<typename Tree::Node *, std::pair<double, double> > q;

		typename Tree::Node::reverse_iterator rit = tree->rbegin();
		while (rit != tree->rend()) {
			if (rit->taxon() == 0 || tval(rit) == missingTaxonValue) {
				std::pair<double, double> Q = std::make_pair<double, double>(0.0, 0.0);
				typename Tree::Node * n = rit->leftChild();
				while (n) {
					if (n->taxon() != 0 && tval(n) != missingTaxonValue) {
						Q.first += -1.0/(2 * n->parentEdge()->length());
						Q.second += tval(n) / n->parentEdge()->length();
					} else {
						Q.first += (-2 * q[n].first) / (-2.0 + 4 * q[n].first * n->parentEdge()->length());
						Q.second += (-2 * q[n].second) / (-2.0 + 4 * q[n].first * n->parentEdge()->length());
					}
					n = n->rightSibling();
				}
				q[*rit] = Q;
			}
			++rit;
		}

		typename Tree::Node::iterator it = tree->begin();
		nval(it) = (it->taxon() == 0 || tval(it) == missingTaxonValue) ? -q[*it].second / (2 * q[*it].first) : tval(it);
		++it;
		while (it != tree->end()) {
			nval(it) = (it->taxon() != 0 && tval(it) != missingTaxonValue) ? tval(it) : (nval(it->parent()) + q[*it].second * it->parentEdge()->length()) / (1 - 2 * q[*it].first * it->parentEdge()->length());
			++it;
		}

		// Calculate maximum likelihood variance
		double v = 0;
		it = tree->begin();
		++it;
		while (it != tree->end()) {
			v += pow((nval(it) - nval(it->parent())),2)/it->parentEdge()->length();
			++it;
		}
		result.rate = v/(tree->size()-1);

		// Calculate likelihood of reconstruction
		result.logLikelihood = 0;
		it = tree->begin();
		++it;
		while (it != tree->end()) {
			double x = nval(it) - nval(it->parent());
			double t = it->parentEdge()->length();
			result.logLikelihood += (-(x*x)/(2 * result.rate * t) - log(sqrt(2 * M_PI * result.rate * t)));
			++it;
		}

		return result;
	}

};

}

#endif // PHYCPP_MODELS_BROWNIANMOTION_HPP
