import numpy as np
from bokeh.io import curdoc
from bokeh.layouts import layout
from bokeh.models import Slider
from bokeh.plotting import figure


def f(x, t, k, w, dk, dw):
    envelope = np.cos(dk * x - dw * t)
    carrier = envelope * np.cos(k * x - w * t)
    return carrier, envelope


x = np.linspace(0, 100, 500)
w = k = 1
t = 0
c, e = f(x, t, k, w, 0, 0)

p = figure(height=300, width=900, x_range=(0, 100), y_range=(-1.1, 1.1))

p.toolbar.active_drag = None
p.toolbar.active_scroll = None
p.toolbar.active_tap = None
p.xaxis.axis_label = "x"
p.xaxis.axis_label_text_font_size = "16pt"
p.xaxis.major_label_text_font_size = "16pt"
p.yaxis.axis_label = "u(x,t)"
p.yaxis.axis_label_text_font_size = "16pt"
p.yaxis.major_label_text_font_size = "16pt"

car = p.line(x, c, color="blue", line_width=2)
env0 = p.line(x, e, color="red", line_width=2)
env1 = p.line(x, -e, color="red", line_dash="10 10", line_width=1)

slider_w = Slider(title="w", start=0.0, end=1.0, value=0.5, step=0.001, width=450)
slider_k = Slider(title="k", start=0.0, end=1.0, value=1.0, step=0.001, width=450)
slider_dk = Slider(title="dk", start=-0.1, end=0.1, value=0.1, step=0.001, width=450)
slider_dw = Slider(title="dw", start=-0.4, end=0.4, value=0.4, step=0.001, width=450)

l = layout([[p], [slider_w, slider_k], [slider_dw, slider_dk]])


def update():
    global t
    t += 0.05
    c, e = f(x, t, slider_k.value, slider_w.value, slider_dk.value, slider_dw.value)
    car.data_source.data["y"] = c
    env0.data_source.data["y"] = e
    env1.data_source.data["y"] = -e


curdoc().add_periodic_callback(update, 10)
curdoc().add_root(l)