// NexusWriter.hpp


#if !defined(NEXUS_WRITER_HPP)
#define NEXUS_WRITER_HPP


#include <map>
#include <ostream>
#include <stdexcept>
#include <string>
#include <vector>
#include <cassert>


namespace Cipres {

namespace SequenceFormats {


class NexusWriter {
public:
	enum DataType {
		STANDARD,
		DNA,
		RNA,
		NUCLEOTIDE,
		PROTEIN,
		CONTINUOUS
	};


	class StandardDataBlock;
	class MolecularDataBlock;
	class ContinuousDataBlock;


	class Taxon {
	public:
		typedef std::vector<char> DiscreteCharList;
		typedef std::vector<float> ContinuousCharList;


		virtual ~Taxon() throw() { }

		size_t Ordinal() const throw()
		{
			return m_ordinal;
		}

		const std::string &Name() const throw()
		{
			return m_name;
		}

		size_t NumChars() const throw()
		{
			return m_num_chars;
		}

	protected:
		Taxon(size_t ordinal, const std::string &name) throw(std::bad_alloc) : m_ordinal(ordinal), m_name(name), m_num_chars(0) { }


		const size_t m_ordinal;
		const std::string m_name;
		size_t m_num_chars;
	};

	class DiscreteTaxon : public Taxon {
		friend class StandardDataBlock;
		friend class MolecularDataBlock;

	public:
		~DiscreteTaxon() throw() { }

		void AddChar(char state) throw(std::bad_alloc)
		{
			m_characters.push_back(state + 4);

			m_num_chars += 1;
		}

		void AddPolyChar(const DiscreteCharList states) throw(std::bad_alloc)
		{
			m_characters.push_back(0);

			for(DiscreteCharList::const_iterator i = states.begin() ; i != states.end() ; i++)
				m_characters.push_back(*i + 4);

			m_characters.push_back(1);

			m_num_chars += 1;
		}

	private:
		DiscreteTaxon(size_t ordinal, const std::string name) throw(std::bad_alloc) : Taxon(ordinal, name) { }

		DiscreteTaxon(size_t ordinal, const std::string name, size_t num_chars) throw(std::bad_alloc) : Taxon(ordinal, name)
		{
			m_characters.reserve(num_chars);
		}


		DiscreteCharList m_characters;
	};

	class ContinuousTaxon : public Taxon {
		friend class ContinuousDataBlock;

	public:
		~ContinuousTaxon() throw() { }

		void AddChar(float state) throw(std::bad_alloc)
		{
			m_characters.push_back(state);

			m_num_chars += 1;
		}

	private:
		ContinuousTaxon(size_t ordinal, const std::string name) throw(std::bad_alloc) : Taxon(ordinal, name) { }

		ContinuousTaxon(size_t ordinal, const std::string name, size_t num_chars) throw(std::bad_alloc) : Taxon(ordinal, name)
		{
			m_characters.reserve(num_chars);
		}


		ContinuousCharList m_characters;
	};


	class DataBlock {
		friend class NexusWriter;

	public:
		virtual ~DataBlock() throw() { }

	protected:
		typedef std::map<size_t, DiscreteTaxon *> DiscreteTaxonMap;
		typedef std::map<size_t, ContinuousTaxon *> ContinuousTaxonMap;


		DataBlock(NexusWriter &owner) throw() : m_owner(owner) { }

		virtual size_t MaxNumChars() const throw() = 0;

		virtual void WriteCharBlock(std::ostream &output) const throw(std::ios_base::failure) = 0;

		virtual bool WriteCharSequence(size_t ordinal, std::ostream &output) const throw(std::ios_base::failure) = 0;


		NexusWriter &m_owner;
	};

	class StandardDataBlock : public DataBlock {
		friend class NexusWriter;

	public:
		~StandardDataBlock() throw()
		{
			for(DiscreteTaxonMap::iterator i = m_taxa.begin() ; i != m_taxa.end() ; i++)
				delete (*i).second;
		}

		void SetSymbols(const std::string &symbols) throw(std::bad_alloc, std::runtime_error)
		{
			m_owner.SetSymbols(symbols);
		}

		DiscreteTaxon &AddTaxon(const std::string name) throw(std::bad_alloc, std::runtime_error);

	private:
		StandardDataBlock(NexusWriter &owner) throw() : DataBlock(owner) { }

		size_t MaxNumChars() const throw();

		void WriteCharBlock(std::ostream &output) const throw(std::ios_base::failure);

