#bamToGFF.py
#script to grab reads from a bam that align to a .gff file

import sys
from collections import defaultdict
from bamToGFFutils import *
import os
import string
from string import join,upper,maketrans
import re


def parseSamHeader(samFile):
    '''parses any sam type file with a 3 column tab del header'''
    samDict = {}
    sam = open(samFile,'r')

    for line in sam:
        if line[0] == '@':
            headerLine = line[:-1].split('\t')
            samDict[headerLine[1]] = headerLine[2]
        else:
            break
    sam.close()

    return samDict


def getUniquelyMappingReads(bamFile):
    '''
    scripts designed to return the total number of uniquely mapping sequence tags from a bam file,
    first by looking for a corresponding stats file in the folder, and second by manually computing the number.
    '''
    #a uniquely mapping sequence tag is defined by collapsing all tags that map to the exact same location with the exact same sequence
    #first try to extract the number from a stats file
    fullPath = os.path.abspath(bamFile)
    bamName = fullPath.split('/')[-1].split('.')[0]
    pathFolder = join(fullPath.split('/')[0:-1],'/')
    bamFiles = os.listdir(pathFolder)
    statsFile = filter(lambda x: x.count(bamName) ==1 and x.count('stats.') ==1,bamFiles)
    if len(statsFile) == 1:
        print('USING STATS FILE %s' % (pathFolder+'/'+statsFile[0]))
        samDict = parseSamHeader(pathFolder+'/'+statsFile[0])
        return int(samDict['UniquelyMappingSequenceTags'])
    elif len(statsFile) > 1:
        statsFile = filter(lambda x: x.count(bamName) ==1 and x.count('stats.concise') ==1,bamFiles)
        print('USING STATS FILE %s' % (pathFolder+'/'+statsFile[0]))
        samDict = parseSamHeader(pathFolder+'/'+statsFile[0])
        return int(samDict['UniquelyMappingSequenceTags'])
    else:
        print('no precomputed stats file found for %s.' % (bamFile))
        return None


