#ifndef __TRIE_H
#define __TRIE_H

#include <cstdio>
#include <iostream>
#include <cstdlib>
#include <vector>
#include <string>
#include <stack>

int ALPHABET_SIZE = 26;

template <typename T>
class Trie
{
 private:
    int count;
    std::vector<Trie<T> *> children;
    T value;
    bool isEndOfWord;
  public:
    Trie() {}

    Trie(int capacity, T value)
      : count(0),
        children(capacity, NULL),
        value(value) { }

    Trie(int capacity)
      : count(0),
        children(capacity, NULL) { }

    ~Trie() {}
 	void insert(std::string key, T value) {
		// TODO 1 implementati functia de inserare in Trie
        Trie<T> *currentNode = this;
        for (int i = 0; i < key.size(); ++i) {
            if (currentNode->children[key[i] - 'a'] == NULL) {
                currentNode->children[key[i] - 'a'] = new Trie<int>(ALPHABET_SIZE);
            }
            currentNode = currentNode->children[key[i] - 'a'];
        }
        currentNode->count++;
        currentNode->value = value;
    }
	bool search(std::string key, T &val) {
 		// TODO 1 implementati functia de cautare in Trie
        Trie<T> *currentNode = this;
        for (int i = 0; i < key.size(); ++i) {
            if (currentNode->children[key[i] - 'a'] == NULL) {
                return false;
            }
            currentNode = currentNode->children[key[i] - 'a'];
        }
        if (currentNode->count == 0) {
            return false;
        }
        val = currentNode->value;
        return true;
    }

    bool remove(std::string key) {
    	// TODO 2 implementati functia de stergere din Trie
        Trie<T> *currentNode = this;
        for (int i = 0; i < key.size(); ++i) {
            if (currentNode->children[key[i] - 'a'] == NULL) {
                return false;
            }
            currentNode = currentNode->children[key[i] - 'a'];
        }
        currentNode->count = 0;
        currentNode->value = 0;
        return true;
    }

    int numWordsWithPrefix(std::string prefix) {
    	// TODO BONUS implementati gasirea numarului de cuvinte din Trie avand prefixul dat
        int ans = 0;
        Trie<T> *currentNode = this;
        for (int i = 0; i < prefix.size(); ++i) {
            if (currentNode->children[prefix[i] - 'a'] == NULL) {
                return 0;
            }
            currentNode = currentNode->children[prefix[i] - 'a'];
        }
        std::stack<Trie<T>* >stk;
        stk.push(currentNode);
        while(!stk.empty()) {
            currentNode = stk.top();
            stk.pop();
            ans += currentNode->count;
            for (int i = 0; i < ALPHABET_SIZE; ++i) {
                if (currentNode->children[i] != NULL) {
                    stk.push(currentNode->children[i]);
                }
            }
        }
        return ans;
    }

};
#endif