#!/usr/bin/env python
"""
Converts BAM data to sorted SAM data.
usage: bam_to_sam.py [options]
   --input1: SAM file to be converted
   --output1: output dataset in bam format
"""

import optparse, os, sys, subprocess, tempfile, shutil, tarfile, random
#from galaxy import eggs
#import pkg_resources; pkg_resources.require( "bx-python" )
#from bx.cookbook import doc_optparse
#from galaxy import util

def stop_err( msg ):
    sys.stderr.write( '%s\n' % msg )
    sys.exit()
    
def toTar(tarFileName, samOutputNames):
    dir = os.path.dirname(tarFileName)    
    tfile = tarfile.open(tarFileName + ".tmp.tar", "w")
    currentPath = os.getcwd()
    os.chdir(dir)
    for file in samOutputNames:
        relativeFileName = os.path.basename(file)
        tfile.add(relativeFileName)
    os.system("mv %s %s" % (tarFileName + ".tmp.tar", tarFileName))
    tfile.close()
    os.chdir(currentPath)    


def __main__():
    #Parse Command Line
    parser = optparse.OptionParser()
    parser.add_option('-t', '--tar', dest='outputTar', default=None, help='output all SAM results in a tar file.' )
    parser.add_option( '', '--input1', dest='input1', help='The input list of BAM datasets on txt format.' )
    #parser.add_option( '', '--input1', dest='input1', help='The input BAM dataset' )
    parser.add_option( '', '--output1', dest='output1', help='The output list of SAM datasets on txt format.' )
    #parser.add_option( '', '--output1', dest='output1', help='The output SAM dataset' )
    parser.add_option( '', '--header', dest='header', action='store_true', default=False, help='Write SAM Header' )
    ( options, args ) = parser.parse_args()


    #Parse the input txt file and read a list of BAM files.
    file = open(options.input1, "r")
    lines = file.readlines()
    inputFileNames = []
    samOutputNames = []
    outputName = options.output1
    resDirName = os.path.dirname(outputName) + '/'
    #Write output txt file and define all output sam file names.
    out = open(outputName, "w")
    for line in lines:
        tab = line.split()
        inputFileNames.append(tab[1])
        samOutName = resDirName + tab[0] + '_samOutput_%s.sam' % random.randrange(0, 10000)
        samOutputNames.append(samOutName)
        out.write(tab[0] + '\t' + samOutName  + '\n')
    file.close()
    out.close()

    # output version # of tool
    try:
        tmp_files = []
        tmp = tempfile.NamedTemporaryFile().name
        tmp_files.append(tmp)
        tmp_stdout = open( tmp, 'wb' )
        proc = subprocess.Popen( args='samtools 2>&1', shell=True, stdout=tmp_stdout )
        tmp_stdout.close()
        returncode = proc.wait()
        stdout = None
        for line in open( tmp_stdout.name, 'rb' ):
            if line.lower().find( 'version' ) >= 0:
                stdout = line.strip()
                break
        if stdout:
            sys.stdout.write( 'Samtools %s\n' % stdout )
        else:
            raise Exception
    except:
        sys.stdout.write( 'Could not determine Samtools version\n' )



    tmp_dirs = []
    for i in range(len(inputFileNames)):
        try:
            # exit if input file empty
            if os.path.getsize( inputFileNames[i] ) == 0:
                raise Exception, 'Initial input txt file is empty.'
            # Sort alignments by leftmost coordinates. File <out.prefix>.bam will be created. This command
            # may also create temporary files <out.prefix>.%d.bam when the whole alignment cannot be fitted
            # into memory ( controlled by option -m ).
            tmp_dir = tempfile.mkdtemp()
            tmp_dirs.append(tmp_dir)
            tmp_sorted_aligns_file = tempfile.NamedTemporaryFile( dir=tmp_dir )
            tmp_sorted_aligns_file_base = tmp_sorted_aligns_file.name
            tmp_sorted_aligns_file_name = '%s.bam' % tmp_sorted_aligns_file.name
            tmp_files.append(tmp_sorted_aligns_file_name)
            tmp_sorted_aligns_file.close()
            
            command = 'samtools sort %s %s' % ( inputFileNames[i], tmp_sorted_aligns_file_base )
            tmp = tempfile.NamedTemporaryFile( dir=tmp_dir ).name
            tmp_stderr = open( tmp, 'wb' )
            proc = subprocess.Popen( args=command, shell=True, cwd=tmp_dir, stderr=tmp_stderr.fileno() )
            returncode = proc.wait()
            tmp_stderr.close()
            # get stderr, allowing for case where it's very large
            tmp_stderr = open( tmp, 'rb' )
            stderr = ''
            buffsize = 1048576
            try:
                while True:
                    stderr += tmp_stderr.read( buffsize )
                    if not stderr or len( stderr ) % buffsize != 0:
                        break
            except OverflowError:
                pass
            tmp_stderr.close()
            if returncode != 0:
                raise Exception, stderr
            # exit if sorted BAM file empty
            if os.path.getsize( tmp_sorted_aligns_file_name) == 0:
                raise Exception, 'Intermediate sorted BAM file empty'
        except Exception, e:
            stop_err( 'Error sorting alignments from (%s), %s' % ( inputFileNames[i], str( e ) ) )
            
        try:
            # Extract all alignments from the input BAM file to SAM format ( since no region is specified, all the alignments will be extracted ).
            if options.header:
                view_options = "-h"
            else:
                view_options = ""
            command = 'samtools view %s -o %s %s' % ( view_options, samOutputNames[i], tmp_sorted_aligns_file_name )
            tmp = tempfile.NamedTemporaryFile( dir=tmp_dir ).name
            tmp_stderr = open( tmp, 'wb' )
            proc = subprocess.Popen( args=command, shell=True, cwd=tmp_dir, stderr=tmp_stderr.fileno() )
            returncode = proc.wait()
            tmp_stderr.close()
            # get stderr, allowing for case where it's very large
            tmp_stderr = open( tmp, 'rb' )
            stderr = ''
            buffsize = 1048576
            try:
                while True:
                    stderr += tmp_stderr.read( buffsize )
                    if not stderr or len( stderr ) % buffsize != 0:
                        break
            except OverflowError:
                pass
            tmp_stderr.close()
            if returncode != 0:
                raise Exception, stderr
        except Exception, e:
            stop_err( 'Error extracting alignments from (%s), %s' % ( inputFileNames[i], str( e ) ) )
        if os.path.getsize( samOutputNames[i] ) > 0:
            sys.stdout.write( 'BAM file converted to SAM' )
        else:
            stop_err( 'The output file is empty, there may be an error with your input file.' )
     
    if options.outputTar != None:
        toTar(options.outputTar, samOutputNames)       
    #clean up temp files
    for tmp_dir in tmp_dirs:
        if os.path.exists( tmp_dir ):
            shutil.rmtree( tmp_dir )
    #print tmp_files
    #for tmp in tmp_files:
    #    os.remove(tmp)            
    

if __name__=="__main__": __main__()