def mapBamToGFF(bamFile,gff,sense = 'both',unique = 0,extension = 200,floor = 0,density = False,rpm = False,binSize = 25,clusterGram = None,matrix = None,raw = False,includeJxnReads = False):
    '''maps reads from a bam to a gff'''
    floor = int(floor)
    bam = Bam(bamFile)
    #if cluster is specified, override certain flags
    if clusterGram:
        density =True
        binSize = int(binSize)
    else:
        binSize = 25

    if matrix:
        density = True
    #new GFF to write to
    newGFF = []
    #millionMappedReads

    if float(unique) == 0.0:
        unique = False
        if rpm:    
            MMR= round(float(bam.getTotalReads('mapped'))/1000000,4)
        else:
            MMR = 1
    else:

        if rpm:
            MMR = float(unique)
        else:
            MMR = 1
        unique = True



    print('using a MMR value of %s' % (MMR))
    senseTrans = maketrans('-+.','+-+')
    if type(gff) == str:
        gff = parseTable(gff,'\t')

    #setting up a clustergram table
    if clusterGram:
        #first grab a header line
        line = gff[0]
        gffLocus = Locus(line[0],int(line[3]),int(line[4]),line[6],line[1])

        nBins = gffLocus.len()/binSize
        binSizeList = [nBins]

        #now go through each line of the gff and make sure they're all the same length
        for i in range(0,len(gff),1):
            line = gff[i]
            gffLocus = Locus(line[0],int(line[3]),int(line[4]),line[6],line[1])
            binSizeList.append(gffLocus.len()/binSize)
        binSizeList = uniquify(binSizeList)
        if len(binSizeList) > 1: 
            print('WARNING: lines in gff are of different length. Output clustergram will have variable row length')
        newGFF.append(['GENE_ID','locusLine'] + [str(x*binSize)+'_'+bamFile.split('/')[-1] for x in range(1,max(binSizeList)+1,1)])
    #setting up a maxtrix table
    if matrix:
        newGFF.append(['GENE_ID','locusLine'] + ['bin_'+str(n)+'_'+bamFile.split('/')[-1] for n in range(1,int(matrix)+1,1)])

    #getting and processing reads for gff lines
    ticker = 0
    print('Number lines processed')
    for line in gff:
        line = line[0:9]
        if ticker%100 == 0:
            print ticker
        ticker+=1
        gffLocus = Locus(line[0],int(line[3]),int(line[4]),line[6],line[1])
        searchLocus = makeSearchLocus(gffLocus,int(extension),int(extension))
        
        reads = bam.getReadsLocus(searchLocus,'both',unique,'none',includeJxnReads)
        #now extend the reads and make a list of extended reads
        extendedReads = []
        for locus in reads:
            if locus.sense() == '+' or locus.sense() == '.':
                locus = Locus(locus.chr(),locus.start(),locus.end()+extension,locus.sense(), locus.ID())
            if locus.sense() == '-':
                locus = Locus(locus.chr(),locus.start()-extension,locus.end(),locus.sense(),locus.ID())
            extendedReads.append(locus)
        if gffLocus.sense() == '+' or gffLocus.sense == '.':
            senseReads = filter(lambda x:x.sense() == '+' or x.sense() == '.',extendedReads)
            antiReads = filter(lambda x:x.sense() == '-',extendedReads)
        else:
            senseReads = filter(lambda x:x.sense() == '-' or x.sense() == '.',extendedReads)
            antiReads = filter(lambda x:x.sense() == '+',extendedReads)

        #at this point can output starts onto the GFF unless density is called
        if density:
            senseHash = defaultdict(int)
            antiHash = defaultdict(int)

            #filling in the readHashes             
            if sense == '+' or sense == 'both' or sense =='.':
                for read in senseReads:
                    for x in range(read.start(),read.end()+1,1):
                        senseHash[x]+=1
            if sense == '-' or sense == 'both' or sense == '.':
                for read in antiReads:
                    for x in range(read.start(),read.end()+1,1):
                        antiHash[x]+=1

            #now apply flooring and filtering for coordinates
            keys = uniquify(senseHash.keys()+antiHash.keys())
            if floor > 0:

                keys = filter(lambda x: (senseHash[x]+antiHash[x]) > floor,keys)
            #coordinate filtering
            keys = filter(lambda x: gffLocus.start() < x < gffLocus.end(),keys)

            if clusterGram or matrix:
                clusterLine = [gffLocus.ID(),gffLocus.__str__()]
                if matrix:
                    binSize = (gffLocus.len()-1)/int(matrix)
                    nBins = int(matrix)
                if clusterGram:
                    nBins = gffLocus.len()/binSize
                if binSize == 0:
                    clusterLine+=['NA']*int(matrix)
                    newGFF.append(clusterLine)
                    continue
                n=0
                if gffLocus.sense() == '+' or gffLocus.sense() =='.' or gffLocus.sense() == 'both':
                    i = gffLocus.start()
                    
                    while n <nBins:
                        n+=1
                        binKeys = filter(lambda x: i < x < i+binSize,keys)
                        binDen = float(sum([senseHash[x]+antiHash[x] for x in binKeys]))/binSize
                        clusterLine+=[round(binDen/MMR,4)]
                        i = i+binSize
                else:
                    i = gffLocus.end()
                    while n < nBins:
                        n+=1
                        binKeys = filter(lambda x: i-binSize < x < i,keys)
                        binDen = float(sum([senseHash[x]+antiHash[x] for x in binKeys]))/binSize
                        clusterLine+=[round(binDen/MMR,4)]
                        i = i-binSize
                newGFF.append(clusterLine)

            #for regular old density calculation
            else:
                senseTotalDen = float(sum([senseHash[x] for x in keys]))/gffLocus.len()
                antiTotalDen = float(sum([antiHash[x] for x in keys]))/gffLocus.len()
                if rpm:
                    senseTotalDen = senseTotalDen/MMR
                    antiTotalDen = antiTotalDen/MMR
                if sense == 'both' or sense == '.':
                    if gffLocus.sense() == '+' or gffLocus.sense() == '.':
                        readLine = '+'+':%s' % (round(senseTotalDen,4)) + ';' +'-' + ':%s' % (round(antiTotalDen,4))
                    else:
                        readLine = '+'+':%s' % (round(antiTotalDen,4)) + ';' +'-' + ':%s' % (round(senseTotalDen,4))
                elif sense == '+':
                    readLine = '+'+':%s' % (round(senseTotalDen,4))
                elif sense == '-':
                    readLine = '-' + ':%s' % (round(antiTotalDen,4))
                newGFF.append(line + [readLine])             
        #if not cluster or density simply return reads 
        elif raw:
            if sense == 'both' or sense == '.':
                if gffLocus.sense() == '+' or gffLocus.sense() == '.':
                    readLine = '+'+':'+ join([str(locus.start()) for locus in senseReads],',') +';' + '-'+':'+ join([str(locus.start()) for locus in antiReads],',')
                else:
                    readLine = '+'+':'+ join([str(locus.start()) for locus in antiReads],',')+';'+'-'+':'+ join([str(locus.start()) for locus in senseReads],',')
            elif sense == '+':
                readLine = gffLocus.sense()+':'+ join([str(locus.start()) for locus in senseReads],',')
            elif sense == '-':
                readLine = string.translate(gffLocus.sense(),senseTrans)+':'+ join([str(locus.start()) for locus in antiReads],',')
            newGFF.append(line+[readLine])
        #if not raw and not density gives total
        else:
            if sense == 'both' or sense == '.':
                readLine = str((len(antiReads) + len(senseReads))/MMR)   
            elif sense == '+':
                readLine = str(len(senseReads)/MMR)
            elif sense == '-':
                readLine = str(len(antiReads)/MMR)
            newGFF.append(line+[readLine])
            
    return newGFF

