def htseq_count(sortedBam,countFile,annotation,strand,annotationSource):
    """This function run htseq_count to count reads given bam file
    * sortedBam: str. Bamfile name
    * countFile: outputfilename
    * annotation: annotation file
    * outputpath: path to store the result files
    * annotation: source. 'ncbi','ensembl'
    """
    # 2. check the annotation source
    if annotationSource == 'ncbi':
        seqType = 'exon'
        id_attr = 'gene'
    elif annotationSource == 'ensembl':
        seqType = 'exon'
        id_attr = 'gene_id'
    elif annotationSource == 'genedb':
        seqType = 'CDS'
        id_attr = 'Parent'
    elif annotationSource == 'plasmodium':
        seqType = 'exon'
        id_attr = 'Parent'
    # 3. run htseq-count
    cmd = ('htseq-count -f bam -s {strand} -t {type} -i {gene} {bam} {annotation} > {output}').format(strand=strand,
         type=seqType,gene=id_attr,bam=sortedBam,annotation=annotation,output=countFile)#os.path.join(outpath,countFile))
    print(cmd);sys.stdout.flush()
    sarge.run(cmd)

def git_dir(dir_with_file):
    run('git init', cwd=str(dir_with_file))
    run('git config user.email "[email protected]"', cwd=str(dir_with_file))
    run('git config user.name "You"', cwd=str(dir_with_file))
    run('git add .', cwd=str(dir_with_file))
    run('git commit -m "Initial commit"', cwd=str(dir_with_file))
    return dir_with_file

def remove_cref(tex, cwd):
    """
    """
    sty_file = cwd+'/dynlearn.sty'

    # add poorman to cleveref import
    cref_regex1 = (r'usepackage\{cleveref\}', r'usepackage[poorman]{cleveref}')
    cref_regex2 = (r'usepackage\[(\w*)\]\{cleveref\}', r'usepackage[\g<1>,poorman]{cleveref}')

    with open(sty_file) as dynlearn:
        source = dynlearn.read()

    for exp, repl in [cref_regex1, cref_regex2]:
        source = re.sub(exp, repl, source)

    with open(sty_file, 'w') as dynlearn:
        dynlearn.write(source)

    # generate and run the sed script
    compile_tex(tex, cwd)

    sed_cmd = sarge.shell_format("sed -f {0}.sed {0}.tex > {0}.tex.temp", tex.split('.')[0])
    sarge.run(sed_cmd, cwd=cwd)
    sarge.run(sarge.shell_format("mv {0}.temp {0}", tex), cwd=cwd)

    # remove the cleveref import
    cref_regex3 = r'\\usepackage\[.*\]\{cleveref\}'
    comment_lines(sty_file, [cref_regex3])

def bwa_Db(db_path,ref_fa):
    """build bwa index"""
    if not os.path.exists(db_path):
        os.mkdir(db_path)
    cmd = ('bwa index -p {db_path}/bwa -a bwtsw {fa}').format(fa=ref_fa,db_path=db_path)
    print(cmd);sys.stdout.flush()
    sarge.run(cmd)

def Message(string,email):
    """
    This function send message to email when it run. 
    Used to calculate the time code runs.
    """
    cmd = ('echo {quote}|mailx -s "{string}" {email}').format(quote="",string=string,email=email)
    sarge.run(cmd)

def evm_cmd_list(out_fn,cmd_fn,evm,ref_fa,weight_fn,partition,gffs=['']):
    '''create cmd list for evm'''
    cmd = ('{evm} --genome {ref} --weights {w} {gffs} --output_file_name {out_fn} \
     --partitions {par} >  {cmd_l}').format(evm=evm,ref=ref_fa,
        w=weight_fn,gffs=' '.join(gffs),out_fn=out_fn,par=partition,cmd_l=cmd_fn)
    print(cmd)
    sarge.run(cmd)

 def test_capture_bytes(self):
     with Capture() as err:
         self.assertEqual(run("cat >&2", stderr=err, shell=True, input="bar").returncode, 0)
     self.assertEqual(err.bytes, b"bar")
     with Capture() as err:
         self.assertEqual(run("cat >&2", stderr=err, shell=True, input="bar").returncode, 0)
     self.assertEqual(err.text, "bar")

def run_cmd(cmd):
    try:
        print(cmd);sys.stdout.flush()
        sarge.run(cmd)
    except:
        print cmd,'error'
        assert False

