from __future__ import print_function
import argparse
import operator

from ngskit.utils.dna import translate2aa
from ngskit.utils.fasta_tools import write_fasta_sequence

__author__ = "C. corbi-Verge"
__email__ = "carles.corbi@gmail.com"
__about__ = ''' This is a small script to count and rank reads from a
             fasta file. This script is agnostic, do not compare or map
             the reads with any source or library'''


def count_seqs_na(filename, seqs=dict()):
    '''.

    From a fasta file with na or aa sequences,
    count sequence and return a dict seq:counts

    '''
    with open(filename, 'r') as input_file:
        for line in input_file:
            if line.startswith('>'):
                pass
            else:
                seqs[line.strip()] = seqs.get(line.strip(), 0) + 1

    return seqs


def count_seqs_aa(filename, seqs=dict()):
    '''.

     From a fasta file with na seq,
    transform to aminoacids and count sequences,
    finally  return a dict seq:counts

    '''
    with open(filename, 'r') as input_file:
        for line in input_file:
            if line.startswith('>'):
                pass
            else:
                try:
                    aa = translate2aa(line.strip())

                except KeyError:
                    # so far ignore it found a N
                    # on the future add fuzzy logic to add read to a sequence
                    continue

                seqs[aa] = seqs.get(aa, 0) + 1
                # print >> fasta_2_logo,translate2aa(line.strip())
    # sorted_seqs = rank_reads(seqs)
    return seqs


def get_options():
    '''.

    '''
    parser = argparse.ArgumentParser(description="""
    Simple Reads counts


    Usage Demultiplexation:
    %prog -f [Fasta file]  -p count AA space  -n count in NA space -o outputname

    """)

    parser.add_argument('-p', help='Translate Oligos to Aminoacids and count',
                        dest='protein', default=False, action='store_true')
    #
    # parser.add_argument('-n', help='Count oligos', dest='dna', default=False,
    #                     action='store_true')

    parser.add_argument('-o', '--output', action="store", dest="output",
                        required=True, help='Output Name')

    parser.add_argument('-i', '--input_fasta', action="store", dest="input_fasta",
                        required=True, help='input fasta file from  demultiplex',
                        nargs='+')

    parser.add_argument('-f', '--output_fasta', action="store_true", dest="output_fasta",
                        default=False, help='Output as  "fasta",by default  "csv" file')

    options = parser.parse_args()

    return options


def write_output(sequences, fileouthandel,  otype ):

    for seq in sequences:

        if otype == "csv":
            print('{}\t{}'.format(seq[0], seq[1]), file=fileouthandel)
        elif otype == "fasta":
            write_fasta_sequence([ seq[0]+'_'+str(seq[1]), seq[0]], fileouthandel)

    fileouthandel.close()


if __name__ == '__main__':

    # setup
    opts = get_options()

    if opts.output_fasta:
        output_type = 'fasta'
    else:
        output_type = 'csv'

    if opts.protein:
        counter_fn = count_seqs_aa
        output_prefix = 'aa_'

    else:
        counter_fn = count_seqs_na
        output_prefix = 'na_'

    # do the job for each file
    counts = dict()
    for fname in opts.input_fasta:
        counts = counter_fn(fname, counts)

    out_handelr = open(output_prefix + opts.output, 'w')
    sorted_seqs = sorted(counts.items(), key=operator.itemgetter(1),
                         reverse=True)

    write_output(sequences=sorted_seqs, fileouthandel=out_handelr, otype=output_type)

@@ -460,32 +460,46 @@
 ame, seqs=dict()
+, min_lenght=0
 ):%0A    '''.%0A%0A   
@@ -743,32 +743,235 @@
           else:%0A
+                # do not count seq under a length, usefull to filter framshifts, etc%0A                if min_lenght and min_lenght %3C= len(line.strip()):%0A                    pass%0A                else:%0A    
                 
@@ -1079,16 +1079,30 @@
 s=dict()
+, min_lenght=0
 ):%0A    '
@@ -1653,17 +1653,209 @@
 ontinue%0A
-%0A
+                # do not count seq under a length, usefull to filter framshifts, etc%0A                if min_lenght and min_lenght %3C= len(aa):%0A                    pass%0A                else:%0A    
         
@@ -2409,141 +2409,8 @@
 ue')
-%0A    #%0A    # parser.add_argument('-n', help='Count oligos', dest='dna', default=False,%0A    #                     action='store_true')
 %0A%0A  
@@ -2778,28 +2778,29 @@
  '--output_f
-asta
+ormat
 ', action=%22s
@@ -2803,21 +2803,16 @@
 n=%22store
-_true
 %22, dest=
@@ -2816,28 +2816,29 @@
 st=%22output_f
-asta
+ormat
 %22,%0A         
@@ -2860,22 +2860,23 @@
 default=
-False,
+%22csv%22, 
  help='O
@@ -2916,16 +2916,247 @@
 v%22 file'
+, type=str,%0A                        choices= %5B'csv', 'fasta'%5D)%0A%0A    parser.add_argument('-m', '--min-length', action=%22store%22, dest=%22min_length%22,%0A                        type=int, default=0, help='Ignore sequence under this lenghth'
 )%0A%0A    o
@@ -3598,103 +3598,8 @@
 ()%0A%0A
-    if opts.output_fasta:%0A        output_type = 'fasta'%0A    else:%0A        output_type = 'csv'%0A%0A
     
@@ -3884,16 +3884,44 @@
 , counts
+, min_lenght=opts.min_lenght
 )%0A%0A    o
@@ -4158,17 +4158,24 @@
 ype=
+opts.
 output_
-type
+format
 )%0A

add minimum length feature