import argparse
from subprocess import Popen, PIPE
from glob import glob
import os
from subprocess import call
import gzip
import logging
import json
import csv
import os
import subprocess
import zipfile


def fastqc_runner(input_dir,output_dir,numthreads = 1,batch = False):
    errors = []
    badfiles = []
    out = []
    file_path = '%s/*.fastq' % input_dir
    file_lst = []
    file_lst.extend(glob(file_path))
    files = ' '.join(file_lst)
    logging.info("start FASTQC on files: %s" % json.dumps(files))
    print 'files are', files
    # create qc dir
    qc_dir = output_dir + '/QC'
    if not os.path.exists(qc_dir):
        call('mkdir %s' % (qc_dir) ,shell=True)
    if batch == True:
        fastqc_cmd = 'fastqc -o %s --threads %s -b %s' % (qc_dir,numthreads,files)
        fastqc_run = Popen(fastqc_cmd,shell=True,stderr=PIPE,stdout=PIPE)
        fastqc_out,fastqc_err=fastqc_run.communicate()

        if fastqc_run.returncode!=0:
            errors.append(fastqc_err)
            for file in files.split():
                badfiles.append(os.pathbasename(file))
        out.append(fastqc_out)
    else:
        for file in files.split():
            print 'file is', os.path.basename(file)
            print('Working with' + file)
            fastqc_cmd = 'fastqc -o %s --noextract --threads 1 %s' % (qc_dir,file)
            logging.info("FASTQC: " + fastqc_cmd)
            fastqc_run = Popen(fastqc_cmd,shell=True,stderr=PIPE,stdout=PIPE)
            fastqc_out,fastqc_err=fastqc_run.communicate()
            if fastqc_run.returncode!=0:
                logging.info("FASTQC failed for: " + file)
            else:
                logging.info("FASTQC complete successfully for: " + file)
            badfiles.append(os.path.basename(file))
            errors.append(fastqc_err)
            out.append(fastqc_out)
    logging.info("end FASTQC")
    return badfiles, errors, out


def summaryCollector(output_dir):
    modules = ['Basic_Statistics',
           'Per_base_sequence_quality',
           'Per_tile_sequence_quality',
           'Per_sequence_quality_scores',
           'Per_base_sequence_content',
           'Per_sequence_GC_content',
           'Per_base_N_content',
           'Sequence_Length_Distribution',
           'Sequence_Duplication_Levels',
           'Overrepresented_sequences',
           'Adapter_Content', 'Kmer_Content']
    # Set dict to convert module results to integer scores:
    scores = {'pass': 1,
              'warn': 0,
              'fail': -1}
    #os.chdir(output_dir)
    # Get list of '_fastqc.zip' files generated by FastQC:
    files = [file for file in os.listdir(output_dir) if file.endswith('_fastqc.zip')]
    print(str(files)+'!!!!!')

    # List to collect module scores for each '_fastqc.zip' file:
    all_mod_scores = []

        # Read fastqc_data.txt file in each archive:
    for file in files:
        try:
            with  zipfile.ZipFile(os.path.join(output_dir,file), 'r') as archive:# open '_fastqc.zip' file 
                members = archive.namelist() # return list of archive members
                fname = [member for member in members if 'fastqc_data.txt' in member][0] # find 'fastqc_data.txt' in members
                print(fname)
                with archive.open(fname) as data: # open 'fastqc_data.txt'
                    # Get module scores for this file:
                    mod_scores = [file]
                    for line in data:
                       text = line.decode('utf-8')
                      # print(text)
                       if 'Total Sequences' in text:
                           print(text)
                           totalSeq = text.split()[-1]
                       if '>>' in text and '>>END' not in text:
                           #print(text)
                           text = text.lstrip('>>').split()
                           module = '_'.join(text[:-1])
                           result = text[-1]
                           mod_scores.append(scores[result])
                           # Append to all module scores list:
                           all_mod_scores.append(mod_scores)
        except:
            print file + " is bad and it should feel bad"

    # Write scores out to a CSV file:
    with open('all_mod_scores.csv', 'w') as f:
        writer = csv.writer(f)
        for mod_scores in all_mod_scores:
            writer.writerow(mod_scores)
    return int(totalSeq)





# Get current working directory:
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="options for running fastqc")
    parser.add_argument('-o','--output_dir',dest='outdir',type=str,help='where to write fastqc output')
    parser.add_argument('-i','--fastq_dir',dest='indir',type=str,help='directory where fastq files are locate')
    parser.add_argument('-t','--threads',dest='nthreads',type=int,default=1,help="num files to process simultaneously")
    parser.add_argument('-b','--batch',dest='batch',type=str,default=False,help="whether to run as a batch")
    opts = parser.parse_args()

    logging.basicConfig(filename='logfile.log',level=logging.DEBUG)

    badfiles,errors, out= fastqc_runner(opts.indir, opts.outdir,batch=opts.batch)
    if len(errors) != 0:
        for i in range(len(errors)):
            print '%s\n%s\n' % (badfiles[i],errors[i])
    summaryCollector(os.path.join(opts.outdir,'QC'))