def rnaseq_map_and_extract_by_chr(path,target_path,batch,rna_pipeline_file,rna_pipeline_param,chrom=''):
    os.chdir(path)
    folders = [f for f in os.listdir(path) if os.path.isdir(f)]
    folders = natsorted(folders)
    sub_folders = chunk(folders,batch)
    
    for sub_dir in sub_folders:
        # 1. copy files
        copy_files(path,target_path,sub_dir)
        # 2. map using STAR
        os.chdir(target_path)
        cmd = ('python {pipe} {pipe_param}').format(pipe=rna_pipeline_file,pipe_param=rna_pipeline_param)
        sarge.run(cmd)
        # 3. extract chromosome
        if chrom != '':
            bam_path = target_path + '/sortBam'
            os.chdir(bam_path)
            if not os.path.exists(bam_path+'/chr'): os.mkdir(bam_path+'/chr')
            bams = [f for f in os.listdir(bam_path) if f.endswith('.sort.bam')]
            for bam in bams:
                out = bam[5:]
                cmd = ('samtools view {input} {chr} > chr/{out}').format(input=bam,chr=chrom,out=out)
                print(cmd)
                sarge.run(cmd)
                os.remove(bam)
                os.remove(bam+'.bai')
            shutil.rmtree(target_path+'/bam')

def evm_partition(ref_fa,evm,gffs=[''],otherParams=['']):
    '''run evm to merge all the gff files'''
    cmd = ('{evm} --genome {ref} {gffs} {other} --segmentSize 50000000 \
     --overlapSize 10000 --partition_listing partitions_list.out').format(evm=evm,ref=ref_fa,
                                        gffs=' '.join(gffs),other=' '.join(otherParams))
    print(cmd)
    sarge.run(cmd)

def STAR(fastqFiles,outSamFile,db_path,thread=1,annotation='',otherParameters=['']):
    """STAR for single end read"""
    if annotation != '':
        otherParameters.extend(['--sjdbGTFfile {gff}'.format(gff=annotation)])
    if annotation.endswith('gff') or annotation.endswith('gff3'):
        otherParameters.append('--sjdbGTFtagExonParentTranscript Parent')
    # generate command
    if len(fastqFiles) == 1:
        starCmd = ('STAR --genomeDir {ref} --readFilesCommand zcat '
                     '--readFilesIn {fq1} --runThreadN {thread} '
                     '--outFileNamePrefix {output} --outSAMstrandField intronMotif '
                     '--outFilterIntronMotifs RemoveNoncanonical').format(
                    ref=db_path,fq1=fastqFiles[0],
                    thread=thread,output=outSamFile)
    elif len(fastqFiles) == 2:
        starCmd = ('STAR --genomeDir {ref} --readFilesCommand zcat '
                     '--readFilesIn {fq1} {fq2} --runThreadN {thread} '
                     '--outFileNamePrefix {output} --outSAMstrandField intronMotif '
                     '--outFilterIntronMotifs RemoveNoncanonical').format(
                    ref=db_path,fq1=fastqFiles[0],fq2=fastqFiles[1],
                    thread=thread,output=outSamFile)
    cmd = starCmd + ' ' + ' '.join(otherParameters)
    print(cmd);sys.stdout.flush()
    sarge.run(cmd)
    if 'SortedByCoordinate' in otherParameters:
        outFile = outSamFile+'Aligned.sortedByCoord.out.bam'
    else:
        outFile = outSamFile+'Aligned.out.bam'
    os.rename(outFile,outSamFile)
    if os.path.exists(outSamFile+'_STARgenome'):
        shutil.rmtree(outSamFile+'_STARgenome')

def sniffle(bam,outVCF,otherParameters=['']):
    """run sniffle to detect SV using pacbio"""
    cmd = ('sniffles -m {bam} -v {outVCF} ').format(bam=bam,outVCF=outVCF)
    if otherParameters != ['']:
        cmd = cmd + ' '.join(otherParameters)
    print(cmd);sys.stdout.flush()
    sarge.run(cmd)

def RNA_BaseRecalibrator4(realiBam,recalBam,gatk,table,ref_fa,gold_vcf,thread='1'):    
    '''Step 4 of base recalibration'''
    cmd = ('java -jar {gatk} -T PrintReads -R {ref_fa} '
           '-I {input} -BQSR {table} -o {output} -nct {thread}').format(gatk=gatk,
            ref_fa=ref_fa,input=realiBam,table=table,output=recalBam,thread=str(thread))
    print(cmd);sys.stdout.flush()
    sarge.run(cmd)

