/*
 * NormalizeTreeLabels.h
 *
 *  Created on: 10. 5. 2017
 *      Author: Aleksandr Shatrovskii
 */

#ifndef _ARBOLOGY_NORMALIZE_TREE_LABELS_H_
#define _ARBOLOGY_NORMALIZE_TREE_LABELS_H_

#include <alib/map>
#include <alib/vector>
#include <alib/tree>

#include <tree/ranked/RankedTree.h>

namespace tree {

namespace simplify {

/**
 * Simple computation of subtree repeats
 */
class NormalizeTreeLabels {
	template < class SymbolType >
	static ext::tree < common::ranked_symbol < unsigned > > normalize ( const ext::tree < common::ranked_symbol < SymbolType > > & tree, ext::map < common::ranked_symbol < SymbolType >, unsigned > & mapping, unsigned & counter );

public:
	/**
	 * Compute a same shaped tree with nodes containing unique subtree ids.
	 * @return Tree of repeats
	 */
	template < class SymbolType >
	static tree::RankedTree < unsigned > normalize ( const tree::RankedTree < SymbolType > & tree );

};

template < class SymbolType >
ext::tree < common::ranked_symbol < unsigned > > NormalizeTreeLabels::normalize ( const ext::tree < common::ranked_symbol < SymbolType > > & tree, ext::map < common::ranked_symbol < SymbolType >, unsigned > & mapping, unsigned & counter ) {

	ext::vector < ext::tree < common::ranked_symbol < unsigned > > > children;

	for ( const ext::tree < common::ranked_symbol < SymbolType > > & child : tree )
		children.push_back ( normalize < SymbolType > ( child, mapping, counter ) );

	auto newLabelIt = mapping.find ( tree.getData ( ) );
	unsigned newLabel;

	if ( newLabelIt == mapping.end ( ) ) {
		mapping.insert ( std::make_pair ( tree.getData ( ), counter ) );
		newLabel = counter;
		counter += 1;
	} else {
		newLabel = newLabelIt->second;
	}

	return ext::tree < common::ranked_symbol < unsigned > > ( common::ranked_symbol < unsigned > ( newLabel, tree.getData ( ).getRank ( ) ), std::move ( children ) );

}

template < class SymbolType >
tree::RankedTree < unsigned > NormalizeTreeLabels::normalize ( const tree::RankedTree < SymbolType > & tree ) {

	ext::map < common::ranked_symbol < SymbolType >, unsigned > mapping;
	unsigned counter = 1;

	return tree::RankedTree < unsigned > ( normalize ( tree.getContent ( ), mapping, counter ) );
}

} /* namespace simplify */

} /* namespace tree */

#endif /* _ARBOLOGY_NORMALIZE_TREE_LABELS_H_ */
