import numpy as np
from bokeh.layouts import row, column
from bokeh.models import ColumnDataSource, Slider, CustomJS, TextAnnotation
from bokeh.plotting import Figure, show
from bokeh.io import output_notebook 
# Define data
x = [x*0.05 for x in range(0, 500)]
trigonometric_functions = {
    '0': np.sin(x),
    '1': np.cos(x),
    '2': np.tan(x),
    '3': np.arctan(x)}
initial_function = '0'



# Wrap the data in two ColumnDataSources
source_visible = ColumnDataSource(data=dict(
    x=x, y=trigonometric_functions[initial_function]))
source_available = ColumnDataSource(data=trigonometric_functions)

# Define plot elements
plot = Figure(plot_width=400, plot_height=400)
plot.line('x', 'y', source=source_visible, line_width=3, line_alpha=0.6)

slider = Slider(title='Trigonometric function',
                value=int(initial_function),
                start=np.min([int(i) for i in trigonometric_functions.keys()]),
                end=np.max([int(i) for i in trigonometric_functions.keys()]),
                step=1)


# 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), code="""
        var selected_function = 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[selected_function];
        // Update the plot
        source_visible.change.emit();
    """)

layout = row(plot, slider)

output_notebook()

show(layout)

def c1t(t, lam = 1., R = .25, c10 = 1.):
    
    expt = lam * t / 2
    
    if R == .5:
        output = c10 * np.exp(-expt) * (1 + expt)
    elif R == 0:
        output = c10 * np.exp(-expt)
    elif R < .5:
        sqt = np.sqrt(1-2*R)
        output = c10 * np.exp(-expt) * (np.cosh(expt * sqt) + np.sinh(expt*sqt) / sqt)
    elif R > .5:
        sqt = np.sqrt(-1+2*R)
        output = c10 * np.exp(-expt) * (np.cos(expt * sqt) + np.sin(expt*sqt) / sqt)

        
    return output

ts = [t*0.02 for t in range(0, 500)]

Rmin = 0
Rmax = 10
Rstep = .2
Rrange = np.arange(Rmin, Rmax, Rstep)
Rrange_str = [str(i) for i in range(len(Rrange))]
#Rrange_str = ['{:.1f}'.format(i) for i in Rrange] # truncate to two decimals
#Rrange_dict = {Rrange_str[i]:Rrange[i] for i,_ in enumerate(Rrange)}


ys = {r_str:[c1t(t, R = Rrange[int(r_str)])**2 for t in ts] for r_str in Rrange_str}
#ys = {r_str:[c1t(t, R = Rrange_dict[r_str])**2 for t in ts] for r_str in Rrange_str}

initial_y = Rrange_str[1]


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

# Define plot elements
plot = Figure(plot_width=400, plot_height=400)
plot.line('x', 'y', source=source_visible, line_width=3, line_alpha=0.6)

slider = Slider(value=int(initial_function),
                start=np.min([int(i) for i in ys.keys()]),
                end=np.max([int(i) for i in ys.keys()]),
                step=1,
                show_value = False)


#slider = Slider(title='R = ',
#                value=float(initial_function),
#                start=float(Rrange_str[0]),
#                end=float(Rrange_str[-1]),
#                step=Rstep)


# 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), code="""
        var r_value = 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_value];
        // Update the plot
        source_visible.change.emit();
    """)

layout = row(plot, slider)

output_notebook()

show(layout)



