# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_browser.ipynb.

# %% auto 0
__all__ = ['GenomeBrowser', 'Track']

# %% ../nbs/00_browser.ipynb 4
from fastcore.basics import *

from genomenotebook.plotting import (
    get_genome_annotations,
    get_gene_patches, 
    get_genes_from_annotation, 
    create_genome_browser_plot,
    get_all_glyphs)
from bokeh.models import (
    CustomJS,
    Range1d,
    ColumnDataSource,
    LabelSet, 
    TextInput,
    Rect,
    Div,
    Styles,
    NumeralTickFormatter
)
from bokeh.plotting import show, figure
from bokeh.layouts import column, row

from Bio import SeqIO
from .js_callback_code import x_range_change_callback_code, search_callback_code
from bokeh.io import output_notebook
import pyBigWig
import numpy as np
import pandas as pd
output_notebook()


# %% ../nbs/00_browser.ipynb 5
class GenomeBrowser:
    def __init__(self,
                 genome_path: str, #path to the fasta file of the genome sequence
                 gff_path: str, #path to the gff3 file of the annotations
                 init_pos=None, #initial position to display
                 init_win=10000, #initial window size (max=20000)
                 bounds=None, #bounds can be specified. This helps preserve memory by not loading the whole genome if not needed.
                 show_seq=True, #shows the sequence when zooming in
                 search=True, #enables a search bar to lookup a gene name or a DNA sequence
                 **kwargs):
        self.genome_path=genome_path
        genome = next(SeqIO.parse(self.genome_path, 'fasta'))
        self.genome_id = genome.id
        self.genome_size=len(genome.seq)
        
        if not bounds:
            self.bounds=(0,self.genome_size)
        else:
            self.bounds=bounds
            
        self.seq=genome.seq[self.bounds[0]:self.bounds[1]]
   
        self.gff_path=gff_path

        if not init_pos:
            self.init_pos=sum(self.bounds)//2
        elif init_pos>bounds[1] or init_pos<bounds[0]:
            print("WARNING: Requested an initial position outside of the browser bounds")
            self.init_pos=sum(self.bounds)//2
        else:
            self.init_pos=init_pos

        self.init_win=min(init_win,self.bounds[1]-self.bounds[0])

        self.show_seq=show_seq
        self.max_glyph_loading_range=20000
        self.frame_width=600

        

        self.elements=self.get_browser(show_seq=show_seq)

        if search:
            self.search=self.get_search_box()
            self.elements=[self.search]+self.elements
        else:
            self.search=None
            
        self.tracks=[]


    def get_browser(self,show_seq=True):

        annotation = get_genome_annotations(self.gff_path,self.bounds)
        annotation.gene = annotation.gene.fillna("")
        annotation.gene = annotation[['gene', 'locus_tag']].apply(lambda x: f"{x[0]} {x[1]}", axis=1)
        genes = get_genes_from_annotation(annotation)

        semi_win = self.init_win / 2
        x_range = Range1d(
            max(self.bounds[0],self.init_pos - semi_win), min(self.bounds[1],self.init_pos + semi_win), 
            bounds=self.bounds, 
            max_interval=100000,
            min_interval=40
        )

        self.glyph_source = ColumnDataSource(get_gene_patches(genes, x_range.start, x_range.end))

        #This will contain the glyphs plotted by bokeh
        loaded_glyph_source = ColumnDataSource({"start":[x_range.start],"end":[x_range.end], "range":[self.max_glyph_loading_range]})

        #This contains the glyphs for the whole genome
        self.all_glyphs=get_all_glyphs(genes, self.bounds)

        p = create_genome_browser_plot(self.glyph_source, x_range, output_backend="webgl")
        p.frame_width=self.frame_width

        sty=Styles(font_size='14px',
                font_family="Courrier",
                color="black",
                display="inline-block",
                background_color = "white",
                margin="0",
                margin_left= "2px",
                )
        
        ## Adding the ability to display the sequence when zooming in
        sequence = {
            'seq': str(self.seq).upper(),
            'bounds':self.bounds
        }

        self.div = Div(height=18, height_policy="fixed", 
                    width=600, width_policy="fixed",
                    styles = sty
                    )

        xcb = CustomJS(
            args={
                "x_range": p.x_range,
                "sequence": sequence,
                "all_glyphs":self.all_glyphs,
                "glyph_source": self.glyph_source,
                "div": self.div,
                "loaded_glyph_source":loaded_glyph_source,
            },
            code=x_range_change_callback_code
        )

        p.x_range.js_on_change('start', xcb)
        self.x_range=p.x_range
        self.gene_track=p

        if self.show_seq:
            return [p,self.div]
        else:
            return [p]

    def get_search_box(self):

        ## Create a text input widget for search
        text_input = TextInput(value="")

        ## Adding BoxAnnotation to highlight search results
        search_span_source = ColumnDataSource({"x":[],"width":[]})#"y":[]
        h=Rect(x='x',y=-2,width='width',height=self.gene_track.height,fill_color='green',fill_alpha=0.2,line_alpha=0)
        self.gene_track.add_glyph(search_span_source, h)

        ## Adding the ability to display the sequence when zooming in
        sequence = {
            'seq': str(self.seq).upper(),
            'bounds':self.bounds
        }

        call_back_search = CustomJS(
            args={
                "x_range": self.x_range,
                "glyph_source": self.glyph_source,
                "sequence": sequence,
                "all_glyphs": self.all_glyphs,
                "text_input": text_input,
                "search_span_source": search_span_source,
                "div": self.div,
            },
            code=search_callback_code
        )

        text_input.js_on_change('value',call_back_search)#,xcb)

        return text_input
    
    def show(self):
        show(column(self.elements + [t.fig for t in self.tracks]))