		bool WriteCharSequence(size_t ordinal, std::ostream &output) const throw(std::ios_base::failure);


		DiscreteTaxonMap m_taxa;
	};

	class MolecularDataBlock : public DataBlock {
		friend class NexusWriter;

	public:
		~MolecularDataBlock() throw()
		{
			for(DiscreteTaxonMap::iterator i = m_taxa.begin() ; i != m_taxa.end() ; i++)
				delete (*i).second;
		}

		DiscreteTaxon &AddTaxon(const std::string name) throw(std::bad_alloc, std::runtime_error);

	private:
		MolecularDataBlock(NexusWriter &owner) throw() : DataBlock(owner) { }

		char CodeToSymbol(unsigned char code) const throw()
		{
			static const char symbols[40] = {
				// either
				'(', ')', '-', '?',
				// nucleotides
				'A', 'C', 'G', 'T', 'U', 'M', 'R', 'W', 'S', 'Y', 'K', 'V', 'H', 'D', 'B', 'N',
				// amino acids
				'A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y'
			};

			assert(code < sizeof(symbols) / sizeof(char));

			return symbols[code];
		}

		size_t MaxNumChars() const throw();

		void WriteCharBlock(std::ostream &output) const throw(std::ios_base::failure);

		bool WriteCharSequence(size_t ordinal, std::ostream &output) const throw(std::ios_base::failure);


		DiscreteTaxonMap m_taxa;
	};

	class ContinuousDataBlock : public DataBlock {
		friend class NexusWriter;

	public:
		~ContinuousDataBlock() throw()
		{
			for(ContinuousTaxonMap::iterator i = m_taxa.begin() ; i != m_taxa.end() ; i++)
				delete (*i).second;
		}

		ContinuousTaxon &AddTaxon(const std::string name) throw(std::bad_alloc, std::runtime_error);

	private:
		ContinuousDataBlock(NexusWriter &owner) throw() : DataBlock(owner) { }

		size_t MaxNumChars() const throw();

		void WriteCharBlock(std::ostream &output) const throw(std::ios_base::failure);

		bool WriteCharSequence(size_t ordinal, std::ostream &output) const throw(std::ios_base::failure);


		ContinuousTaxonMap m_taxa;
	};


	NexusWriter() throw() { }

	~NexusWriter() throw()
	{
		for(DataBlockList::iterator i = m_data_blocks.begin() ; i != m_data_blocks.end() ; i++)
			delete *i;
	}

	size_t NumDataBlocks() const throw()
	{
		return m_data_blocks.size();
	}

	size_t NumTaxa() const throw()
	{
		return m_taxa_by_name.size();
	}

	void SetTitle(const std::string &title) throw(std::bad_alloc)
	{
		m_title.assign(title);
	}

	StandardDataBlock &AddStandardDataBlock() throw(std::bad_alloc, std::runtime_error);

	MolecularDataBlock &AddDnaDataBlock() throw(std::bad_alloc, std::runtime_error);

	MolecularDataBlock &AddProteinDataBlock() throw(std::bad_alloc, std::runtime_error);

	ContinuousDataBlock &AddContinuousDataBlock() throw(std::bad_alloc, std::runtime_error);

	void Write(std::ostream &output) const throw(std::bad_alloc, std::runtime_error, std::ios_base::failure);

private:
	typedef std::map<std::string, size_t> TaxonOrdinalMap;
	typedef std::vector<std::string> TaxonNameList;
	typedef std::vector<DataBlock *> DataBlockList;


	void SetSymbols(const std::string &symbols) throw(std::bad_alloc, std::runtime_error)
	{
		if(symbols.find_first_of(" \t\n\r-.?()") != std::string::npos)
			throw std::runtime_error("invalid state symbol");

		m_symbols.reserve(symbols.size() + 4);

		m_symbols.assign("()-?", 4);
		m_symbols.append(symbols);
	}

	char StateToSymbol(unsigned char state) const throw()
	{
		assert(state < m_symbols.size());

		return m_symbols[state];
	}

	size_t AddTaxon(const std::string name) throw(std::bad_alloc);

	void WriteTaxaBlock(std::ostream &output) const throw(std::bad_alloc, std::ios_base::failure);


	unsigned int m_data_type;
	std::string m_title;
	std::string m_symbols;
	DataBlockList m_data_blocks;
	TaxonNameList m_taxa_by_ordinal;
	TaxonOrdinalMap m_taxa_by_name;
};


} // namespace SequenceFormats

} // namespace Cipres


#endif // NEXUS_WRITER_HPP