#!/usr/bin/env python

'''
Convert qiime labels to smile train labels
'''

import argparse, re, sys
from SmileTrain import util

def convert_qiime_label(fasta_entries):
    ''' ">samplename_123 lots of garbage stuff" -> ">samplename;123" '''
    for sid, sequence in fasta_entries:
        # split into "samplename" and "123"
        m = re.match("(.+)_(\d+)", sid)
        if m is None: raise RuntimeError("line %s did not parse" %(sid))
        sample_name, sequence_number = m.groups()
        
        # create a new sid
        new_sid = "sample=%s;%s" %(sample_name, sequence_number)
        
        yield new_sid, sequence
    

if __name__ == '__main__':
    # parse command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('input', help='input fasta')
    parser.add_argument('--output', '-o', default=sys.stdout, type=argparse.FileType('w'), help='output fasta (default stdout)')
    args = parser.parse_args()
    
    with open(args.input) as f:
        for sid, seq in convert_qiime_label(util.fasta_entries(f)):
            args.output.write(">%s\n%s\n" %(sid, seq))
        
@@ -62,16 +62,85 @@
 n labels
+%0A%0Ae.g., %22%3Esamplename_123 lots of garbage stuff%22 -%3E %22%3Esample=name;123%22
 %0A'''%0A%0Aim
@@ -162,16 +162,38 @@
 re, sys%0A
+from Bio import SeqIO%0A
 from Smi
@@ -237,29 +237,22 @@
 e_label(
-fasta_entries
+record
 ):%0A    '
@@ -323,52 +323,8 @@
 '''%0A
-    for sid, sequence in fasta_entries:%0A    
     
@@ -359,20 +359,16 @@
 d %22123%22%0A
-    
     m = 
@@ -394,17 +394,19 @@
 )%22, 
-sid)%0A    
+record.id)%0A
     
@@ -443,15 +443,22 @@
 or(%22
-line 
+record id '
 %25s
+'
  did
@@ -477,20 +477,16 @@
 %25(sid))%0A
-    
     samp
@@ -528,28 +528,24 @@
 ()%0A        %0A
-    
     # create
@@ -555,25 +555,51 @@
 new 
-s
 id
-%0A        new_s
+ and throw out the other junk%0A    record.
 id =
@@ -654,44 +654,58 @@
     
-    %0A        yield new_sid, sequence
+record.description = ''%0A        %0A    return record
 %0A   
@@ -833,21 +833,21 @@
 gument('
-input
+fasta
 ', help=
@@ -1034,163 +1034,122 @@
     
-with open(args.input) as f:%0A        for sid, seq in convert_qiime_label(util.fasta_entries(f)):%0A            args.output.write(%22%3E%25s%5Cn%25s%5Cn%22 %25(sid, seq))%0A        
+for record in SeqIO.parse(args.fasta, 'fasta'):%0A        SeqIO.write(convert_qiime_label(record), args.output, 'fasta')

qiime label: fixed