# %% ../nbs/00_browser.ipynb 7
class Track:
    def __init__(self,
                 height: int = 200, #size of the track
                ):
        self.height = height
        self.fig = figure(tools="xwheel_zoom,xpan,save,reset",
                          active_scroll="xwheel_zoom",
                          height=height,
                          y_axis_location="right", #this is required in order to keep a proper alignment with the sequence
                          output_backend="webgl")
        self.fig.xaxis[0].formatter = NumeralTickFormatter(format="0,0")
        


# %% ../nbs/00_browser.ipynb 9
@patch
def add_track(self:GenomeBrowser,
             height:int = 200 #size of the track
             ) -> Track:
    """Adds a track to the the GenomeBrowser. Ensures that the x_range are shared and figure widths are identical."""
    t = Track(height=height)
    t.fig.x_range = self.x_range
    t.fig.frame_width = self.frame_width
    t.bounds = self.bounds
    self.tracks.append(t)
    return t
    

# %% ../nbs/00_browser.ipynb 13
@patch
def line(self:Track,
         source: pd.DataFrame, #pandas DataFrame containing the data
         pos: str, #name of the column containing the positions along the genome
         y: str, #name of the column containing the data to be plotted on the y-axis
        ):
    source=source.loc[(self.bounds[0] < source[pos]) & (source[pos] < self.bounds[1])]
    self.fig.line(source=source, x=pos, y=y)


# %% ../nbs/00_browser.ipynb 17
from bokeh.transform import linear_cmap, factor_cmap

# %% ../nbs/00_browser.ipynb 18
@patch
def scatter(self:Track,
         source: pd.DataFrame, #pandas DataFrame containing the data
         pos: str, #name of the column containing the positions along the genome
         y: str, #name of the column containing the data to be plotted on the y-axis
         z: str = None, #name of a column containing numerical data rendered as a linear color map (cannot be used for line plots)
         factors: str = None, #name of a column of values to be used as factors
        ):
    source=source.loc[(self.bounds[0] < source[pos]) & (source[pos] < self.bounds[1])]
    
    if factors!=None:
        color=factor_cmap(factors,"Category10_3",tuple(set(source[factors].values)))
        
        self.fig.scatter(source=source, x=pos, y=y, color=color, legend_group=factors)

        self.fig.legend.location = "top_left"
        self.fig.legend.title = "ori"
    elif z!=None:
        pass
    else:
        self.fig.scatter(source=source, x=pos, y=y)


# %% ../nbs/00_browser.ipynb 21
@patch
def bar(self:Track,
         source: pd.DataFrame, #pandas DataFrame containing the data
         pos: str, #name of the column containing the positions along the genome
         y: str, #name of the column containing the data to be plotted on the y-axis
         z: str = None, #name of a column containing numerical data rendered as a linear color map (cannot be used for line plots)
         factors: str = None, #name of a column of values to be used as factors
        ):
    source=source.loc[(self.bounds[0] < source[pos]) & (source[pos] < self.bounds[1])]
    
    if factors!=None:
        color=factor_cmap(factors,"Category10_3",tuple(set(source[factors].values)))
        
        self.fig.vbar(source=source, x=pos, top=y, color=color, legend_group=factors)

        self.fig.legend.location = "top_left"
        self.fig.legend.title = "ori"
    elif z!=None:
        pass
    else:
        self.fig.vbar(source=source, x=pos, top=y)
        
