%pylab inline

"""
Explanation: Set the matplotlib magic to notebook enable inline plots
End of explanation
"""


import subprocess
import matplotlib.pyplot as plt
import random
import numpy as np

"""
Explanation: Calculate the Nonredundant Read Fraction (NRF)
SAM format example:
SRR585264.8766235       0       1       4       15      35M     *       0       0       CTTAAACAATTATTCCCCCTGCAAACATTTTCAAT     GGGGGGGGGGGGGGGGGGGGGGFGGGGGGGGGGGG    XT:A:U   NM:i:1  X0:i:1  X1:i:6  XM:i:1  XO:i:0  XG:i:0  MD:Z:8T26
Import the required modules
End of explanation
"""


plt.style.use('ggplot')
figsize(10,5)

"""
Explanation: Make figures prettier and biger
End of explanation
"""


file = "/ngschool/chip_seq/bwa/input.sorted.bam"

"""
Explanation: Parse the SAM file and extract the unique start coordinates.
First store the file name in the variable
End of explanation
"""


p = subprocess.Popen(["samtools", "view", "-q10", "-F260", file],
                      stdout=subprocess.PIPE)
coords = []
for line in p.stdout:
    flag, ref, start = line.decode('utf-8').split()[1:4]
    coords.append([flag, ref, start])

coords[:3]

"""
Explanation: Next we read the file using samtools. From each read we need to store the flag, chromosome name and start coordinate.
End of explanation
"""


len(coords)

"""
Explanation: What is the total number of our unique reads?
End of explanation
"""


random.seed(1234)
sample = random.sample(coords, 1000000)

len(sample)

"""
Explanation: Randomly sample the coordinates to get 1M for NRF calculations
End of explanation
"""


uniqueStarts = {'watson': set(), 'crick': set()}
for coord in sample:
    flag, ref, start = coord
    if int(flag) & 16:
        uniqueStarts['crick'].add((ref, start))
    else:
        uniqueStarts['watson'].add((ref, start))

"""
Explanation: How many of those coordinates are unique? (We will use the set python object which only the unique items.)
End of explanation
"""


len(uniqueStarts['watson'])

"""
Explanation: How many on the Watson strand?
End of explanation
"""


len(uniqueStarts['crick'])

"""
Explanation: And on the Crick?
End of explanation
"""


NRF_input = (len(uniqueStarts['watson']) + len(uniqueStarts['crick']))*1.0/len(sample)
print(NRF_input)

"""
Explanation: Calculate the NRF
End of explanation
"""


def calculateNRF(filePath, pickSample=True, sampleSize=10000000, seed=1234):
    p = subprocess.Popen(['samtools', 'view', '-q10', '-F260', filePath],
                          stdout=subprocess.PIPE)
    coordType = np.dtype({'names': ['flag', 'ref', 'start'],
                          'formats': ['uint16', 'U10', 'uint32']})
    coordArray = np.empty(10000000, dtype=coordType)
    i = 0
    for line in p.stdout:
        if i >= len(coordArray):
            coordArray = np.append(coordArray, np.empty(1000000, dtype=coordType), axis=0)
        fg, rf, st = line.decode('utf-8').split()[1:4]
        coordArray[i] = np.array((fg, rf, st), dtype=coordType)
        i += 1
    coordArray = coordArray[:i]
    sample = coordArray
    if pickSample and len(coordArray) > sampleSize:
        np.random.seed(seed)
        sample = np.random.choice(coordArray, sampleSize, replace=False)
    uniqueStarts = {'watson': set(), 'crick': set()}
    for read in sample:
        flag, ref, start = read
        if flag & 16:
            uniqueStarts['crick'].add((ref, start))
        else:
            uniqueStarts['watson'].add((ref, start))
    NRF = (len(uniqueStarts['watson']) + len(uniqueStarts['crick']))*1.0/len(sample)
    return NRF

"""
Explanation: Lets create a function from what we did above and apply it to all of our files!
To use our function on the real sequencing datasets (not only on a small subset) we need to optimize our method a bit- we will use python module called numpy.
End of explanation
"""


