#include <catch2/catch.hpp>

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

TEST_CASE ( "LL1 First", "[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 < ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > >, ext::set < ext::vector < DefaultSymbolType > > > first;

		first[rhsE1] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tA }, { tL }
		};
		first[rhsE2] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tA }, { tL }
		};
		first[rhsT1] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tA }, { tL }
		};
		first[rhsT2] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tA }, { tL }
		};
		first[rhsF1] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tA }
		};
		first[rhsF2] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tL }
		};

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

		ext::map < ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > >, ext::set < ext::vector < DefaultSymbolType > > > firstAlgo;

		for ( const auto & rule : grammar::RawRules::getRawRules ( grammar ) )
			for ( const auto & rhs : rule.second )
				firstAlgo[rhs] = grammar::parsing::First::first ( grammar, rhs );

		CHECK ( first == firstAlgo );
	}

	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 < ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > >, ext::set < ext::vector < DefaultSymbolType > > > first;

		first[rhsS1] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tA }, { tB }
		};
		first[rhsS2] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tD }
		};
		first[rhsS3] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tC }
		};
		first[rhsA1] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tA }
		};
		first[rhsA2] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tB }
		};
		first[rhsB1] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tA }
		};
		first[rhsB2] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tB }
		};
		first[rhsC1] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tA }
		};
		first[rhsC2] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tE }
		};
		first[rhsD1] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tC }
		};
		first[rhsE1] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tC }
		};
		first[rhsE2] = ext::set < ext::vector < DefaultSymbolType > > {
			ext::vector < DefaultSymbolType > ( )
		};
		first[rhsF1] = ext::set < ext::vector < DefaultSymbolType > > {
			{ tE }
		};
		first[rhsF2] = ext::set < ext::vector < DefaultSymbolType > > {
			ext::vector < DefaultSymbolType > ( )
		};
		// --------------------------------------------------

		ext::map < ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > >, ext::set < ext::vector < DefaultSymbolType > > > firstAlgo;

		for ( const auto & rule : grammar::RawRules::getRawRules ( grammar ) )
			for ( const auto & rhs : rule.second )
				firstAlgo[rhs] = grammar::parsing::First::first ( grammar, rhs );

		CHECK ( first == firstAlgo );
	}

	SECTION ( "Test 3" ) {
		DefaultSymbolType A = DefaultSymbolType ( 'A' );
		DefaultSymbolType c = DefaultSymbolType ( 'c' );
		DefaultSymbolType d = DefaultSymbolType ( 'd' );

		grammar::CFG < > grammar ( A );

		grammar.setTerminalAlphabet ( { c, d } );
		grammar.setNonterminalAlphabet ( { A } );
		grammar.setInitialSymbol ( A );

		grammar.addRule ( A, ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > { A, c } );
		grammar.addRule ( A, ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > { d } );

		ext::map < ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > >, ext::set < ext::vector < DefaultSymbolType > > > res = { { { d }, { { d } } }, { { A, c }, { { d } } } };
		CHECK ( res == grammar::parsing::First::first ( grammar ) );
	}

	SECTION ( "Test 4" ) {
		DefaultSymbolType S = DefaultSymbolType ( 'S' );
		DefaultSymbolType A = DefaultSymbolType ( 'A' );
		DefaultSymbolType B = DefaultSymbolType ( 'B' );
		DefaultSymbolType a = DefaultSymbolType ( 'a' );
		DefaultSymbolType b = DefaultSymbolType ( 'b' );
		DefaultSymbolType c = DefaultSymbolType ( 'c' );
		DefaultSymbolType d = DefaultSymbolType ( 'd' );
		DefaultSymbolType f = DefaultSymbolType ( 'f' );

		grammar::CFG < > grammar ( S );

		grammar.setTerminalAlphabet ( { a, b, c, d, f } );
		grammar.setNonterminalAlphabet ( { S, A, B } );
		grammar.setInitialSymbol ( S );

		grammar.addRule ( S, ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > { A, a } );
		grammar.addRule ( S, ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > { b, S } );
		grammar.addRule ( A, ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > { c, A, d } );
		grammar.addRule ( A, ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > { B } );
		grammar.addRule ( B, ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > { f, S } );
		grammar.addRule ( B, ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > > { } );

		ext::map < ext::vector < ext::variant < DefaultSymbolType, DefaultSymbolType > >, ext::set < ext::vector < DefaultSymbolType > > > res =
		{
			{ { A, a }, { { c }, { f }, { a } } }, { { b, S }, { { b } } }, { { c, A, d }, { { c } } }, { { B }, { { f }, ext::vector < DefaultSymbolType > ( ) } }, { { f, S }, { { f } } }, { { }, { ext::vector < DefaultSymbolType > ( ) } }
		};
		CHECK ( res == grammar::parsing::First::first ( grammar ) );
	}
}
