#!/usr/bin/env python3
#
# Convert fasta file from UNITE into
# (i) fasta file with only sequence names and sequences
# (ii) Taxonomy file with 7 levels D_0__ to D_6__
#
# By Jeanette Tångrot 2020-09-02

#--- Import libraries, do initializations  ---#
import sys
from Bio import SeqIO

usage = """create_unite_taxfile.py <unite.fa> <unite_fasta.fa> <unite_tax.txt>
  <unite.fa> : Input. Fasta file from UNITE containing taxonomies in description line
  <unite_fasta.fa> : Output. Name of fasta file.
  <unite_tax.txt> : Output. Name of text file with taxonomies.
"""

#--- Check and read arguments ---#
if len(sys.argv) != 4:
    exit("Usage: " + usage )

fasta_in = sys.argv[1]
fasta_out = sys.argv[2]
tax_out = sys.argv[3]

#--- Read sequence file and create new records ---#
replace_dict = {';p__': ';D_1__', ';c__': ';D_2__', ';o__': ';D_3__', ';f__': ';D_4__', ';g__': ';D_5__', ';s__': ';D_6__'}

fh_fasta = open( fasta_out, mode = 'w' )
fh_tax = open( tax_out, mode = 'w' )

for entry in SeqIO.parse( fasta_in, "fasta" ):
    (name, tax) = entry.id.split('|k__')
    tax = 'D_0__' + tax
    tax = tax.replace('unidentified','')
    for n1, n2 in replace_dict.items():
        tax = tax.replace( n1, n2 )
    tax = tax.replace('|SH','_SH')
    fh_fasta.write('>' + name + '\n' + str(entry.seq).upper() + '\n')
    fh_tax.write( name + '\t' + tax + '\n' )

fh_fasta.close()
fh_tax.close()

