#!/usr/bin/env python3
import re
import os
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from glob import glob
from Bio import SeqIO
import numpy as np
from collections import defaultdict

fsa_files = sorted(glob('./*G2-?+*.fsa'), key = lambda x:re.findall(r"\d+.", x))
fsb_files = sorted(glob('./*G2-?-*.fsa'), key = lambda x:re.findall(r"\d+.", x))
fsc_files = sorted(glob('./*C21-?+*.fsa'), key = lambda x:re.findall(r"\d+.", x))
fsd_files = sorted(glob('./*C21-?-*.fsa'), key = lambda x:re.findall(r"\d+.", x))
#fsd_files = sorted(glob('./77-.fsa'), key = lambda x:re.findall(r"\d+.", x))
#fsc_files = sorted(glob('./77+.fsa'), key = lambda x:re.findall(r"\d+.", x))
#fsd_files = sorted(glob('./77-.fsa'), key = lambda x:re.findall(r"\d+.", x))

for i in range(len(fsa_files)):
    print('processing File #{:03d}'.format(i))
    plt.subplot(2,2,1)
    record = SeqIO.read(fsa_files[i], "abi")
    channels = ["DATA1", "DATA2", "DATA3", "DATA4","DATA5","DATA105"]
    trace = defaultdict(list)
    for c in channels:
        trace[c] = record.annotations["abif_raw"][c]
    dataS = trace["DATA105"]
    peakS, _ = find_peaks(dataS, width=1)
    a = np.array(peakS[204])
    b = np.array(peakS[505])
    data1 = trace["DATA1"]#[(a+401):(b+71)]
    data2 = trace["DATA2"]#[(a+401):(b+71)]
    #peaks, _ = find_peaks(data1, width=1)
    #plt.plot(peaks, np.array(data1)[peaks], "x", color = 'blue', label = re.findall(r"_(.+?)_",fsa_files[i]))
    plt.plot(data1,  label = fsa_files[i])
    plt.plot(data2,  label = 'seq' + fsa_files[i])
    plt.legend(fontsize=6)
    #plt.plot(data1)
    #print (len(peaks))
    #plt.hist(np.array(data1)[peaks], bins=100, label = re.findall(r"_(.+?)_",fsa_files[i]))
for i in range(len(fsb_files)):
    plt.subplot(2,2,2)
    record = SeqIO.read(fsb_files[i], "abi")
    channels = ["DATA1", "DATA2", "DATA3", "DATA4","DATA5","DATA105"]
    trace = defaultdict(list)
    for c in channels:
        trace[c] = record.annotations["abif_raw"][c]
    data1 = trace["DATA1"]#[(a+401):(b+71)]
    data2 = trace["DATA2"]#[(a+401):(b+71)]
    plt.plot(data1,  label = fsb_files[i])
    plt.plot(data2,  label = 'seq' + fsb_files[i])
    plt.legend(fontsize=6)

for i in range(len(fsc_files)):
    plt.subplot(2,2,3)
    record = SeqIO.read(fsc_files[i], "abi")
    channels = ["DATA1", "DATA2", "DATA3", "DATA4","DATA5","DATA105"]
    trace = defaultdict(list)
    for c in channels:
        trace[c] = record.annotations["abif_raw"][c]
    data1 = trace["DATA1"]#[(a+401):(b+71)]
    data2 = trace["DATA2"]#[(a+401):(b+71)]
    plt.plot(data1,  label = fsc_files[i])
    plt.plot(data2,  label = 'seq' + fsc_files[i])
    plt.legend(fontsize=6)

for i in range(len(fsd_files)):
    plt.subplot(2,2,4)
    record = SeqIO.read(fsd_files[i], "abi")
    channels = ["DATA1", "DATA2", "DATA3", "DATA4","DATA5","DATA105"]
    trace = defaultdict(list)
    for c in channels:
        trace[c] = record.annotations["abif_raw"][c]
    data1 = trace["DATA1"]#[(a+401):(b+71)]
    data2 = trace["DATA2"]#[(a+401):(b+71)]
    plt.plot(data1,  label = fsd_files[i])
    plt.plot(data2,  label = 'seq' + fsd_files[i])
    plt.legend(fontsize=6)
#plt.legend(fontsize=6)
plt.savefig('raw_1m7-fsa.pdf')
plt.show()
