#!/usr/bin/env python
# arima_mapping 0.0.1
# Generated by dx-app-wizard.
#
# Basic execution pattern: Your app will run on a single machine from
# beginning to end.
#
# See https://wiki.dnanexus.com/Developer-Portal for documentation and
# tutorials on how to modify this file.
#
# DNAnexus Python Bindings (dxpy) documentation:
#   http://autodoc.dnanexus.com/bindings/python/current/

import os
import dx_utils
import re
import multiprocessing

import dxpy

@dxpy.entry_point('run_bwa')
def run_bwa(reads, ref_genome):
    dx_utils.download_and_gunzip_file(ref_genome)
    bwt_filename = dx_utils.run_cmd('ls *.bwt', returnOutput=True).strip()
    genome_prefix = os.path.splitext(bwt_filename)[0]

    reads = dx_utils.download_and_gunzip_file(reads, skip_decompress=True, create_named_pipe=True)
    ofn = re.sub('.fastq[.gz]*$', '.bam', reads)   

    cmd = 'set -e -o pipefail; bwa mem -t {0} -B 8 {1} {2} '.format(multiprocessing.cpu_count(), genome_prefix, reads)
    cmd += '| perl ./filter_five_end.pl | samtools view -@{0} -Sb - > {1} '.format(multiprocessing.cpu_count(), ofn)
    dx_utils.run_cmd(cmd)

    return {'output_bam': dxpy.dxlink(dxpy.upload_local_file(ofn))}

@dxpy.entry_point('combine_bams')
def combine_bams(fwd_bam, rev_bam):
    fwd_bam = dx_utils.download_and_gunzip_file(fwd_bam, create_named_pipe=True)
    rev_bam = dx_utils.download_and_gunzip_file(rev_bam, create_named_pipe=True)

    ofn = re.sub('.bam$', '.combined.bam', fwd_bam)

    #cmd = 'perl two_read_bam_combiner.pl {0} {1} | samtools view -@{2} -Sb - | samtools sort -o {3} - '
    cmd = 'set -e -o pipefail; perl two_read_bam_combiner.pl {0} {1} | /opt/biobambam2/bin/bamsormadup inputformat=sam indexfilename={2}.bai > {2} '
    cmd = cmd.format(fwd_bam, rev_bam, ofn)
    dx_utils.run_cmd(cmd)

    return {'output_bam': dxpy.dxlink(dxpy.upload_local_file(ofn, name=fwd_bam)),
            'output_bai': dxpy.dxlink(dxpy.upload_local_file(ofn + '.bai', name=fwd_bam + '.bai'))}


@dxpy.entry_point('main')
def main(fwd_reads, rev_reads, ref_genome):
    fwd_bam = dxpy.new_dxjob({'reads': fwd_reads, 'ref_genome': ref_genome}, 'run_bwa', name='Map forward reads').get_output_ref('output_bam')
    rev_bam = dxpy.new_dxjob({'reads': rev_reads, 'ref_genome': ref_genome}, 'run_bwa', name='Map reverse reads').get_output_ref('output_bam')

    merge_job = dxpy.new_dxjob({'fwd_bam': fwd_bam, 'rev_bam': rev_bam}, 'combine_bams', name='Combine bams')

    output = {'output_bam': merge_job.get_output_ref('output_bam'),
              'output_bai': merge_job.get_output_ref('output_bai')}

    return output

dxpy.run()
