parse_ncbi_tax.jl

#=
Parse the NCBI dump data to a more regular format, with useless clades removed
etc.
The input should be the latest from https://ftp.ncbi.nih.gov/pub/taxonomy/

It creates a list of child/parent pairs, with the following changes from the raw data:

* Removes all descendants from known catch-all trash clades like "environmental samples"

* For any pair (child, parent) where the parent is not a canonical rank (i.e. one of
  the standard seven ranks used to e.g. GTDB), set the parent to the nearest canonical
  ancestor.
  E.g. if a genus has a subfamily parent, set its parent to the family instead.
  For this purpose, the universal ancestor of all life ("LUCA", tax number 1) is canonical.
  
* Remove all nodes which are not descendant from LUCA

* If a clade of a canonical rank has a parent which is not its real immediate parent
  (e.g. a genus level clade with an order level parent), add the missing clades as
  dummy clades.

* For nodes with multiple names, pick names according to a preference order, e.g.
  scientific name before historical names

* Make sure names are unique in the dataset - EXCEPT when two taxids actually point to the
  same clade. This can happen when two taxids are merged in later NCBI revisions.

* Also add in renamed (merged) nodes, as duplicates of the new node, but under the old taxid.
=#

if VERSION < v"1.11"
    macro main()
        return :main
    end
end

function (@main)(args)
    if length(args) != 4
        println(stderr, "Usage: julia parse_ncbi_tax.jl outfile names.dmp nodes.dmp merged.dmp")
        exit(1)
    end
    (outpath, namespath, nodespath, mergedpath) = args
    ispath(outpath) && error("Outpath exists: \"$(outpath)\"")
    for f in [namespath, nodespath, mergedpath]
        isfile(f) || error("No such file: \"$(f)\"")
    end

    fields = make_ncbi(namespath, nodespath, mergedpath)

    open(outpath, "w") do out
        println(out, "child_id\tchild_rank\tparent_id\tname")
        for i in fields
            println(out, join(i, '\t')::String)
        end
    end
    return
end

function make_ncbi(
        names_path::String,
        nodes_path::String,
        merged_path::String
    )::Vector{Tuple{Int, Rank, Int, String}}
    nodes = open(parse_nodes, nodes_path)
    parent_dict = build_parent_dict(nodes)

    # We remove these nodes, because they don't refer to real clades.
    to_remove = [
        57727, # environmental samples
    ] |> Set

    # Remove the nodes in to_remove and all their descendants
    parent_dict = remove_descendants(parent_dict, to_remove)

    # Get names, but only for those we haven't filtered away
    keep_ids = Set([c.id for c in keys(parent_dict)])
    names = open(names_path) do io
        get_names(io, keep_ids)
    end

    # Rename some important names
    # This is called "Bacteria <bacteria>" for some reason.
    names[2] = (NameTypes.scientific_name, "Bacteria")

    # Remove non-canonical ranks like infraorder.
    parent_dict = make_parent_canonical(parent_dict, RANK_TO_CANONICAL_INDEX, names)

    # Remove clades not directly descendant from the universal common ancestor
    parent_dict = remove_holey_descendants(parent_dict)

    merged_list = open(parse_merged, merged_path)
    id_to_node = Dict{Int, Node}()

    # We don't care about parents, because all nodes also appear as children, except id 1
    # which we assume is not deprecated because that would be weird.
    for child in keys(parent_dict)
        L = length(id_to_node)
        id_to_node[child.id] = child
        length(id_to_node) == L && error(lazy"Two nodes have id $(child.id)")
    end

    fields = [
        (child.id, child.rank, parent.id, names[child.id][2])
            for (child, parent) in parent_dict
    ]

    # Add deprecated (old merged nodes) to result.
    # We add it here and not in the parent dict, because we need to use the new
    # id to look up the name.
    for (; old, new) in merged_list
        child = get(id_to_node, new, nothing)
        child === nothing && continue
        parent = parent_dict[child]
        push!(fields, (old, child.rank, parent.id, names[child.id][2]))
    end

    # Make the output file compress better by sorting according to name
    sort!(fields; by = last)
    return fields
