##################################
#                                #
# Last modified 2018/04/24       # 
#                                #
# Georgi Marinov                 #
#                                # 
##################################

import sys
import string
import pysam
import numpy as np
import random
import os
from scipy.cluster.hierarchy import dendrogram, linkage, fcluster
from sets import Set

# FLAG field meaning
# 0x0001 1 the read is paired in sequencing, no matter whether it is mapped in a pair
# 0x0002 2 the read is mapped in a proper pair (depends on the protocol, normally inferred during alignment) 1
# 0x0004 4 the query sequence itself is unmapped
# 0x0008 8 the mate is unmapped 1
# 0x0010 16 strand of the query (0 for forward; 1 for reverse strand)
# 0x0020 32 strand of the mate 1
# 0x0040 64 the read is the first read in a pair 1,2
# 0x0080 128 the read is the second read in a pair 1,2
# 0x0100 256 the alignment is not primary (a read having split hits may have multiple primary alignment records)
# 0x0200 512 the read fails platform/vendor quality checks
# 0x0400 1024 the read is either a PCR duplicate or an optical duplicate

def getReverseComplement(preliminarysequence):
    
    DNA = {'A':'T','T':'A','G':'C','C':'G','N':'N','a':'t','t':'a','g':'c','c':'g','n':'n'}
    sequence=''
    for j in range(len(preliminarysequence)):
        sequence=sequence+DNA[preliminarysequence[len(preliminarysequence)-j-1]]
    return sequence

def FLAG(FLAG):

    Numbers = [0,1,2,4,8,16,32,64,128,256,512,1024]

    FLAGList=[]

    MaxNumberList=[]
    for i in Numbers:
        if i <= FLAG:
            MaxNumberList.append(i)

    Residual=FLAG
    maxPos = len(MaxNumberList)-1

    while Residual > 0:
        if MaxNumberList[maxPos] <= Residual:
            Residual = Residual - MaxNumberList[maxPos]
            FLAGList.append(MaxNumberList[maxPos])
            maxPos-=1
        else:
            maxPos-=1
  
    return FLAGList

def run():

    if len(sys.argv) < 4:
        print 'usage: python %s BAM genome.fa sequence(s) outfilename' % sys.argv[0]
        print '\tThe BAM file should come from Bismarck, but is expected to be sorted and deduped'
        print '\tsequences should be comma separated if more than one is listed, and of the same length'
        sys.exit(1)

    BAM = sys.argv[1]
    fasta = sys.argv[2]
    SeqContexts = {}
    SClens = []
    for SC in sys.argv[3].split(','):
        SeqContexts[SC] = {}
        SeqContexts[SC]['total'] = 0
        SeqContexts[SC]['meth'] = 0
        SClens.append(len(SC))
    SClens = list(Set(SClens))
    if len(SClens) > 1:
        print 'sequence contexts have to be of the same length, exiting'
        sys.exit(1)
    SCL = SClens[0]
    print SeqContexts.keys()
    outfilename = sys.argv[4]

    GenomeDict = {}
    sequence = ''
    inputdatafile = open(fasta)
    for line in inputdatafile:
        if line[0]=='>':
            if sequence != '':
                GenomeDict[chr] = ''.join(sequence).upper()
            chr = line.strip().split('>')[1]
            sequence=[]
            Keep=False
            continue
        else:
            sequence.append(line.strip())
    GenomeDict[chr] = ''.join(sequence).upper()

    print 'finished inputting genomic sequence'

    samfile = pysam.Samfile(BAM, "rb" )

    R = 0
    for chr in GenomeDict.keys():
        start = 1
        end = len(GenomeDict[chr])
        for alignedread in samfile.fetch(chr, start, end):
            R+=1
            if R % 100000 == 0:
                print R, 'reads processed', chr, pos, end
            fields=str(alignedread).split('\t')
            FLAGfields = FLAG(int(fields[1]))
            pos = int(fields[3])
            matepos = int(fields[7])
#            readseq = alignedread.seq
            readseq_temp = alignedread.seq
            readseq = ''
            rpos = 0
            for (m,bp) in alignedread.cigar:
# soft-clipped bases:
                if m == 4:
                    rpos += bp
# matches:
                if m == 0:
                    readseq += readseq_temp[rpos:rpos+bp]
                    rpos += bp
# insertions:
# note: not handled properly, as the junction remaining after excising the insertion might be a CG or GC
# but there is no good way to deal with it
                if m == 1:
                    rpos += bp
# deletions:
                if m == 2:
                    for D in range(bp):
                        readseq += 'N'
            if alignedread.is_reverse:
                strand = '-'
            else:
                strand = '+'
            if matepos < pos + len(readseq):
                readseq = readseq[0:matepos-pos]
            for i in range(pos,pos + len(readseq) - SCL + 1):
                SC = GenomeDict[chr][i:i+SCL]
                if SeqContexts.has_key(SC):
                    SeqContexts[SC]['total'] += 1
                    if readseq[i-pos:i-pos+SCL] == SC:
                        SeqContexts[SC]['meth'] += 1
    outfile = open(outfilename, 'w')

    outline = '#SequenceContext\tTotal\tMethylated\tFraction'

    for SC in SeqContexts.keys():
        outline = SC + '\t' + str(SeqContexts[SC]['total']) + '\t' + str(SeqContexts[SC]['meth']) + '\t' + str(SeqContexts[SC]['meth']/(SeqContexts[SC]['total'] + 0.0))
        outfile.write(outline+'\n')

    outfile.close()
           
run()