from utils.protein_group_sql_db import ProteinGroupSqlDbInteractor


def get_tax_from_taxmaps(accs, taxmap_files, n_levels):
    if len(accs) == 0:
        return {}

    taxa = defaultdict(set)
    for taxmap_file in taxmap_files:
        with open(taxmap_file, 'r') as handle:
            fname = os.path.basename(taxmap_file)
            for line in handle:
                line = line.strip()
                fields = line.split("\t")
                acc = fields[0]
                tax = fields[-n_levels:]
                if acc in accs:
                    taxa[acc].add(fname + "\t" + "\t".join(tax))
    return dict(taxa)

rule tax_by_taxmap:
  input:
    user_taxmaps = expand('{fname}', fname=config['input']['taxmaps']),
    db_taxmaps = lambda w: config['taxmaps'] if 'taxmaps' in config else [],
    pg_db = 'pgs/protein_groups_db.sql',
    faa = 'seqs/all.faa'
  output:
    'tax/taxmap.txt',
    'seqs/missing_taxa.faa'
  message:
    "mapping taxa to accessions ..."
  version:
    "0.1"
  run:
    # these files are not always produced and can therefore not be listed under output
    missing_taxa = os.path.join(os.getcwd(), 'tax/missing_taxa_map.txt')
    ambig_taxa = os.path.join(os.getcwd(), 'tax/ambiguous_taxa.txt')
    #file reset
    for f in [missing_taxa, ambig_taxa]:
        if os.path.isfile(f):
            os.remove(f)

    #get accs
    accs = {}
    for orig_acc in ProteinGroupSqlDbInteractor(input.pg_db).get_all_accs_from_pg_db():
        acc = orig_acc.split("|")[-1]
        if acc in accs:
            sys.exit("error: accessions not unique after removing accession tags.")
        accs[acc] = orig_acc

    #screen taxmap files submitted by user
    taxa = get_tax_from_taxmaps(accs.keys(), input.user_taxmaps, len(config['taxlevel']))
    ambig = [x for x in accs.keys() if x in taxa and len(taxa[x]) > 1]
    if ambig:
        with open(ambig_taxa, "w") as handle:
            handle.write("\n".join(ambig))
            trace("error: ambiguous taxonomic information in submitted fasta files.", True)

    #screen built-in fasta files if necessary
    missing = [x for x in accs.keys() if x not in taxa]
    if missing:
        taxa = get_tax_from_taxmaps(missing, input.db_taxmaps, len(config['taxlevel']))
        ambig = [x for x in accs.keys() if x in taxa and len(taxa[x]) > 1]

        if ambig:
            with open(ambig_taxa, "w") as handle:
                handle.write("\n".join(ambig))
                trace("error: ambiguous taxonomic information in submitted fasta files.", True)

    #prepare fasta for taxblast
    missing = set([x for x in accs.keys() if x not in taxa])
    if missing:
        #if 'taxblast' not config['taxblast']:
        #    trace("error: missing taxonomic information in provided taxmaps (see " + config['output']['missing_taxa_map'] + " for accessions). Complement taxmaps or activate taxblast", True)
        #else:
        with open(missing_taxa, "w") as outhandle:
            outhandle.write("\n".join(missing))

        with open(output[1], "w") as outhandle, open(input.faa, "r") as inhandle:
            for line in inhandle:
                if len(line.strip()) == 0:
                    continue
                if line[0] == ">":
                    seq = inhandle.readline()
                    acc = line[1:].strip().split("|")[-1]
                    if acc in missing:
                        outhandle.write(line + seq)
                    elif acc not in taxa:
                        sys.exit("error: " + acc + " not found in " + input.faa)
    #write taxdump
    with open(output[0], "w") as handle:
        for acc, tax in taxa.items():
            if len(tax) > 0:
                handle.write(accs[acc] + "\t" + tax.pop() + "\n")
