# All valid characters of the English alphabet from 'a' to 'z'
const VALID_CHARS = collect('a':'z')

mutable struct Node
    son::Vector{Int}   # child vertex numbers for each character ('a'-'z')
    ans::Int           # the number of matches ending here
    fail::Int          # a link to the "previous" node with the same suffix
    du::Int            # the number of dependent nodes to be processed after the current one
    idx::Int           # a unique index for the detected word
end

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

mutable struct ACAutomaton
    tr::Vector{Node}     # an array of all nodes of the prefix boron/ automaton tree
    tot::Int             # total number of used/existing nodes
    final_ans::Vector{Int} # the result of counting the number of occurrences of each pattern
    pidx::Int            # a counter for unique word indexes
end

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

function insert(ac::ACAutomaton, pattern)
    u = 1                      # the initial node is the root
    for ch in pattern          # passing through each character of the template
        char_code = findfirst(==(ch), VALID_CHARS)
        isnothing(char_code) && continue
        if ac.tr[u].son[char_code] == 0         # if there is no such child node
            ac.tot += 1                          # we select a new node
            ac.tr[u].son[char_code] = ac.tot
        end
        u = ac.tr[u].son[char_code]              # moving on to the next node
    end

    # We assign a unique identifier to the pattern, if it has not yet been assigned.
    if ac.tr[u].idx == 0
        ac.pidx += 1
        ac.tr[u].idx = ac.pidx
    end
    
    # returning a unique index
    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] # Making a greedy sequel
            end
        end
    end
end

function query(ac::ACAutomaton, text)
    u = 1
    for ch in text             # we follow the characters of the text
        char_code = findfirst(==(ch), VALID_CHARS)
        isnothing(char_code) && continue
        u = ac.tr[u].son[char_code]
        ac.tr[u].ans += 1       # increasing the count of matches ending at this vertex
    end
end

function calculate_final_answers(ac::ACAutomaton)
    ac.final_ans = zeros(Int, ac.pidx + 1)
    
    # We start with nodes that have no dependencies.
    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)  # if this node is now ready for processing, we add it to the queue.
        end
    end
end

function test_aho_corasick(input, text)
    max_nodes, n = 200000 + 6, 5    # limit on the maximum number of nodes
    ac = ACAutomaton(max_nodes)      # creating a new slot machine
    ids = zeros(Int, n + 1)          # an array for storing the IDs of each template
    
    # Adding all the templates to the machine
    for i in 1:n
        pattern = input[i]
        ids[i+1] = insert(ac, pattern)
    end

    build(ac)                        # we have built suffix links
    query(ac, text)                  # completed the search
    calculate_final_answers(ac)      # We have collected the answers

    # We display the results
    for i in 1:n
        uniqueid = ids[i]
        println("Number of instances of '$(input[i])': " => ac.final_ans[uniqueid+1])
    end
end

# Several test templates
input = ["a", "bb", "aa", "abaa", "abaaa"]   
text = "abaaabaa"               # Text for analysis

test_aho_corasick(input, text)