#include <catch2/catch.hpp>

#include <grammar/generate/GenerateUpToLength.h>

#include <grammar/ContextFree/EpsilonFreeCFG.h>
#include <string/LinearString.h>
#include <string/string/LinearString.h>

#include <factory/StringDataFactory.hpp>

static unsigned countStrings ( const ext::trie < DefaultSymbolType, bool > & node ) {
	unsigned res = node.getData ( );

	for ( const auto & child : node.getChildren ( ) )
		res += countStrings ( child.second );

	return res;
}

TEST_CASE ( "Generate Up To", "[unit][algo][grammar][generate]" ) {
	SECTION ( "Test eps-free CFG" ) {
		{
			DefaultSymbolType S = DefaultSymbolType("S");
			DefaultSymbolType A = DefaultSymbolType("A");
			DefaultSymbolType B = DefaultSymbolType("B");
			DefaultSymbolType C = DefaultSymbolType("C");
			DefaultSymbolType D = DefaultSymbolType("D");

			DefaultSymbolType a = DefaultSymbolType('a');
			DefaultSymbolType b = DefaultSymbolType('b');

			grammar::EpsilonFreeCFG < > grammar1(S);
			grammar1.setNonterminalAlphabet({S, A, B, C, D});
			grammar1.setTerminalAlphabet({a, b});
			grammar1.setGeneratesEpsilon(false);

			ext::set<string::LinearString < >> strings;

			ext::trie < DefaultSymbolType, bool > generated = grammar::generate::GenerateUpToLength::generate(grammar1, 5);

			std::ostringstream oss;
			generated.nicePrint ( oss );
			CAPTURE ( oss.str ( ) );

			for(const string::LinearString < > & str : strings) {
				bool flag = true;
				ext::trie < DefaultSymbolType, bool > * node = & generated;
				for ( const DefaultSymbolType & symbol : str.getContent ( ) ) {
					auto iter = node->getChildren ( ).find ( symbol );
					if ( iter == node->getChildren ( ).end ( ) ) {
						flag = false;
						break;
					}
					node = & iter->second;
				}

				CAPTURE ( factory::StringDataFactory::toString ( str ) );
				CHECK ( ( flag && node->getData ( ) == true ) );
			}

			CHECK (strings.size ( ) == countStrings ( generated ) );
		}
		{
			DefaultSymbolType S = DefaultSymbolType("S");
			DefaultSymbolType A = DefaultSymbolType("A");
			DefaultSymbolType B = DefaultSymbolType("B");
			DefaultSymbolType C = DefaultSymbolType("C");
			DefaultSymbolType D = DefaultSymbolType("D");

			DefaultSymbolType a = DefaultSymbolType('a');
			DefaultSymbolType b = DefaultSymbolType('b');
			DefaultSymbolType c = DefaultSymbolType('c');
			DefaultSymbolType d = DefaultSymbolType('d');

			grammar::EpsilonFreeCFG < > grammar1(S);
			grammar1.setNonterminalAlphabet({S, A, B, C, D});
			grammar1.setTerminalAlphabet({a, b, c, d});
			grammar1.addRule(S, {A, D});
			grammar1.addRule(S, {A});
			grammar1.addRule(S, {D});
			grammar1.addRule(A, {a, A, d});
			grammar1.addRule(A, {a, d});
			grammar1.addRule(A, {B, C});
			grammar1.addRule(A, {B});
			grammar1.addRule(A, {C});
			grammar1.addRule(B, {b, B, c});
			grammar1.addRule(B, {b, c});
			grammar1.addRule(C, {c, C});
			grammar1.addRule(C, {c});
			grammar1.addRule(D, {d, D});
			grammar1.addRule(D, {d});
			grammar1.setGeneratesEpsilon(true);

			ext::set<string::LinearString < >> strings;
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{c}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{c, c}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{c, c, c}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{d, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{d, d, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{a, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{a, d, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{b, c}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{b, c, c}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{c, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{c, c, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{c, d, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{a, c, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{b, c, d}));

			ext::trie < DefaultSymbolType, bool > generated = grammar::generate::GenerateUpToLength::generate(grammar1, 3);

			std::ostringstream oss;
			generated.nicePrint ( oss );
			CAPTURE ( oss.str ( ) );

			for(const string::LinearString < > & str : strings) {
				bool flag = true;
				ext::trie < DefaultSymbolType, bool > * node = & generated;
				for ( const DefaultSymbolType & symbol : str.getContent ( ) ) {
					auto iter = node->getChildren ( ).find ( symbol );
					if ( iter == node->getChildren ( ).end ( ) ) {
						flag = false;
						break;
					}
					node = & iter->second;
				}

				CAPTURE ( factory::StringDataFactory::toString ( str ) );
				CHECK ( ( flag && node->getData ( ) == true ) );
			}

			CHECK (strings.size ( ) == countStrings ( generated ) );
		}
		{
			DefaultSymbolType S = DefaultSymbolType("S");
			DefaultSymbolType A = DefaultSymbolType("A");
			DefaultSymbolType B = DefaultSymbolType("B");
			DefaultSymbolType C = DefaultSymbolType("C");
			DefaultSymbolType D = DefaultSymbolType("D");

			DefaultSymbolType a = DefaultSymbolType('a');
			DefaultSymbolType b = DefaultSymbolType('b');
			DefaultSymbolType c = DefaultSymbolType('c');
			DefaultSymbolType d = DefaultSymbolType('d');

			grammar::EpsilonFreeCFG < > grammar1(S);
			grammar1.setNonterminalAlphabet({S, A, B, C, D});
			grammar1.setTerminalAlphabet({a, b, c, d});
			grammar1.addRule(S, {A, D});
			grammar1.addRule(S, {D});
			grammar1.addRule(A, {a, A, d});
			grammar1.addRule(A, {a, d});
			grammar1.addRule(A, {B, C});
			grammar1.addRule(A, {B});
			grammar1.addRule(A, {C});
			grammar1.addRule(B, {b, B, c});
			grammar1.addRule(B, {b, c});
			grammar1.addRule(C, {c, C});
			grammar1.addRule(C, {c});
			grammar1.addRule(D, {d, D});
			grammar1.addRule(D, {d});

			ext::set<string::LinearString < >> strings;
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{d, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{d, d, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{a, d, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{c, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{c, c, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{c, d, d}));
			strings.insert(string::LinearString < >(ext::vector<DefaultSymbolType>{b, c, d}));

			ext::trie < DefaultSymbolType, bool > generated = grammar::generate::GenerateUpToLength::generate(grammar1, 3);

			std::ostringstream oss;
			generated.nicePrint ( oss );
			CAPTURE ( oss.str ( ) );

			for(const string::LinearString < > & str : strings) {
				bool flag = true;
				ext::trie < DefaultSymbolType, bool > * node = & generated;
				for ( const DefaultSymbolType & symbol : str.getContent ( ) ) {
					auto iter = node->getChildren ( ).find ( symbol );
					if ( iter == node->getChildren ( ).end ( ) ) {
						flag = false;
						break;
					}
					node = & iter->second;
				}

				CAPTURE ( factory::StringDataFactory::toString ( str ) );
				CHECK ( ( flag && node->getData ( ) == true ) );
			}

			CHECK(strings.size ( ) == countStrings ( generated ) );
		}
	}
}
