#ifndef BRANCHGROUP_H
#define BRANCHGROUP_H 1

#include "Common/Algorithms.h"
#include "Common/Exception.h"
#include <algorithm> // for swap
#include <map>
#include <utility>

enum BranchGroupStatus
{
	BGS_ACTIVE,
	BGS_NOEXT,
	BGS_JOINED,
	BGS_TOOLONG,
	BGS_TOOMANYBRANCHES
};

/** A container of BranchRecord. */
class BranchGroup
{
	public:
		typedef std::vector<BranchRecord> BranchGroupData;
		typedef BranchGroupData::iterator iterator;
		typedef BranchGroupData::const_iterator const_iterator;

		BranchGroup()
			: m_dir(SENSE), m_maxNumBranches(0),
			m_noExt(false), m_status(BGS_ACTIVE)
			{ }

		BranchGroup(extDirection dir, size_t maxNumBranches,
				const BranchRecord::V &origin)
			: m_dir(dir), m_origin(origin),
			m_maxNumBranches(maxNumBranches), m_noExt(false),
			m_status(BGS_ACTIVE)
		{
			m_branches.reserve(m_maxNumBranches);
		}

		BranchGroup(extDirection dir, size_t maxNumBranches,
				const BranchRecord::V &origin, const BranchRecord& branch)
			: m_dir(dir), m_origin(origin),
			m_maxNumBranches(maxNumBranches), m_noExt(false),
			m_status(BGS_ACTIVE)
		{
			m_branches.reserve(m_maxNumBranches);
			m_branches.push_back(branch);
		}

		BranchGroup(const BranchGroup& o)
			: m_branches(o.m_branches), m_dir(o.m_dir),
			m_origin(o.m_origin),
			m_maxNumBranches(o.m_maxNumBranches),
			m_noExt(o.m_noExt), m_status(o.m_status)
		{
			m_branches.reserve(m_maxNumBranches);
		}

		/** Add a branch to this group. */
		BranchRecord& addBranch(const BranchRecord& branch)
		{
			assert(m_branches.size() < m_maxNumBranches);
			m_branches.push_back(branch);
			return m_branches.back();
		}

		/** Add a branch to this group and extend the new branch with
		 * the given k-mer. */
		void addBranch(const BranchRecord& branch,
				const BranchRecord::V& kmer)
		{
			if (m_branches.size() < m_maxNumBranches)
				addBranch(branch).push_back(
						std::make_pair(kmer, BranchRecord::VP()));
			else
				m_status = BGS_TOOMANYBRANCHES;
		}

		/** Return the specified branch. */
		BranchRecord& operator [](unsigned id)
		{
			return m_branches[id];
		}

		/** Return the number of branches in this group. */
		size_t size() const { return m_branches.size(); }

		/** Return whether a branch contains the specified k-mer at
		 * the index i. */
		bool exists(unsigned i, const BranchRecord::V& kmer) const
		{
			for (BranchGroupData::const_iterator it
					= m_branches.begin();
					it != m_branches.end(); ++it)
				if (it->exists(i, kmer))
					return true;
			return false;
		}

		// return the current status of the branch
		BranchGroupStatus getStatus() const { return m_status; }

		// set the no extension flag
		void setNoExtension() { m_noExt = true; }

		// is the no extension flag set?
		bool isNoExt() const { return m_noExt; }

		// return the direction of growth
		extDirection getDirection() const { return m_dir; }

		iterator begin() { return m_branches.begin(); }
		iterator end() { return m_branches.end(); }
		const_iterator begin() const { return m_branches.begin(); }
		const_iterator end() const { return m_branches.end(); }

// Check the stop conditions for the bubble growth
BranchGroupStatus
updateStatus(unsigned maxLength)
{
	assert(m_branches.size() <= m_maxNumBranches);

	if (m_status != BGS_ACTIVE)
		return m_status;

	// Check if the no extension flag is set
	if(m_noExt)
	{
		m_status = BGS_NOEXT;
		return m_status;
	}

	// Check if any branches are too long or any sequence has a loop
	for (BranchGroupData::const_iterator iter = m_branches.begin();
			iter != m_branches.end(); ++iter) {
		if (iter->isTooLong(maxLength)) {
			m_status = BGS_TOOLONG;
			return m_status;
		}
	}

	BranchGroupData::const_iterator it = m_branches.begin();
	const BranchRecord::V& lastSeq = it->back().first;
	while (++it != m_branches.end())
		if (it->back().first != lastSeq)
			return m_status = BGS_ACTIVE;

	// All the branches of the bubble have joined.
	// Remove the last base, which is identical for every branch.
	std::for_each(m_branches.begin(), m_branches.end(),
			std::mem_fun_ref(&BranchRecord::pop_back));

	// Sort the branches by coverage.
	sort_by_transform(m_branches.begin(), m_branches.end(),
			std::mem_fun_ref(&BranchRecord::calculateBranchMultiplicity));
	reverse(m_branches.begin(), m_branches.end());

	return m_status = BGS_JOINED;
}

/** Return whether any branches of this group are active. */
bool
isActive() const
{
	for (BranchGroupData::const_iterator it = m_branches.begin();
			it != m_branches.end(); ++it)
		if (it->isActive())
			return true;
	return false;
}

/** Return whether this branch is extendable. */
bool
isExtendable()
{
	if (m_noExt)
		return false;

	// A group is extendable when all the branches are the same
	// length. All the branches are lockstepped for growth.
	BranchGroupData::iterator it = m_branches.begin();
	unsigned length = it++->size();
	for (; it != m_branches.end(); ++it)
		if (it->size() != length)
			return false;
	return true;
}

/** Return whether this branch is ambiguous at its origin. Also
 * returns false if the origin of the branch has since been deleted.
 */
bool
isAmbiguous(const SequenceCollectionHash& g) const
{
	// Get fresh data from the collection to check that this bubble
	// does in fact still exist.
	const BranchRecord::VP& data = g.getSeqAndData(m_origin).second;
	return data.deleted() ? false : data.isAmbiguous(m_dir);
}

	private:
		BranchGroup& operator =(const BranchGroup& o);

		BranchGroupData m_branches;
		extDirection m_dir;
		BranchRecord::V m_origin;
		size_t m_maxNumBranches;
		bool m_noExt;
		BranchGroupStatus m_status;
};

#endif