#include <catch2/catch.hpp>

#include "grammar/parsing/Follow.h"
#include "grammar/ContextFree/CFG.h"

TEST_CASE ( "LL1 Follow", "[unit][grammar]" ) {
	SECTION ( "Test 1" ) {
		DefaultSymbolType nE = DefaultSymbolType ( 'E' );
		DefaultSymbolType nT = DefaultSymbolType ( 'T' );
		DefaultSymbolType nF = DefaultSymbolType ( 'F' );

		DefaultSymbolType tP = DefaultSymbolType ( '+' );
		DefaultSymbolType tS = DefaultSymbolType ( '*' );
		DefaultSymbolType tL = DefaultSymbolType ( '(' );
		DefaultSymbolType tR = DefaultSymbolType ( ')' );
		DefaultSymbolType tA = DefaultSymbolType ( 'a' );

		grammar::CFG < > grammar ( nE );
		grammar.setTerminalAlphabet ( ext::set < DefaultSymbolType > { tP, tS, tL, tR, tA } );
		grammar.setNonterminalAlphabet ( ext::set < DefaultSymbolType > { nE, nT, nF } );

		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsE1 ( { nE, tP, nT } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsE2 ( { nT } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsT1 ( { nT, tS, nF } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsT2 ( { nF } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsF1 ( { tA } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsF2 ( { tL, nE, tR } );

		grammar.addRule ( nE, rhsE1 );
		grammar.addRule ( nE, rhsE2 );
		grammar.addRule ( nT, rhsT1 );
		grammar.addRule ( nT, rhsT2 );
		grammar.addRule ( nF, rhsF1 );
		grammar.addRule ( nF, rhsF2 );

		 // --------------------------------------------------
		ext::map < DefaultSymbolType, ext::set < ext::vector < DefaultSymbolType > > > follow;

		follow[nE] = ext::set < ext::vector < DefaultSymbolType > > {
			ext::vector < DefaultSymbolType > ( ), { tP }, { tR }
		};
		follow[nT] = ext::set < ext::vector < DefaultSymbolType > > {
			ext::vector < DefaultSymbolType > ( ), { tP }, { tR }, { tS }
		};
		follow[nF] = ext::set < ext::vector < DefaultSymbolType > > {
			ext::vector < DefaultSymbolType > ( ), { tP }, { tR }, { tS }
		};

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

		ext::map < DefaultSymbolType, ext::set < ext::vector < DefaultSymbolType > > > followAlgo;

		for ( const auto & nt : grammar.getNonterminalAlphabet ( ) )
			followAlgo[nt] = grammar::parsing::Follow::follow ( grammar, nt );

		CHECK ( follow == followAlgo );
	}

	SECTION ( "Test 2" ) {
		DefaultSymbolType nS = DefaultSymbolType ( 'S' );
		DefaultSymbolType nA = DefaultSymbolType ( 'A' );
		DefaultSymbolType nB = DefaultSymbolType ( 'B' );
		DefaultSymbolType nC = DefaultSymbolType ( 'C' );
		DefaultSymbolType nD = DefaultSymbolType ( 'D' );
		DefaultSymbolType nE = DefaultSymbolType ( 'E' );
		DefaultSymbolType nF = DefaultSymbolType ( 'F' );

		DefaultSymbolType tA = DefaultSymbolType ( 'a' );
		DefaultSymbolType tB = DefaultSymbolType ( 'b' );
		DefaultSymbolType tC = DefaultSymbolType ( 'c' );
		DefaultSymbolType tD = DefaultSymbolType ( 'd' );
		DefaultSymbolType tE = DefaultSymbolType ( 'e' );

		grammar::CFG < > grammar ( nS );
		grammar.setTerminalAlphabet ( ext::set < DefaultSymbolType > { tA, tB, tC, tD, tE } );
		grammar.setNonterminalAlphabet ( ext::set < DefaultSymbolType > { nS, nA, nB, nC, nD, nE, nF } );

		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsS1 ( { nB, tD, nS } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsS2 ( { tD, tD, nC } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsS3 ( { tC, nA } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsA1 ( { tA, tE, nE } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsA2 ( { tB, tB, nE } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsB1 ( { tA, nF } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsB2 ( { tB, tB, nD } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsC1 ( { tA, nB, tD } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsC2 ( { tE, nA } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsD1 ( { tC, tA, nF } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsE1 ( { tC, tA, tE, nE } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsE2 ( { } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsF1 ( { tE, nD } );
		ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > rhsF2 ( { } );

		grammar.addRule ( nS, rhsS1 );
		grammar.addRule ( nS, rhsS2 );
		grammar.addRule ( nS, rhsS3 );
		grammar.addRule ( nA, rhsA1 );
		grammar.addRule ( nA, rhsA2 );
		grammar.addRule ( nB, rhsB1 );
		grammar.addRule ( nB, rhsB2 );
		grammar.addRule ( nC, rhsC1 );
		grammar.addRule ( nC, rhsC2 );
		grammar.addRule ( nD, rhsD1 );
		grammar.addRule ( nE, rhsE1 );
		grammar.addRule ( nE, rhsE2 );
		grammar.addRule ( nF, rhsF1 );
		grammar.addRule ( nF, rhsF2 );

		 // --------------------------------------------------
		ext::map < DefaultSymbolType, ext::set < ext::vector < DefaultSymbolType > > > follow;

		follow[nS] = ext::set < ext::vector < DefaultSymbolType > > {
			ext::vector < DefaultSymbolType > ( )
		};
		follow[nA] = ext::set < ext::vector < DefaultSymbolType > > {
			ext::vector < DefaultSymbolType > ( )
		};
		follow[nB] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tD }
		};
		follow[nC] = ext::set < ext::vector < DefaultSymbolType > > {
			ext::vector < DefaultSymbolType > ( )
		};
		follow[nD] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tD }
		};
		follow[nE] = ext::set < ext::vector < DefaultSymbolType > > {
			ext::vector < DefaultSymbolType > ( )
		};
		follow[nF] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tD }
		};

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

		ext::map < DefaultSymbolType, ext::set < ext::vector < DefaultSymbolType > > > followAlgo;

		for ( const auto & nt : grammar.getNonterminalAlphabet ( ) )
			followAlgo[nt] = grammar::parsing::Follow::follow ( grammar, nt );

		CHECK ( follow == followAlgo );
	}
}