end

# All ranks in NCBI file
module Ranks
    names = [
        "acellular root",
        "biotype",
        "cellular root",
        "clade",
        "class",
        "cohort",
        "domain",
        "family",
        "forma",
        "forma specialis",
        "genotype",
        "genus",
        "infraclass",
        "infraorder",
        "isolate",
        "kingdom",
        "morph",
        "no rank",
        "order",
        "parvorder",
        "pathogroup",
        "phylum",
        "realm",
        "section",
        "series",
        "serogroup",
        "serotype",
        "species",
        "species group",
        "species subgroup",
        "strain",
        "subclass",
        "subcohort",
        "subfamily",
        "subgenus",
        "subkingdom",
        "suborder",
        "subphylum",
        "subsection",
        "subspecies",
        "subtribe",
        "subvariety",
        "superclass",
        "superfamily",
        "superorder",
        "superphylum",
        "tribe",
        "varietas",
    ]
    syms = map(i -> replace(i, " " => "_"), names)
    @eval @enum Rank::UInt8 $(Symbol.(syms)...)
    const STR_TO_ENUM = Dict(n => i for (n, i) in zip(names, instances(Rank)))
end

# All clade name types in NCBI
module NameTypes
    # Sorted in order from worst to best
    names = [
        "in-part",
        "acronym",
        "type material",
        "blast name",
        "includes",
        "synonym",
        "authority",
        "genbank acronym",
        "equivalent name",
        "genbank common name",
        "common name",
        "scientific name",
        "generated",
    ]
    syms = map(i -> replace(i, " " => "_", "-" => "_"), names)
    @eval @enum NameType::UInt8 $(Symbol.(syms)...)
    const STR_TO_ENUM = Dict(n => i for (n, i) in zip(names, instances(NameType)))
end

using .Ranks: Ranks, Rank
using .NameTypes: NameTypes, NameType

# These are the ranks we care to keep
const RANK_TO_CANONICAL_INDEX = Dict(
    Ranks.species => 1,
    Ranks.genus => 2,
    Ranks.family => 3,
    Ranks.order => 4,
    Ranks.class => 5,
    Ranks.phylum => 6,
    Ranks.domain => 7,
)

const PREFIXES = [c * '_' for c in "sgfocph"]

const INDEX_TO_CANONICAL_RANK = [
    Ranks.species,
    Ranks.genus,
    Ranks.family,
    Ranks.order,
    Ranks.class,
    Ranks.phylum,
    Ranks.domain,
]

const ID_COUNTER = Threads.Atomic{Int}(typemax(Int32))

const NAME_COUNTERS = fill(0, length(INDEX_TO_CANONICAL_RANK))
const NAME_COUNTERS_LOCK = ReentrantLock()

function create_new_taxon(rank_index::Integer)
    id = Threads.atomic_sub!(ID_COUNTER, 1)
    local counter
    @lock NAME_COUNTERS_LOCK begin
        counter = NAME_COUNTERS[rank_index]
        NAME_COUNTERS[rank_index] += 1
    end
    name = PREFIXES[rank_index] * "dummy_" * string(counter)
    return (id, name)
end

Base.tryparse(::Type{Rank}, s::AbstractString) = get(Ranks.STR_TO_ENUM, s, nothing)

function Base.parse(::Type{Rank}, s::AbstractString)
    r = tryparse(Rank, s)
    r === nothing && error(lazy"Could not parse as Rank: \"$(s)\"")
    return r
end

Base.tryparse(::Type{NameType}, s::AbstractString) = get(NameTypes.STR_TO_ENUM, s, nothing)

struct Node
    id::Int
    parent::Int
    rank::Rank
end

is_top(x::Node) = x.id == 1

"Remove all descendants of the nodes in `to_remove`"
function remove_descendants(
        parents::Dict{Node, Node},
        to_remove::Set{Int},
    )::Dict{Node, Node}
    result = empty(parents)
    for (child, parent) in parents
        should_add = true
        p = child
        while !is_top(p)
            if p.id ∈ to_remove
                should_add = false
                break
            end
            p = parents[p]
        end
        if should_add
            result[child] = parent
        end
    end
    @assert length(result) <= length(parents)
    return result
