#!/usr/bin/env python3
import sys
def parseNames(dmp):
    tax2name = dict()
    with open(dmp, 'r') as f:
        for line in f:
            fields = line.rstrip('\t|\n').split('\t|\t')
            if len(fields) != 4:
                print('skipped {}'.format(line))
                continue
            taxid = fields[0]
            name = fields[1]
            scientific = fields[3] == 'scientific name'
            if not taxid in tax2name:
                tax2name[taxid] = name
            elif scientific:
                tax2name[taxid] = name
    return tax2name


def parseACC(dmp, ref=None):
    if ref is not None:
        filterC = set()
        with open(ref, 'r') as f:
            for line in f:
                filterC.add(line.strip())
    acc2tax = dict()
    if dmp != '-':
        with open(dmp, 'r') as f:
            for line in f:
                fields = [x.strip() for x in line.split('\t')]
                acc = fields[1]
                if ref is not None:
                    if acc in filterC:
                        taxid = fields[2]
                        acc2tax[acc] = taxid
                else:
                    taxid = fields[2]
                    acc2tax[acc] = taxid
    else:
        f = sys.stdin
        for line in f:
            fields = [x.strip() for x in line.split('\t')]
            acc = fields[1]
            if ref is not None:
                if acc in filterC:
                    taxid = fields[2]
                    acc2tax[acc] = taxid
            else:
                taxid = fields[2]
                acc2tax[acc] = taxid
    return acc2tax


def generateMinimalCatalog(acc2taxid, taxid2name, outfile, ref=None):
    with open(outfile, 'w') as out:
        acc2tax = parseACC(acc2taxid, ref)
        tax2name = parseNames(taxid2name)
        for acc in sorted(acc2tax):
            try:
                out.write('{}\t{}\t{}\n'.format(acc, acc2tax[acc], tax2name[acc2tax[acc]]))
            except KeyError:
                continue

if __name__ == '__main__':
   import argparse
   parser = argparse.ArgumentParser(description='Generate catalog for altered MicrobeGPS')
   parser.add_argument('acc2taxid', type=str, help='nucl_gb.accession2taxid')
   parser.add_argument('taxid2name', type=str, help='names.dmp')
   parser.add_argument('out', type=str, help='Output location')
   parser.add_argument('-ref', type=str, default=None, help='File containing all accessions to be included in catalog')
   args = parser.parse_args()
   generateMinimalCatalog(args.acc2taxid, args.taxid2name, args.out, args.ref)
