# 英文字母从'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)          # 用于存储每个模板的Id的数组
    
    # 将所有模板添加到计算机
    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)