#!/usr/bin/env python
# coding: utf-8

# In[5]:


from Bio import SeqIO
from bokeh.io import output_notebook, push_notebook
from bokeh.plotting import figure, show
from bokeh.models.tools import BoxZoomTool
from bokeh.models.tickers import FixedTicker
from bokeh.models import Range1d,LinearAxis
import numpy as np
output_notebook()


# In[10]:


seq = SeqIO.read('/home/tboland/SAD7743_P01_E09_LC_E09.ab1', 'abi')
quality = seq.letter_annotations['phred_quality']
peak_calls = seq.annotations['abif_raw']['PLOC2']
order = seq.annotations['abif_raw']['FWO_1']
sequence = str(seq.seq)

# Extract color data
peaks = {}
for i, (base, r) in enumerate(zip(order, [9, 10, 11, 12])):
    peaks[base] = seq.annotations['abif_raw']['DATA{}'.format(r)]

# Plot
p = figure(
    plot_width=900,
    plot_height=400,
    output_backend="webgl",
    tools=['xpan', 'reset','xwheel_zoom', BoxZoomTool(dimensions='width')],
    active_drag='xpan',
    active_scroll='xwheel_zoom',
    title=seq.id
)

# Plot quality
p.extra_y_ranges = {"q": Range1d(start=0, end=100)}
p.add_layout(LinearAxis(y_range_name="q"), 'right')
p.rect(x=peak_calls, y=np.array(quality)/2, width=14, height=quality, y_range_name='q', color='#e7ecf6')

# Plot colors
for color, base in zip(['red', 'green', 'blue', 'black'], ['T', 'A', 'C', 'G']):
    peak = peaks[base]
    p.line(range(len(peak)), peak, line_color=color, legend=base)

# Adjust tickers/ranges
p.xaxis.ticker = FixedTicker(ticks=peak_calls)
p.xaxis.major_label_overrides = dict(zip(peak_calls, sequence))
p.y_range = Range1d(0, max(max(x) for x in peaks.values()))

show(p)