/*
 * GlushkovFollow.h
 *
 *  Created on: 26. 7. 2017
 *	  Author: Tomas Pecka
 */

#ifndef RTE_GLUSHKOV_FOLLOW_H_
#define RTE_GLUSHKOV_FOLLOW_H_

#include <alib/map>
#include <alib/set>
#include <alib/vector>
#include <common/ranked_symbol.hpp>

#include <rte/formal/FormalRTE.h>
#include <rte/formal/FormalRTEElements.h>

#include "GlushkovFirst.h"

namespace rte {


	template < class SymbolType >
	using TFollowTuple = ext::vector < ext::set < common::ranked_symbol < ext::pair < SymbolType, unsigned > > > >;

	template < class SymbolType >
	using TSetOfSymbols = ext::set < common::ranked_symbol < ext::pair < SymbolType, unsigned > > >;

	template < class SymbolType >
	using TSubstMap = ext::map < common::ranked_symbol < ext::pair < SymbolType, unsigned > >, TSetOfSymbols < SymbolType > >;

	template < class SymbolType >
	using TFollowMap = ext::map < common::ranked_symbol < ext::pair < SymbolType, unsigned > >, TFollowTuple < SymbolType > >;

class GlushkovFollow {


	template < class SymbolType >
	static TFollowTuple < SymbolType > replaceConstants ( const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, const ext::vector < TSetOfSymbols < SymbolType > > & children, const TSubstMap < SymbolType > & subMap );

	template < class SymbolType >
	static void preprocessSubMap ( TSubstMap < SymbolType > & subMap, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte );

public:
	/**
	 * @param re rte to probe
	 * @return all symbols that can follow specific symbol in tree
	 */
	template < class SymbolType >
	static ext::map < common::ranked_symbol < ext::pair < SymbolType, unsigned > >, TFollowTuple < SymbolType > > follow ( const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte );

	template < class SymbolType >
	class Formal {
	public:
		static void visit ( const FormalRTEElement        < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & substMap, TFollowMap < SymbolType > & res );
		static void visit ( const FormalRTEAlternation    < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & res );
		static void visit ( const FormalRTESubstitution   < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & res );
		static void visit ( const FormalRTEIteration      < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & res );
		static void visit ( const FormalRTESymbolAlphabet < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & res );
		static void visit ( const FormalRTESymbolSubst    < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & res );
		static void visit ( const FormalRTEEmpty          < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & res );
	};
};

template < class SymbolType >
ext::map < common::ranked_symbol < ext::pair < SymbolType, unsigned > >, TFollowTuple < SymbolType > > GlushkovFollow::follow ( const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte ) {
	TSubstMap < SymbolType > subMap;
	TFollowMap < SymbolType > res;

	/* initialize subMap keys
	   \forall a \in K: sub[a] = \emptyset
	*/
	for ( const common::ranked_symbol < ext::pair < SymbolType, unsigned > > & ssymb : rte.getSubstitutionAlphabet ( ) )
		subMap.insert ( std::make_pair ( ssymb, TSetOfSymbols < SymbolType > { } ) );

	/* initialize follow keys
	   \forall a \in F: follow[a] = \emptyset
	*/
	/*
	for ( const common::ranked_symbol < ext::pair < SymbolType, unsigned > > & symb : rte.getAlphabet ( ) )
		res.insert ( std::make_pair ( symb, TFollowTuple < SymbolType > ( ) ) );
	*/

	rte.getRTE ( ).getStructure ( ).template accept < void, GlushkovFollow::Formal < SymbolType > > ( rte, subMap, res );
	return res;
}

// -----------------------------------------------------------------------------

template < class SymbolType >
TFollowTuple < SymbolType > GlushkovFollow::replaceConstants ( const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, const ext::vector < TSetOfSymbols < SymbolType > > & children, const TSubstMap < SymbolType > & subMap ) {
	TFollowTuple < SymbolType > follow;

	for ( const ext::set < common::ranked_symbol < ext::pair < SymbolType, unsigned > > > & s : children ) {
		ext::set < common::ranked_symbol < ext::pair < SymbolType, unsigned > > > processed;
		for ( const common::ranked_symbol < ext::pair < SymbolType, unsigned > > & e : s ) {
			if ( rte.getSubstitutionAlphabet ( ).count ( e ) > 0 )
				processed.insert ( subMap.at ( e ).begin ( ), subMap.at ( e ).end ( ) );
			else
				processed.insert ( e );
		}
		follow.push_back ( processed );
	}

	return follow;
}

/**
 * Preprocessing:
 *  - Let k1, k2 be elements of alphabet K.
 *  - If k1 is an element of substMap[k2], then copy content of substMap[k1] into substMap[k2]
 */
template < class SymbolType >
void GlushkovFollow::preprocessSubMap ( TSubstMap < SymbolType > & subMap, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte ) {
	const TSetOfSymbols < SymbolType > & alphabetK = rte.getSubstitutionAlphabet ( );
	for ( bool change = true; change; change = false )
		for ( std::pair < const common::ranked_symbol < ext::pair < SymbolType, unsigned > >, TSetOfSymbols < SymbolType > > & kv : subMap ) {
			TSetOfSymbols < SymbolType > & substSet = kv.second;

			for ( auto eIter = substSet.begin ( ); eIter != substSet.end ( ); ) {
				if ( alphabetK.count ( * eIter ) == 0 ) {
					++eIter;
				} else {
					auto it = subMap.find ( * eIter );
					size_t oldSize = substSet.size ( );
					substSet.insert ( it->second.begin ( ), it->second.end ( ) );
					change = ( oldSize != substSet.size ( ) ); // something was added
					eIter = substSet.erase ( eIter );
				}
			}
		}
}


// -----------------------------------------------------------------------------

template < class SymbolType >
void GlushkovFollow::Formal < SymbolType > ::visit ( const rte::FormalRTEAlternation < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & res ) {

	/* update substitution mapping */
	// none

	// just recurse here
	node.getLeftElement  ( ).template accept < void, GlushkovFollow::Formal < SymbolType > > ( rte, subMap, res );
	node.getRightElement ( ).template accept < void, GlushkovFollow::Formal < SymbolType > > ( rte, subMap, res );
}

template < class SymbolType >
void GlushkovFollow::Formal < SymbolType > ::visit ( const rte::FormalRTESubstitution < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & res ) {
	preprocessSubMap ( subMap, rte );

	/* update substitution mapping */
	// prepare left map
	auto itMap = subMap.find ( node.getSubstitutionSymbol ( ).getSymbol ( ) );
	const ext::set < common::ranked_symbol < ext::pair < SymbolType, unsigned > > > backup = std::move ( itMap->second );
	itMap->second.clear();

	for ( const auto & s : node.getRightElement ( ).template accept < ext::set < common::ranked_symbol < ext::pair < SymbolType, unsigned > > >, GlushkovFirst::Formal < ext::pair < SymbolType, unsigned > > > ( ) )
		itMap->second.insert ( s );

	node.getLeftElement ( ).template accept < void, GlushkovFollow::Formal < SymbolType > > ( rte, subMap, res );

	// restore original map
	itMap->second = std::move ( backup );

	node.getRightElement ( ).template accept < void, GlushkovFollow::Formal < SymbolType > > ( rte, subMap, res );
}

template < class SymbolType >
void GlushkovFollow::Formal < SymbolType > ::visit ( const rte::FormalRTEIteration < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & res ) {
	preprocessSubMap ( subMap, rte );

	/* update substitution mapping */
	auto itMap = subMap.find ( node.getSubstitutionSymbol ( ).getSymbol ( ) );
	const TSetOfSymbols < SymbolType > backup = itMap->second;
	for ( const auto & s : node.getElement ( ).template accept < TSetOfSymbols < SymbolType >, GlushkovFirst::Formal < ext::pair < SymbolType, unsigned > > > ( ) ) {
		subMap [ node.getSubstitutionSymbol ( ).getSymbol ( ) ].insert ( s );
	}

	node.getElement ( ).template accept < void, GlushkovFollow::Formal < SymbolType > > ( rte, subMap, res );

	/* restore substitution mapping */
	itMap->second = backup;
}

template < class SymbolType >
void GlushkovFollow::Formal < SymbolType > ::visit ( const rte::FormalRTESymbolAlphabet < ext::pair < SymbolType, unsigned > > & node, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & res ) {
	preprocessSubMap ( subMap, rte );

	/* update substitution mapping */
	// none

	/* compute first for all children and replace substitution symbols */
	ext::vector < ext::set < common::ranked_symbol < ext::pair < SymbolType, unsigned > > > > children;
	for ( const rte::FormalRTEElement < ext::pair < SymbolType, unsigned > > & c : node.getElements ( ) )
		children.push_back ( c.template accept < TSetOfSymbols < SymbolType >, GlushkovFirst::Formal < ext::pair < SymbolType, unsigned > > > ( ) );

	res.insert ( std::make_pair ( node.getSymbol ( ), replaceConstants ( rte, children, subMap ) ) );

	// call follow on children
	for ( const auto & childNode : node.getElements ( ) ) {
		childNode . template accept < void, GlushkovFollow::Formal < SymbolType > > ( rte, subMap, res );
	}
}

template < class SymbolType >
void GlushkovFollow::Formal < SymbolType > ::visit ( const rte::FormalRTESymbolSubst < ext::pair < SymbolType, unsigned > > & /* node */, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & /* res */ ) {
	preprocessSubMap ( subMap, rte );
}

template < class SymbolType >
void GlushkovFollow::Formal < SymbolType > ::visit ( const rte::FormalRTEEmpty < ext::pair < SymbolType, unsigned > > & /* node */, const rte::FormalRTE < ext::pair < SymbolType, unsigned > > & rte, TSubstMap < SymbolType > & subMap, TFollowMap < SymbolType > & /* res */ ) {
	preprocessSubMap ( subMap, rte );
}

} /* namespace rte */

#endif /* RTE_GLUSHKOV_FOLLOW_H_ */
