from trie import Trie

def print_grid(grid):
    for i in grid:
        output = '   '.join(i)
        print("\t", output)

def find_strings(grid, words):
    trie_for_words = Trie()
    result = []
    # Inserting strings in the dictionary
    for word in words:
        trie_for_words.insert(word)
    # Calling dfs for all the cells in the grid
    for j in range(len(grid)):
        for i in range(len(grid[0])):
            dfs(trie_for_words, trie_for_words.root, grid, j, i, result)       
    return result

def dfs(words_trie, node, grid, row, col, result, word=''):
    # Checking if we found the string
    if node.is_string:
        result.append(word)
        node.is_string = False
        # remove the characters in the word that are not shared
        words_trie.remove_characters(word)
    
    if 0 <= row < len(grid) and 0 <= col < len(grid[0]):
        char = grid[row][col]
        # Getting child node of current node from Trie
        child = node.children.get(char)
        # if child node exists in Trie
        if child is not None:
            word += char
            # Marking it as visited before exploration
            grid[row][col] = None
            # Recursively calling DFS to search in all four directions
            for row_offset, col_offset in [(0, 1), (1, 0), (0, -1), (-1, 0)]:
                dfs(words_trie, child, grid, row + row_offset, col + col_offset, result, word)

            # Restoring state after exploration
            grid[row][col] = char
        

# Driver Code
def main():
    test_case_grid = [
        [['B', 'S', 'L', 'I', 'M'], 
        ['R', 'I', 'L', 'M', 'O'], 
        ['O', 'L', 'I', 'E', 'O'], 
        ['R', 'Y', 'I', 'L', 'N'], 
        ['B', 'U', 'N', 'E', 'C']],

        [['C', 'S', 'L', 'I', 'M'], 
        ['O', 'I', 'B', 'M', 'O'], 
        ['O', 'L', 'U', 'E', 'O'], 
        ['N', 'L', 'Y', 'S', 'N'], 
        ['S', 'I', 'N', 'E', 'C']],

        [['C', 'O', 'L', 'I', 'M'], 
        ['I', 'N', 'L', 'M', 'O'], 
        ['A', 'L', 'I', 'E', 'O'], 
        ['R', 'T', 'A', 'S', 'N'], 
        ['S', 'I', 'T', 'A', 'C']],

        [['P', 'S', 'L', 'A', 'M'], 
        ['O', 'P', 'U', 'R', 'O'], 
        ['O', 'L', 'I', 'E', 'O'], 
        ['R', 'T', 'A', 'S', 'N'], 
        ['S', 'I', 'T', 'A', 'C']],

        [['O', 'A', 'A', 'N'],
        ['E', 'T', 'A', 'E'],
        ['I', 'H', 'K', 'R'],
        ['I', 'F', 'L', 'V']],

        [['S', 'T', 'R', 'A', 'C'],
        ['I', 'R', 'E', 'E', 'E'],
        ['N', 'G', 'I', 'T', 'C'],
        ['I', 'T', 'S', 'R', 'A']],

        [['A', 'A', 'A'],
        ['A', 'A', 'A'],
        ['A', 'A', 'A']]
    ]

    strings_to_search = [
        ["BUY", "SLICK", "SLIME", "ONLINE", "NOW"],
        ["BUY", "STUFF", "ONLINE", "NOW"],
        ["REINDEER", "IN", "RAIN"],
        ["TOURISM", "DESTINY", "POPULAR"],
        ["OATH", "PEA", "EAT", "RAIN"],
        ["STREET", "STREETCAR", "STRING", "STING", "RING", "RACECAR"],
        ["A", "AA", "AAA", "AAAA"]
    ]

    for i in range(len(test_case_grid)):
            if i > 0:
                print()
            print(i + 1, ".\t 2D Grid:\n", sep="")
            print_grid(test_case_grid[i])
            print("\n\t Input list: ", strings_to_search[i])
            print("\n\t Output: ", find_strings(test_case_grid[i], strings_to_search[i]))
            print("-"*100) 


if __name__ == '__main__':
    main()