from bokeh.layouts import row, column
from bokeh.models import ColumnDataSource, Slider, CustomJS, Text, DataRange1d, Title
from bokeh.plotting import Figure, show, output_file
from bokeh.io import output_notebook 

from amplitude_damping import *

R_values = np.geomspace(start=0.2, stop=500.0, num=100)

R_values_str = [str(i) for i in range(len(R_values))]

R_values_str_R = ['R = {:.2f}'.format(R) for R in R_values] # truncate to two decimals

# make a dictionary of form {'0': 0.0, '1': 0.2, .. }
R_values_dict = {R_values_str[i]:R_values.round(2)[i] for i,_ in enumerate(R_values)} # rounding to two decimals

initial_r = R_values_str[len(R_values)//2]

npoints = 400

ts_cap = np.linspace(0,1.1,200)
ys_cap = {R: np.array([Qa(c1(R_values_dict[R],t)**2) for t in ts_cap]) for R in R_values_str}
rs_cap = {R_values_str[i] : [R_values_str_R[i]] for i,_ in enumerate(R_values)} 


# Wrap the data in two ColumnDataSources
#source_visible = ColumnDataSource(data=dict(
#    x = ts, y = ys[initial_r]))
#source_available = ColumnDataSource(data=ys)

source_visible = ColumnDataSource(data=dict(
    x = ts_cap, y = ys_cap[initial_r]))

source_available = ColumnDataSource(data=ys_cap)


# Define plot elements
plot = Figure(plot_width=400, plot_height=400, x_range = DataRange1d(), y_range=(-.01, 1.01))
plot.line('x', 'y', source=source_visible, legend_label="Q(Φ(t))", line_width=3, line_alpha=0.6)

# Add text
text_source = ColumnDataSource({'r_value': [rs_cap[initial_r]]})
r_available = ColumnDataSource(data=rs_cap)
text = Text(x=0, x_offset=315, y=.8, text='r_value', text_font_size='15pt', text_align='right')
plot.add_glyph(text_source, text)

# Add slider
slider = Slider(value=int(initial_r),
                start=np.min([int(i) for i in ys_cap.keys()]),
                end=np.max([int(i) for i in ys_cap.keys()]),
                step=1,
                show_value = False,
                title = 'R')

# Define CustomJS callback, which updates the plot based on selected function
# by updating the source_visible ColumnDataSource.
slider.callback = CustomJS(
    args=dict(source_visible=source_visible,
              source_available=source_available,
              text_source = text_source,
              r_available = r_available), code="""
        var r_idx = cb_obj.value;
        
        // Get the data from the data sources
        var data_visible = source_visible.data;
        var data_available = source_available.data;
        
        // Change y-axis data according to the selected value
        data_visible.y = data_available[r_idx];
                
        // text
        text_source.data = {'r_value': [String(r_available.data[r_idx])]};
        
        // Update the plot
        source_visible.change.emit();
        
    """)

layout = column(plot,slider)

output_file("channel_cap.html", title="Channel Capacity")
output_notebook()
show(layout)