end

"""
    parse_nodes(io::IO)::Vector{Node}

Parse the NCBI nodes dump file, and return a vector of `Node`s.
"""
function parse_nodes(io::IO)::Vector{Node}
    result = Node[]
    for line in eachline(io)
        line = chopsuffix(rstrip(line), "\t|")
        isempty(line) && continue
        (id_str, parent_id_str, rankname) = split(line, "\t|\t")
        node = Node(parse(Int, id_str), parse(Int, parent_id_str), parse(Rank, rankname))
        push!(result, node)
    end
    return result
end

function build_parent_dict(nodes::Vector{Node})::Dict{Node, Node}
    parent_of = Dict{Node, Node}() # only top node has no parent
    node_by_id = Dict{Int, Node}()
    for node in nodes
        if get!(node_by_id, node.id, node) !== node
            error("Duplicate id")
        end
    end
    for node in nodes
        is_top(node) && continue
        parent_of[node] = node_by_id[node.parent]
    end
    return parent_of
end

"""
Change the parent to be the closest ancestor of a canonical clade.

Lots of the entries in the tax database are partially unknown, e.g.
it may be a genus which is the child of an unlabeled clade, which is
the child of a known family.
In this function, we remove all intermediate non-canonical clades, such that
every child's parent is either the top node, or a canonical clade.

Further, some canonical children have parents of the wrong rank, e.g.
a species whose direct parent is a family (skipping the genus rank).
We remove these.
"""
function make_parent_canonical(
        parent_of::Dict{Node, Node},
        rank_to_index::Dict{Rank, <:Integer},
        names::Dict{Int, Tuple{NameType, String}},
    )
    new_parents = empty(parent_of)
    for (child, parent) in parent_of
        new_parent = parent
        parent_index = get(rank_to_index, new_parent.rank, nothing)
        # Set parent to the closest canonical rank, e.g. if parent was infraorder before
        # set it to order.
        while !is_top(new_parent) && isnothing(parent_index)
            new_parent = parent_of[new_parent]
            parent_index = get(rank_to_index, new_parent.rank, nothing)
        end

        # Tax id 1 (the top) is "all life"
        parent_index = is_top(new_parent) ? 8 : parent_index
        child_index = get(rank_to_index, child.rank, nothing)

        # If child is not canonical, we accept whatever parent, since it's
        # not possible to skip a canonical rank, from the logic above
        if isnothing(child_index)
            @assert (child.rank == Ranks.no_rank || child.rank != new_parent.rank)
            new_parents[child] = new_parent
            continue
        end

        # If child is the true canonical child of the parent, just add it
        @assert child_index < parent_index
        if child_index + 1 == parent_index
            new_parents[child] = new_parent
            continue
        end

        # Else, we having missing ranks, which we create synthetically
        @assert child_index + 1 < parent_index
        ids = Int[]
        for index in (child_index + 1):(parent_index - 1)
            (id, name) = create_new_taxon(index)
            @assert !haskey(names, id)
            push!(ids, id)
            names[id] = (NameTypes.generated, name)
        end
        missings = Node[]
        for (i, index) in enumerate((child_index + 1):(parent_index - 1))
            parent_id = index == parent_index - 1 ? parent_index : ids[i + 1]
            rank = INDEX_TO_CANONICAL_RANK[index]
            node = Node(ids[i], parent_id, rank)
            push!(missings, node)
        end
        new_parents[child] = first(missings)
        new_parents[last(missings)] = new_parent
        for i in 1:(lastindex(missings) - 1)
            new_parents[missings[i]] = missings[i + 1]
        end
    end
    return new_parents
end