def RNA_BaseRecalibrator3(table,plot,post_table,gatk,ref_fa):
    '''Step 3 of base recalibration, compare the two tables'''
    cmd = ('java -jar {gatk} -T AnalyzeCovariates -R {ref_fa} '
           '-before {table} -after {post_table} -plots {output}').format(
            gatk=gatk,ref_fa=ref_fa,table=table,post_table=post_table,output=plot)
    print(cmd);sys.stdout.flush()
    sarge.run(cmd)

    def start_component(self, py_path, name, program_dir, on_keyboard_interrupt=None):
        """ Starts a component in background or foreground, depending on the 'fg' flag.
        """
        tmp_path = mkstemp('-zato-start-{}.txt'.format(name.replace(' ','')))[1]
        stdout_redirect = '' if self.args.fg else '1> /dev/null'
        stderr_redirect = '2> {}'.format(tmp_path)
        program = '{} -m {} {} {} {}'.format(get_executable(), py_path, program_dir, stdout_redirect, stderr_redirect)
        try:
            _stderr = _StdErr(
                tmp_path, stderr_sleep_fg if self.args.fg else stderr_sleep_bg)

            run(program, async=False if self.args.fg else True)

            # Wait a moment for any potential errors
            _err = _stderr.wait_for_error()
            if _err:
                self.logger.warn(_err)
                sys.exit(self.SYS_ERROR.FAILED_TO_START)

        except KeyboardInterrupt:
            if on_keyboard_interrupt:
                on_keyboard_interrupt()
            sys.exit(0)

        if self.show_output:
            if not self.args.fg and self.verbose:
                self.logger.debug('Zato {} `{}` starting in background'.format(name, self.component_dir))
            else:
                self.logger.info('OK')

def build_fa_dict(ref_fa,picard):
    '''build dictionary file for fa file '''
    out = '.'.join(ref_fa.split('.')[:-1]) + '.dict'
    cmd = ('java -jar {picard} CreateSequenceDictionary R={ref} O={out}').format(
            picard = picard,ref=ref_fa,out=out)
    print(cmd);sys.stdout.flush()
    sarge.run(cmd)

def main_evm(thread):
    os.chdir(evm_path)
    evm_gffs = ['--gene_predictions '+genemark_gff,'--transcript_alignments '+tr_gff,'--protein_alignments '+pr_gff]
    # 1. partition input
    evm_partition(ref_fa,evm+'/EvmUtils/partition_EVM_inputs.pl',evm_gffs)
    # 2. generate command lines
    evm_cmd_out = 'evm.out'
    cmd_fn = 'commands.list'
    evm_cmd_list(evm_cmd_out,cmd_fn,evm+'/EvmUtils/write_EVM_commands.pl',ref_fa,weight_fn,'partitions_list.out',evm_gffs)
    # 3. run commands
    pool = mp.Pool(processes=int(thread))
    cmds = open(cmd_fn).readlines()
    for cmd in cmds:
        pool.apply_async(run_cmd,args=(cmd,))
    pool.close()
    pool.join()
    # 4. combine results
    evm_combine = evm + '/EvmUtils/recombine_EVM_partial_outputs.pl'
    combine_partition(evm_combine,'partitions_list.out')
    # 5. transfer to gff
    to_gff = evm + '/EvmUtils/convert_EVM_outputs_to_GFF3.pl'
    cmd = ('{evm} --partitions partitions_list.out --output evm.out --genome {ref}').format(evm=to_gff,ref=ref_fa)
    sarge.run(cmd)
    # 6. merge gff
    fns = glob.glob('*/*.out.gff3')
    cmd = ('cat {input} > evm.merge.gff').format(input=' '.join(fns))
    sarge.run(cmd)
    # 7. extract genes supported by two algorithm
    filter_evm_gff(evm_path)

def Honey_tails(finalBam,bamTail,otherParams=['']):
    """This function run Honey tail,culster the soft clipped reads
    """
    cmd = ('Honey.py tails -o {out} {input} ').format(input=finalBam,out=bamTail)
    cmd = cmd + ' '.join(otherParams)
    print(cmd)
    sarge.run(cmd)

def exonerate(ref_fa,pr_fn,out_fn):
    '''map protein sequence to dna seq'''
    cmd = ('exonerate -m p2g -q {pr} -t {ref} --showalignment no \
    --showvulgar no --showtargetgff yes --minintron 20 --percent 50 \
    --score 100 --geneseed 250 -n 10 > {gff}').format(pr=pr_fn,ref=ref_fa,gff=out_fn)
    print(cmd)
    sarge.run(cmd)

def splitN(dedupBam,splitBam,gatk,ref_fa):
    '''This function splits reads due to wrong splicng by STAR'''
    cmd = ('java -jar {gatk} -T SplitNCigarReads -R {ref_fa} '
            '-I {input} -o {output} -rf ReassignOneMappingQuality '
            '-RMQF 255 -RMQT 60 -U ALLOW_N_CIGAR_READS').format(
            gatk=gatk,ref_fa=ref_fa,input=dedupBam,output=splitBam)
    print(cmd);sys.stdout.flush()
    sarge.run(cmd)