# Все допустимые символы английского алфавита от 'a' до 'z'
const VALID_CHARS = collect('a':'z')

mutable struct Node
    son::Vector{Int}   # номера дочерних вершин по каждому символу ('a'-'z')
    ans::Int           # количество совпадений, оканчивающихся здесь
    fail::Int          # ссылка на «предыдущий» узел с таким же суффиксом 
    du::Int            # число зависимых узлов, которые должны быть обработаны после текущего 
    idx::Int           # уникальный индекс для обнаруженного слова
end

Node() = Node(zeros(Int, length(VALID_CHARS)), 0, 1, 0, 0)

mutable struct ACAutomaton
    tr::Vector{Node}     # массив всех узлов дерева префиксного бора / автомата
    tot::Int             # общее количество использованных/существующих узлов
    final_ans::Vector{Int} # результат подсчета количества вхождений каждого шаблона
    pidx::Int            # счетчик уникальных индексов слов
end

ACAutomaton(max_nodes) = ACAutomaton([Node() for _ in 1:max_nodes], 1, Int[], 0)

function insert(ac::ACAutomaton, pattern)
    u = 1                      # начальный узел -- корень
    for ch in pattern          # проход по каждому символу шаблона
        char_code = findfirst(==(ch), VALID_CHARS)
        isnothing(char_code) && continue
        if ac.tr[u].son[char_code] == 0         # если такого дочернего узла нет
            ac.tot += 1                          # выделяем новый узел
            ac.tr[u].son[char_code] = ac.tot
        end
        u = ac.tr[u].son[char_code]              # переходим к следующему узлу
    end

    # Присваиваем паттерну уникальный идентификатор, если ещё не назначен
    if ac.tr[u].idx == 0
        ac.pidx += 1
        ac.tr[u].idx = ac.pidx
    end
    
    # возвращаем уникальный индекс
    return ac.tr[u].idx
end

function build(ac::ACAutomaton)
    q = [ac.tr[begin].son[i] for i in eachindex(VALID_CHARS) if ac.tr[begin].son[i] != 0]

    while !isempty(q)
        u = q[begin]
        q = q[begin+1:end]

        for i in eachindex(VALID_CHARS)
            son_node_idx = ac.tr[u].son[i]
            fail_node_idx = ac.tr[u].fail

            if son_node_idx != 0
                ac.tr[son_node_idx].fail = ac.tr[fail_node_idx].son[i]
                ac.tr[ac.tr[son_node_idx].fail].du += 1
                push!(q, son_node_idx)
            else
                ac.tr[u].son[i] = ac.tr[fail_node_idx].son[i] # делаем жадное продолжение
            end
        end
    end
end

function query(ac::ACAutomaton, text)
    u = 1
    for ch in text             # идём по символам текста
        char_code = findfirst(==(ch), VALID_CHARS)
        isnothing(char_code) && continue
        u = ac.tr[u].son[char_code]
        ac.tr[u].ans += 1       # увеличиваем счётчик совпадений, заканчивающихся в этой вершине
    end
end

function calculate_final_answers(ac::ACAutomaton)
    ac.final_ans = zeros(Int, ac.pidx + 1)
    
    # Начинаем с узлов, не имеющих зависимостей
    q = filter(i -> ac.tr[i].du == 0, 1:ac.tot)
    
    while !isempty(q)
        u = pop!(q)
        idx = ac.tr[u].idx
        if idx != 0
            ac.final_ans[idx] = ac.tr[u].ans
        end

        v = ac.tr[u].fail
        if v > 1
            ac.tr[v].ans += ac.tr[u].ans
            ac.tr[v].du -= 1
            ac.tr[v].du == 0 && push!(q, v)  # если теперь этот узел готов к обработке, добавляем в очередь
        end
    end
end

function test_aho_corasick(input, text)
    max_nodes, n = 200000 + 6, 5    # ограничение на максимальное количество узлов
    ac = ACAutomaton(max_nodes)      # создаем новый автомат
    ids = zeros(Int, n + 1)          # массив хранения идентификаторов каждого шаблона
    
    # Добавляем все шаблоны в автомат
    for i in 1:n
        pattern = input[i]
        ids[i+1] = insert(ac, pattern)
    end

    build(ac)                        # построили суффиксные ссылки
    query(ac, text)                  # выполнили поиск
    calculate_final_answers(ac)      # собрали ответы

    # Выводим результаты
    for i in 1:n
        uniqueid = ids[i]
        println("Number of instances of '$(input[i])': " => ac.final_ans[uniqueid+1])
    end
end

# Несколько тестовых шаблонов
input = ["a", "bb", "aa", "abaa", "abaaa"]   
text = "abaaabaa"               # Текст для анализа

test_aho_corasick(input, text)