import pyBigWig
from collections import defaultdict
from matplotlib import pyplot as plt
import numpy as np

site = defaultdict(list)
with open('../metadata/reference/tss_sites.txt') as fin:
    for line in fin:
        line = line.strip().split('-')
        site[line[0]].append([line[1],line[2],int(line[3])])

#ChIP-seq H3K27ac        
bws = {}
# 1-H3K27AC_combined: ChIP-seq_for_HeLa-AB1_H3K27ac_AM_25uw_low
# 2-H3K27AC_combined: ChIP-seq_for_HeLa-AB1_H3K27ac_AM_25uw_high
# 3-H3K27AC_combined: ChIP-seq_for_HeLa-AB1_H3K27ac_Dark_control
prefix = ['1-H3K27AC_combined','2-H3K27AC_combined','3-H3K27AC_combined']
for p in prefix:
    bws[p] = pyBigWig.open('../metadata/chip-seq/{}.SeqDepthNorm.bw'.format(p))

fig, ax = plt.subplots()
length = 1000
data = {}
for p in prefix:
    depth = bws[p]
    total = np.array(depth.values(site['mRN'][0][1],
                                  site['mRN'][0][2]-length,
                                  site['mRN'][0][2]+length))
    for s in site['mRN'][1:]:
        values = depth.values(s[1],s[2]-length,s[2]+length)
        if s[0] == 'forward':
            total += np.array(values)
        else:
            values.reverse()
            total += np.array(values)
    data[p] = total
    plt.plot(np.array(range(0,2*length))-length,total)
plt.legend(prefix)

# #ATAC-seq
# bws = {}
# # high: ATAC-seq_for_HeLa-AB1_AM_25uw_high
# # low: ATAC-seq_for_HeLa-AB1_AM_25uw_low
# # dark: ATAC-seq_for_HeLa-AB1_Dark_control
# prefix = ['high','low','dark']
# for p in prefix:
#     bws[p] = pyBigWig.open('../metadata/atac-seq/{}.SeqDepthNorm.bw'.format(p))
    
# fig, ax = plt.subplots()
# length = 1000
# for p in prefix:
#     depth = bws[p]
#     total = np.array(depth.values(site['mRN'][0][1],
#                                   site['mRN'][0][2]-length,
#                                   site['mRN'][0][2]+length))
#     for s in site['mRN'][1:]:
#         values = depth.values(s[1],s[2]-length,s[2]+length)
#         if s[0] == 'forward':
#             total += np.array(values)
#         else:
#             values.reverse()
#             total += np.array(values)
#     data[p] = total
#     plt.plot(np.array(range(0,2*length))-length,total)
# plt.legend(prefix)

# #scATAC-seq

# prefix = []
# with open('../rawdata/scatac-seq/HGC20191230001-0003_lane7/L7_md5sum.check.out') as fin:
#     pattern = '_R1_001.fastq'
#     for line in fin:
#         if pattern in line and 'HW' in line:
#             line = line.split(' ')[0]
#             line = line.replace(pattern,' ')
#             line = line.split(' ')[0]
#             prefix.append(line)
# with open('../rawdata/scatac-seq/HGC20191230001-0003_lane8/L8_md5sum.check.out') as fin:
#     pattern = '_R1_001.fastq'
#     for line in fin:
#         if pattern in line and 'HW' in line:
#             line = line.split(' ')[0]
#             line = line.replace(pattern,' ')
#             line = line.split(' ')[0]
#             prefix.append(line)

# bws = {}
# for p in prefix:
#     bws[p] = pyBigWig.open('../metadata/scatac-seq/{}.SeqDepthNorm.bw'.format(p))

# # HW1: scATAC-seq_for_HeLa-AB1_150min_light_on_80uw_384
# # HW2: scATAC-seq_for_HeLa-AB1_dark_control_384
# # HW3: scATAC-seq_for_HeLa-AB1_600min_light_on_80uw_384
# # HW4: scATAC-seq_for_HeLa-AB1_750min_light_on_80uw_384
# # HW5: scATAC-seq_for_HeLa-AB1_1200min_light_on_80uw_384
# # HW6: scATAC-seq_for_HeLa-AB1_1350min_light_on_80uw_384
# exp = ['HW1','HW2','HW3','HW4','HW5','HW6']
# length = 1000
# data = {}
# for e in exp:
#     total_exp = []
#     for p in prefix:
#         if e in p:
#             depth = bws[p]
#             try:
#                 total = np.array(depth.values(site['mRN'][0][1],
#                                           site['mRN'][0][2]-length,
#                                           site['mRN'][0][2]+length))
#             except RuntimeError:
#                 print('mRN',e,p,s)
#             for s in site['mRN'][1:]:
#                 try:
#                     values = depth.values(s[1],s[2]-length,s[2]+length)
#                     if s[0] == 'forward':
#                         total += np.array(values)
#                     else:
#                         values.reverse()
#                         total += np.array(values)
#                 except RuntimeError:
#                     print('mRN',e,p,s)
#             total_exp.append(total)
#     total_exp = np.array(total_exp)
#     data[e] = total_exp