"""Remove any descendants that can't be traced directly to the top.
We do this by traversing from the top down, and only include reached nodes"""
function remove_holey_descendants(parent_of::Dict{Node, Node})
    new_parents = empty(parent_of)
    children_of = Dict{Node, Vector{Node}}()
    for (child, parent) in parent_of
        push!(get!(valtype(children_of), children_of, parent), child)
    end
    top = first(Iterators.filter(is_top, values(parent_of)))
    stack = Node[top]
    while !isempty(stack)
        parent = pop!(stack)
        for child in get(children_of, parent, Node[])
            new_parents[child] = parent
            push!(stack, child)
        end
    end
    return new_parents
end


"""
    get_names(io::IO, tax_ids::Union{Nothing, Set{Int}})::Dict{Int, Tuple{NameType, String}}

Parse the NCBI names dump fil, and return a dictionary mapping tax ids to names,
only for the taxids in `tax_ids`. If `tax_ids` is `nothing`, all tax ids are included.
"""
function get_names(
        io::IO,
        tax_ids::Union{Nothing, Set{Int}},
    )::Dict{Int, Tuple{NameType, String}}
    # Store all names for a given tax id (tax ids may have multiple names)
    # E.g. "human", "Homo sapiens" etc.
    names_of_tax = Dict{Int, Vector{Tuple{NameType, String}}}()
    for line in eachline(io)
        line = chopsuffix(rstrip(line), "\t|")
        # The NCBI file has two kinds of names: Unique names, and non-unique names.
        # For the latter, different clades may have the same name.
        # The unique name is empty if the name itself is already unique.
        (tax_str, non_unique_name, unique_name, class_str) = split(line, "\t|\t")
        name = isempty(unique_name) ? non_unique_name : unique_name
        tax_id = parse(Int, tax_str)
        # Only include tax ids the user asked for
        !isnothing(tax_ids) && tax_id ∉ tax_ids && continue
        # The name type is e.g. "common name" or "scientific name".
        # The NCBI file contains multiple rows with different names for the same clade,
        # so if there are multiple, we pick the kind of name we like the best.
        class = tryparse(NameType, class_str)
        isnothing(class) && error(
            lazy"Error: Could not parse name type $class_str for tax"
        )
        push!(
            get!(valtype(names_of_tax), names_of_tax, parse(Int, tax_str)),
            (class, String(name)),
        )
    end
    # Pick the kind of name we like the best, where multiple names are present.
    # E.g. we prefer "Opacifrons coxata" over "Opacifrons coxata (Stenhammar, 1854)".
    result = Dict(
        k::Int => argmax(((class, name),) -> Integer(class), v)::Tuple{NameType, String} for (k, v) in names_of_tax
    )
    # Assert names are unique
    names = Set{String}()
    for (_, (_, name)) in result
        if UInt8('\t') ∈ codeunits(name)
            error(lazy"Cannot have tab in name, but it is present in $(name)")
        end
        L = length(names)
        push!(names, name)
        length(names) == L && error(
            lazy"Error: Name \"$(name)\" is not unique"
        )
    end
    return result
end

"""
Parses the merged.dmp file, which contains a list of taxids which has been renamed, or merged with others.
Returns a list of old taxids and the new id it maps to
"""
function parse_merged(io::IO)::Vector{@NamedTuple{old::Int, new::Int}}
    # Each line looks like this: "12      |       74109   |" with whitespace being single tabs
    v = Vector{@NamedTuple{old::Int, new::Int}}()
    for (lineno, line) in enumerate(eachline(io))
        stripped = rstrip(line)
        isempty(stripped) && continue
        m = match(r"^([0-9]+)\t\|\t([0-9]+)\t\|$", stripped)
        m === nothing && error(
            "In merged file, on line $(lineno), expected line to fit regex " *
                "$(repr(r"^([0-9]+)\t\|\t([0-9]+)\t\|$")), but got $(repr(line))"
        )
        (old, new) = (something(m.captures[1]), something(m.captures[2]))
        push!(v, (; old = parse(Int, old), new = parse(Int, new)))
    end
    return v
end

if VERSION < v"1.11" && abspath(PROGRAM_FILE) == @__FILE__
    main(ARGS)
end