import numpy as np
from bokeh.layouts import row, column
from bokeh.models import ColumnDataSource, Slider, CustomJS, Text
from bokeh.plotting import Figure, show, output_file
from bokeh.io import output_notebook 

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) * (np.cosh(expt * sqt) + np.sinh(expt*sqt) / sqt)
    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

def lorentzian_J(w, R = .25, omega_0 = 0.):
    # assume gamma_0 = 1
    if R == 0.:
        return 1
    lam = 1. / R
    output = 1/(2*np.pi) * lam**2 / ((omega_0 - w)**2 + lam**2)
    return output

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

Rrange = [r*.02 for r in range(0,int(1/.02))] + [r * .1 for r in range(10,100)] + [r for r in range(10,100+1)]
Rrange = np.array(Rrange)

Rrange_str = [str(i) for i in range(len(Rrange))]

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

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


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_r = Rrange_str[len(Rrange)//2]

ws = [t*0.02 for t in range(-250, 250)]
js = {r_str:[lorentzian_J(w, R = Rrange[int(r_str)]) for w in ws] for r_str in Rrange_str}

rs = {Rrange_str[i] : [Rrange_str_R[i]] for i,_ in enumerate(Rrange)} 

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

# Wrap the data in two ColumnDataSources
source_visible2 = ColumnDataSource(data=dict(
    x = ws, y = js[initial_r]))
source_available2 = ColumnDataSource(data=js)


# Define plot elements
plot = Figure(plot_width=300, plot_height=300, x_range=(-.1, 10), y_range=(-.01, 1))
plot.line('x', 'y', source=source_visible, legend_label="ρ₁₁(t)", line_width=3, line_alpha=0.6)

plot2 = Figure(plot_width=300, plot_height=300, x_range=(-5, 5), y_range=(-.001, .2))
plot2.line('x', 'y', source=source_visible2, legend_label="J(ω)", 
           line_width=3, line_alpha=0.6, line_color="#f01001")


# Add text
text_source = ColumnDataSource({'r_value': ['%s' % Rrange_str_R[1]]})
r_available = ColumnDataSource(data=rs)
text = Text(x=9.5, y=.7, 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.keys()]),
                end=np.max([int(i) for i in ys.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,
              source_visible2=source_visible2,
              source_available2=source_available2,
              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;
        var data_visible2 = source_visible2.data;
        var data_available2 = source_available2.data;
        
        // Change y-axis data according to the selected value
        data_visible.y = data_available[r_idx];
        data_visible2.y = data_available2[r_idx];
        
        
        // text
        text_source.data = {'r_value': [String(r_available.data[r_idx])]};
        
        // Update the plot
        source_visible.change.emit();
        source_visible2.change.emit();        
    """)

layout = row(column(plot,slider), plot2)



output_file("jaynescummings.html", title="Jaynes-Cummings Model")
output_notebook()
show(layout)

