from bokeh.client import push_session
from bokeh.io import curdoc
from bokeh.models import (ColumnDataSource, DataRange1d, Plot, LinearAxis, Grid,
                          Circle, HoverTool, BoxSelectTool)
from bokeh.models.widgets import (Select, DataTable, TableColumn, StringFormatter, NumberFormatter,
                                  StringEditor, IntEditor, NumberEditor, SelectEditor)
from bokeh.models.layouts import Row, Column, WidgetBox
from bokeh.plotting import figure
from bkcharts import Bar

from bokeh.sampledata.autompg import autompg as df
# from bokeh.sampledata.autompg2 import autompg2 as df2

df = df.iloc[:10, :]
# df.name = df.name.astype(str)
df = df.astype({'name': str, 'mpg': float,  'cyl': int,  'displ': float, 'hp': int,  'weight': int,
                'accel': float, 'yr': int, 'origin': int})
print(df)


class PlotBarsDataTable():
    def __init__(self, df, bar_col_names=['cyl', 'mpg', 'accel']):
        self.df = df
        self.df_orig = df.copy()
        self.bar_col_names = bar_col_names
        self.document = curdoc()  # container for all bokeh objects
        # initialize filters and data source for data table
        self.name_filter = None
        self.source = ColumnDataSource(df)  # initialize with dataframe
        # create layout
        # add model as the root module of the document ("on_change" callbacks registered)
        self.document.add_root(self.create_layout())
        # push document to the bokeh server
        self.session = push_session(self.document)

    def create_layout(self):
        df = self.df
        # selector
        names = sorted(df['name'].unique())
        name_select = Select(title='Car names:', value='All', options=['All'] + names)
        name_select.on_change('value', self.on_name_change)

        # data table columns, tablecolumn has width arg!!!
        columns = []
        for col_name in df.columns:
            if col_name == 'name':
                columns += [TableColumn(field=col_name, title=col_name,
                                        editor=SelectEditor(options=names))]
            else:
                columns += [TableColumn(field=col_name, title=col_name)]
        data_table = DataTable(source=self.source, columns=columns, editable=False)

        # bar plot new
        p = figure(title='mpg cyl accel for car name')
        p.vbar(x='name', top='mpg', width=None, source=self.source)

        # bar chart old
        # p = Bar(df, label='vars', group='name',
        #         values=blend(*self.bar_col_names, name='values', labels_name='vars'),
        #         title='mpg cyl accel for car name', color='name', legend=True)
        # tooltips, select tool later!!
        selectors = WidgetBox(name_select)
        plot_panel = Row(children=[selectors, p])
        self.layout = Column(children=[plot_panel, data_table])

        return self.layout

    # "on_change" functions
    def on_name_change(self, attr, _, value):
        self.name_filter = None if value == "All" else value
        self.update_data()

    def update_data(self):
        if self.name_filter:
            df = self.df_orig[self.df_orig['name'] == self.name_filter]
        else:
            df = self.df_orig
        self.source.data = ColumnDataSource.from_df(df)

    def run(self, do_view=False):  #, poll_interval=0.5):
        if do_view:
            self.session.show()
        self.session.loop_until_closed()

if __name__ == '__main__':
    plot = PlotBarsDataTable(df)
    plot.run(True)