NRF_chip = calculateNRF("/ngschool/chip_seq/bwa/sox2_chip.sorted.bam", sampleSize=1000000)
print(NRF_chip)

"""
Explanation: Calculate the NRF for the chip-seq sample
End of explanation
"""


plt.bar([0,2],[NRF_input, NRF_chip], width=1)
plt.xlim([-0.5,3.5]), plt.xticks([0.5, 2.5], ['Input', 'ChIP'])
plt.xlabel('Sample')
plt.ylabel('NRF')
plt.ylim([0, 1.25]), plt.yticks(np.arange(0, 1.2, 0.2))
plt.plot((-0.5,3.5), (0.8,0.8), 'red', linestyle='dashed')
plt.show()

"""
Explanation: Plot the NRF!
End of explanation
"""


countList = []
with open('/ngschool/chip_seq/bedtools/input_coverage.bed', 'r') as covFile:
    for line in covFile:
        countList.append(int(line.strip('\n').split('\t')[3]))
countList[0:6]

countList[-15:]

"""
Explanation: Calculate the Signal Extraction Scaling
Load the results from the coverage calculations
End of explanation
"""


plt.plot(range(len(countList)), countList)
plt.xlabel('Bin number')
plt.ylabel('Bin coverage')
plt.xlim([0, len(countList)])
plt.show()

"""
Explanation: Lets see where do our reads align to the genome. Plot the distribution of tags along the genome.
End of explanation
"""


countList.sort()

countList[0:6]

"""
Explanation: Now sort the list- order the windows based on the tag count
End of explanation
"""


countSum = sum(countList)
countSum

"""
Explanation: Sum all the aligned tags
End of explanation
"""


countFraction = []
for i, count in enumerate(countList):
    if i == 0:
        countFraction.append(count*1.0 / countSum)
    else:
        countFraction.append((count*1.0 / countSum) + countFraction[i-1])

"""
Explanation: Calculate the summaric fraction of tags along the ordered windows.
End of explanation
"""


countFraction[-5:]

"""
Explanation: Look at the last five items of the list:
End of explanation
"""


winNumber = len(countFraction)
winNumber

"""
Explanation: Calculate the number of windows.
End of explanation
"""


winFraction = []
for i in range(winNumber):
    winFraction.append(i*1.0 / winNumber)

"""
Explanation: Calculate what fraction of a whole is the position of each window.
End of explanation
"""


winFraction[-5:]

"""
Explanation: Look at the last five items of our new list:
End of explanation
"""


def calculateSES(filePath):
    countList = []
    with open(filePath, 'r') as covFile:
        for line in covFile:
            countList.append(int(line.strip('\n').split('\t')[3]))
    plt.plot(range(len(countList)), countList)
    plt.xlabel('Bin number')
    plt.ylabel('Bin coverage')
    plt.xlim([0, len(countList)])
    plt.show()
    countList.sort()
    countSum = sum(countList)
    countFraction = []
    for i, count in enumerate(countList):
        if i == 0:
            countFraction.append(count*1.0 / countSum)
        else:
            countFraction.append((count*1.0 / countSum) + countFraction[i-1])
    winNumber = len(countFraction)
    winFraction = []
    for i in range(winNumber):
        winFraction.append(i*1.0 / winNumber)
    return [winFraction, countFraction]

"""
Explanation: Now prepare the function!
End of explanation
"""


chipSes = calculateSES("/ngschool/chip_seq/bedtools/sox2_chip_coverage.bed")

"""
Explanation: Use our function to calculate the signal extraction scaling for the Sox2 ChIP sample:
End of explanation
"""


plt.plot(winFraction, countFraction, label='input')
plt.plot(chipSes[0], chipSes[1], label='Sox2 ChIP')
plt.ylim([0,1])
plt.xlabel('Ordered window franction')
plt.ylabel('Genome coverage fraction')
plt.legend(loc='best')
plt.show()

"""
Explanation: Now we can plot the calculated fractions for both the input and ChIP sample:
End of explanation
"""