def convertEnrichedRegionsToGFF(enrichedRegionFile):
    '''converts a young lab enriched regions file into a gff'''
    newGFF = []
    enrichedRegions = open(enrichedRegionFile,'r')
    header = enrichedRegions.readline()
    i = 0
    for line in enrichedRegions:
        line = line[:-1].split('\t')
        newLine = ['chr'+line[0],'row_'+str(i),line[4],line[1],line[2],'','.','','row_'+str(i),'']
        newGFF.append(newLine)
        i+=1
    return newGFF

def main():
    from optparse import OptionParser
    usage = "usage: %prog [options] -b [SORTED BAMFILE] -i [INPUTFILE] -o [OUTPUTFILE]"
    parser = OptionParser(usage = usage)
    #required flags
    parser.add_option("-b","--bam", dest="bam",nargs = 1, default=None,
                      help = "Enter .bam file to be processed.")
    parser.add_option("-i","--input", dest="input",nargs = 1, default=None,
                      help = "Enter .gff or ENRICHED REGION file to be processed.")
    #output flag
    parser.add_option("-o","--output", dest="output",nargs = 1, default=None,
                      help = "Enter the output filename.")
    #additional options
    parser.add_option("-s","--sense", dest="sense",nargs = 1, default='both',
                      help = "Map to '+','-' or 'both' strands. Default maps to both.")
    parser.add_option("-u","--unique", dest="unique",nargs = 1, default=0,
                      help = "Takes only unique sequence tags (avoids pcr repeats). must provide number of million uniquely mapping reads")
    parser.add_option("-d","--density", dest="density",action='store_true', default=False,
                      help = "Calculates a read density for each region, returns a single value per region")
    parser.add_option("-f","--floor", dest="floor",nargs =1, default=0,
                      help = "Sets a read floor threshold necessary to count towards density")    
    parser.add_option("-e","--extension", dest="extension",nargs = 1, default=200,
                      help = "Extends reads by n bp. Default value is 200bp")
    parser.add_option("-r","--rpm", dest="rpm",action = 'store_true', default=False,
                      help = "Normalizes density to reads per million (rpm)")
    parser.add_option("-t","--total", dest="total",action = 'store_true', default=False,
                      help = "Gives the total read count in the region. Overrides density, floor, and rpm flags")
    parser.add_option("-c","--cluster", dest="cluster",nargs = 1, default=None,
                      help = "Outputs a fixed bin size clustergram. user must specify bin size.")
    parser.add_option("-m","--matrix", dest="matrix",nargs = 1, default=None,
                      help = "Outputs a variable bin sized matrix. User must specify number of bins.")
    parser.add_option("-j","--jxn", dest="jxn",action = 'store_true', default=False,
                      help = "if flagged, includes jxn reads")
    (options,args) = parser.parse_args()

    print(options)
    print(args)

    if options.bam:
        bamFile = options.bam
        fullPath = os.path.abspath(bamFile)
        bamName = fullPath.split('/')[-1].split('.')[0]
        pathFolder = join(fullPath.split('/')[0:-1],'/')
        fileList = os.listdir(pathFolder)
        hasBai = False
        for fileName in fileList:
            if fileName.count(bamName) == 1 and fileName.count('.bai') == 1:
                hasBai = True

        if not hasBai:
            print('ERROR: no associated .bai file found with bam. Must use a sorted bam with accompanying index file')
            parser.print_help()
            exit()

    if options.sense:
        if ['+','-','.','both'].count(options.sense) == 0:
            print('ERROR: sense flag must be followed by +,-,.,both')
            parser.print_help()
            exit()

    if options.cluster and options.matrix:
        print('ERROR: Cannot specify both matrix and clustergram flags.')
        parser.print_help()
        exit()

    if options.matrix:
        try:
            int(options.matrix)
        except:
            print('ERROR: User must specify an integer bin number for matrix (try 50)')
            parser.print_help()
            exit()
            
    if options.cluster:
        try:
            int(options.cluster)
        except:
            print('ERROR: User must specify an integer bin size for clustergram (try 25)')
            parser.print_help()
            exit()

    if options.input and options.bam:
        inputFile = options.input
        if inputFile.split('.')[-1] != 'gff':
            print('converting file to a .gff')
            gffFile = convertEnrichedRegionsToGFF(inputFile)
        else:
            gffFile = inputFile

        bamFile = options.bam

        if options.output == None:
            output = os.getcwd() + inputFile.split('/')[-1]+'.mapped'
        else:
            output = options.output
        if options.cluster:
            print('mapping to GFF and making clustergram with fixed bin width')
            newGFF = mapBamToGFF(bamFile,gffFile,options.sense,options.unique,int(options.extension),options.floor,options.density,options.rpm,options.cluster,True,None,False,options.jxn)
        elif options.matrix:
            print('mapping to GFF and making a matrix with fixed bin number')
            newGFF = mapBamToGFF(bamFile,gffFile,options.sense,options.unique,int(options.extension),options.floor,options.density,options.rpm,25,None,options.matrix,False,options.jxn)

        else:
            print('mapping to GFF and returning reads')
            if options.total:

                newGFF = mapBamToGFF(bamFile,gffFile,options.sense,options.unique,int(options.extension),options.floor,options.density,options.rpm,25,None,None,False,options.jxn)
            else:
                newGFF = mapBamToGFF(bamFile,gffFile,options.sense,options.unique,int(options.extension),options.floor,options.density,options.rpm,25,None,None,True)
        unParseTable(newGFF,output,'\t')
    else:
        parser.print_help()

if __name__ == "__main__":
    main()