#!/usr/bin/env python3
"""Annotate a list of species with its full NCBI Taxonomy lineage.

    python3 species_lineage.py species.txt

Writes species_lineage.txt beside the input. Reads the NCBI Taxonomy dump
files directly and walks the tree itself, so it runs offline and needs
nothing beyond the standard library.

Get the dumps with:

    wget https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/taxdump.tar.gz
    tar -xzf taxdump.tar.gz nodes.dmp names.dmp

The sample nodes.dmp and names.dmp shipped alongside this script are a tiny
excerpt in the same format, enough to run the example without the 400 MB
download.
"""

import argparse
import os
import sys

FIELD_SEP = "\t|\t"
ROOT_TAXID = "1"


def read_dump(path, key_col, value_cols, keep_only=None):
    """Map one column of an NCBI .dmp file to one or more of the others.

    NCBI separates fields with "\t|\t" and ends each line with "\t|", so the
    final column carries a trailing separator that has to come off.
    """
    table = {}
    with open(path, encoding="utf-8") as handle:
        for line in handle:
            if keep_only is not None and keep_only not in line:
                continue
            fields = line.rstrip("\n").rstrip("\t|").split(FIELD_SEP)
            try:
                if len(value_cols) == 1:
                    table[fields[key_col]] = fields[value_cols[0]].upper()
                else:
                    table[fields[key_col]] = [fields[c].upper() for c in value_cols]
            except IndexError:
                # a short or malformed line is skipped rather than fatal
                continue
    return table


def invert(table):
    """Name -> taxid, from the taxid -> name mapping."""
    return {value: key for key, value in table.items()}


def walk_to_root(nodes, taxid):
    """Every taxid from this node up to the root, child first.

    Guards against a node that is its own parent and against a cycle in a
    hand-edited dump, either of which would otherwise loop forever.
    """
    lineage = [taxid]
    seen = {taxid}
    while True:
        entry = nodes.get(taxid)
        if not entry:
            break
        parent = entry[0]
        if parent == taxid or parent == ROOT_TAXID or parent in seen:
            break
        lineage.append(parent)
        seen.add(parent)
        taxid = parent
    return lineage


def lineage_string(taxids, names):
    """Root first, which is the order people expect to read a lineage in."""
    return "|".join(names.get(t, "?") for t in reversed(taxids))


def annotate(species_file, nodes, names, names_reversed, out_handle):
    written = 0
    for line in species_file:
        words = line.strip().split()
        if not words:
            continue
        species = " ".join(words).upper()

        if len(words) != 2 and "VIRUS" not in species:
            out_handle.write(
                species + "|Species name should consist of two words "
                "if it is not a virus\n")
            continue

        taxid = names_reversed.get(species)
        if taxid is None:
            out_handle.write(species + "|Not found — most likely a typo\n")
            continue

        out_handle.write(lineage_string(walk_to_root(nodes, taxid), names) + "\n")
        written += 1
    return written


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("species", help="text file, one species per line")
    parser.add_argument("--nodes", default="nodes.dmp")
    parser.add_argument("--names", default="names.dmp")
    parser.add_argument("-o", "--out", help="defaults to <input>_lineage.txt")
    args = parser.parse_args(argv)

    for path in (args.nodes, args.names, args.species):
        if not os.path.exists(path):
            parser.error(f"no such file: {path}")

    nodes = read_dump(args.nodes, 0, [1, 2])
    names = read_dump(args.names, 0, [1], keep_only="scientific name")
    names_reversed = invert(names)

    out_path = args.out or os.path.splitext(args.species)[0] + "_lineage.txt"
    with open(args.species, encoding="utf-8") as species_file, \
            open(out_path, "w", encoding="utf-8") as out_handle:
        written = annotate(species_file, nodes, names, names_reversed, out_handle)

    print(f"wrote {out_path} ({written} lineages resolved)", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())