#!/usr/bin/env python

#############################################################################
#   dictReader.py
#   2015 James A. Stapleton
#
#   This program parses JSON-dumped barcode-hashed short-read dictionaries
#       created by barcodeHasher.py and sends files to SPAdes or Velvet
#       for assembly.
#
#   Options:
#       --makeHistogram: counts the number of reads in each  barcode-defined
#                           group, returns a file histogram.txt with one
#                           count per line that can be used to make a
#                           histogram and estimate the sample complexity.
#                           The constant MIN_NUMBER_OF_SEQUENCES defines the
#                           cutoff below which a bin is ignored.
#
#       --runVelvet: creates a fasta file from the reads in each barcode-
#                       defined group and sends it to Velvet for assembly.
#
#       --diginorm: performs digital normalization on each read group before
#                      sending it for assembly.
#
#       --runSpades: creates fasta or fastq files from the reads in each
#                       barcode-defined group and sends them to SPAdes for
#                       assembly.
#
#       --runTruSpades: creates fastq files from the reads in each
#                          barcode-defined group and sends them to truSPAdes
#                          for assembly.
#
#       --HPCC: calls SPAdes with one thread as required by some HPCCs.
#
#       --quality: creates fastq files rather than fasta files and uses
#                   SPAdes's built-in error correction mode.
#
#       --TRUNCATED_BARCODE_LENGTH: BARCODE_LENGTH - BARCODE_TRUNCATE from
#                                       the barcodeHasher.py run, default 16
#
#       --MIN_NUMBER_OF_SEQUENCES: minimum number of reads in a bin to count
#                                      in --makeHistogram
#
#       --threads: number of threads to use when running SPAdes, default 8,
#                  --HPCC sets to 1
#
#       --runCelera: calls the Celera assembler
#
#       --runTadpole: calls the Tadpole assember from bbmap
#
#############################################################################

from __future__ import division
from itertools import izip
import argparse
import subprocess
import os


def main(infile, makeHistogram, runVelvet, diginorm, runSpades, runTruSpades, HPCC,
        quality, TRUNCATED_BARCODE_LENGTH, MIN_NUMBER_OF_SEQUENCES, threads, runTadpole, runCelera):

    KMER_LENGTH = 99 # for diginorm
    MIN_CONTIG_LENGTH = 350 # for Velvet

    if not runVelvet and not makeHistogram and not runSpades and not runTruSpades and not runTadpole and not runCelera:
        print "Not doing anything!"
        return 0

    if makeHistogram:
        #histogram = collections.Counter()
        histoOut = open("histogram.txt", 'w')
        goodReads = 0
        totalReads = 0

    if (runVelvet + runSpades + runTruSpades):
        if os.path.exists('./contig_list.txt'):
            subprocess.call(['rm', 'contig_list.txt'])

    if runTruSpades:
        quality = 1
        if not os.path.exists('./truspades_input'):
            os.makedirs('./truspades_input')

    if runTadpole or runCelera:
        quality = 1

    with open(infile, 'r') as data:
        # parse JSON-dumped dictionary (uniqueDict.txt, unpairedDict.txt, etc)
        # Consider one barcode-defined group at a time
        seq_list = []
        for line in data:
            if line[-2] == '}':     # end of the file
                break
            elif line[0] == '{':    # start of the file
                continue
            # start of a barcode, example:     "AAAAAACGTTATGCAG": [
            elif line[-2] == '[':
                barcode = line[5:5+TRUNCATED_BARCODE_LENGTH]
            # start of a sequence, example:
            # "GGAAACTATACTAAAACTTGCTAAAAGCCATGATAAACTGAT",
            elif line[-3] == '"':
                sequence = line[9:-3]
                seq_list.append(sequence)
            # start of the last sequence in a group
            # no comma at the end, so -2 instead of -3
            elif line[-2] == '"':
                sequence = line[9:-2]
                seq_list.append(sequence)
            # end of a barcode-defined group, example:     ],
            elif line[-3] == ']' or line[-2] == ']':
                complete_list = seq_list
                seq_list = []

                if makeHistogram:
                    # this is the old way of printing a
                    # histogram of barcode abundance:
                    # returns pairs like 100, 20
                     # but this is hard to visualize
    #                 numBarcodes = len(complete_list)
    #                 histogram[numBarcodes] += 1
                    # new way: just print the number
                    # of sequences in the barcode
                    #  group to a file, plot a histogram with Python
                   # if len(complete_list) > 10 and len(complete_list) < 10000:
                    if quality:
                        if len(complete_list) > 2 * MIN_NUMBER_OF_SEQUENCES:
                            # divide by two here to get number of PE read pairs
                            print >> histoOut, len(complete_list)/4
                        totalReads += len(complete_list)/2
                        if len(complete_list) > 1000:
                            goodReads += 500
                        elif len(complete_list) > 200:
                            goodReads += len(complete_list)/2
                    else:
                        if len(complete_list) > MIN_NUMBER_OF_SEQUENCES:
                            # divide by two here to get number of PE read pairs
                            print >> histoOut, len(complete_list)/2
                        totalReads += len(complete_list)
                        if len(complete_list) > 500:
                            goodReads += 500
                        elif len(complete_list) > 100:
                            goodReads += len(complete_list)

                # write sequences from each items[barcode] to a fasta file,
                # send to assembler, overwrite fasta with next barcode sequnces
                # only run assembler if there are enough reads
                # for good coverage
                if (((not quality) and
                    (len(complete_list) > MIN_NUMBER_OF_SEQUENCES))
                    or (quality and
                        (len(complete_list) > 2 * MIN_NUMBER_OF_SEQUENCES))):

                    if runVelvet:
                            # open an output file that will be
                        # overwritten each time through
                        # the loop with a fasta list of sequences
                        # to feed to Velvet
                        with open('fasta.txt', 'w') as fasta:
                            i = 1
                            for seq in complete_list:
                                print >> fasta, ('>' + str(i))
                                if seq == '':
                                    fasta.write('A\n')
                                else:
                                    fasta.write(seq+'\n')
                                i += 1

                    # send fasta to Velveth
                    # Diginorm
                        if diginorm:
                            subprocess.call(["normalize-by-median.py",
                                             "-k", "20", "-C", "20", "-N", "4",
                                             "-x", "5e8", "fasta.txt"])
                            subprocess.call(["velveth", "seqData",
                                             str(KMER_LENGTH),
                                             "fasta.txt.keep"])
                        else:
                            subprocess.call(["velveth", "seqData",
                                             str(KMER_LENGTH), "-shortPaired",
                                             "fasta.txt"])

                        # get back files from Velveth
                        # send these files to Velvetg
                        subprocess.call(["velvetg", "seqData", "-exp_cov",
                                         "auto", "-cov_cutoff", 'auto',
                                         "-min_contig_lgth",
                                         str(MIN_CONTIG_LENGTH)])

                        # append contigs.fa to a growing file of contigs
                        with open("seqData/contigs.fa", "r") as fin:
                            data = fin.read()
                        with open("contig_list.txt", "a") as fout:
                            fout.write('>' + "Barcode: " + barcode + "\n")
                            fout.write(data + "\n")
                        with open("seqData/stats.txt", "r") as statsin:
                            stats = statsin.read()
                        with open("stats_list.txt", "a") as statsout:
                            statsout.write('>'+"Barcode: " + barcode + "\n")
                            statsout.write(stats + "\n")

                    if runSpades or runTadpole or runCelera: