# The Python Code for The Double Slit Experiment Visualization
# ------------------------------------------------------------
# using the Bokeh library http://bokeh.org/
# by Jana Legerská
# jana.legerska@matfyz.cuni.cz
#-------------------------------------------------------------

import numpy as np
from numpy import math

from bokeh.layouts import column, row
from bokeh.models import ColumnDataSource, CustomJS, Button, RadioButtonGroup, Slider
from bokeh.models import CheckboxButtonGroup, Panel, Tabs, RadioGroup, CheckboxGroup, Div
from bokeh.plotting import figure, show, output_file
from bokeh.models import NumeralTickFormatter

# With Young interferention only, no difraction

max_x = 30  # Screen length [cm]
var_x = 2  #  Gaussian distribution variance
slitdist_factor = 100 # Ratio for exaggerating Gaussian side shift
slitdist = 500 # Slit distance [nm]
slitwidth = 500 # Slit width
wavelength = 100 # Wavelength [nm]
color = "red" # Initial particle traces colour
L = 20 # distance from source to the screen (L >> slitdist) [cm]

source = ColumnDataSource(data=dict(x=[], y=[], colors=[]), tags=[dict(max_x=max_x, slitdist_factor=slitdist_factor, 
                                                                       macro_scale=True, # the applet starts with interference
                                                                       slitwidth=slitwidth, L=L, color=color)])

# Screen with points
p = figure(
    title="Screen of the double slit experiment",
    tools="",         # (de)activate Bokeh tools    
    sizing_mode="stretch_width", 
    width=800,
    height=380,
    y_axis_label=None,
    x_range=(-max_x,max_x),
    y_range=(0, 1)
)

# Delete y-axis ticks on screen
p.yaxis.minor_tick_in = 0
p.yaxis.minor_tick_out = 0
p.yaxis.major_tick_in = 0
p.yaxis.major_tick_out = 0
p.yaxis.major_label_text_color = "white"
p.xaxis.major_label_text_color = "white"
p.yaxis[0].formatter = NumeralTickFormatter(format="0.00")

# Deactivate Bokeh logo
p.toolbar.logo = None

p.circle('x', 'y', source=source, size=7, fill_color='colors', line_color="black")

# Probability denstity function above the screen + histogram
pd = figure(title="Histogram and probability density",
            sizing_mode="stretch_width",  
            width=p.width, height=250, x_range=p.x_range, tools="")
pd.toolbar.logo = None

source_pd = ColumnDataSource(data=dict(x=[], y=[]), 
                             tags=[dict(var_x=var_x, previous_num=0)]) 
                             # previous_num ~ the number of particles sent before the change of distribution
pd_line = pd.line('x', 'y', source=source_pd, line_width=2, line_color="navy", visible=False)

# Delete y-axis ticks on pdf graph
pd.yaxis.minor_tick_in = 0
pd.yaxis.minor_tick_out = 0
pd.yaxis.major_tick_in = 0
pd.yaxis.major_tick_out = 0
pd.xaxis.major_label_text_color = "white"
pd.yaxis[0].formatter = NumeralTickFormatter(format="0.00")

# Histogram -- Young double slit
Nbins = 60 # number of bins
binwidth = 2*max_x/Nbins

# hist, bins = np.histogram([], range=(0, max_x), bins=Nbins) 
# hist = relative frequancies
# bins = array array of division points of bins
# counts = absolute frequencies

source_hist = ColumnDataSource(data=dict(hist=[0 for _ in range(Nbins)], 
                                         bins=[binwidth/2 + i*binwidth - max_x for i in range(Nbins)], # bins centering
                                         counts=[0 for _ in range(Nbins)], # array of absolute frequencies
                                         widths = [binwidth for _ in range(Nbins)] # columns widths
                                        )
                              )  
pd.vbar(top="hist", x="bins", width="widths", source=source_hist, alpha=0.6, color="skyblue")

#------------------------------------------------------------------------------------------------
# Buttons definition

# Add 1 particle
button = Button(label="Send 1 particle", button_type="success")

# Send 100 particles at once
button100 = Button(label="Send 100 particles", button_type="success")

# Reset the screen
resetbutton = Button(label="Reset", button_type="success")

# Continuous stream of particles
streamtoggle = Button(label="Send continuous stream", button_type="success")

# Change distribution
LABELS = ["Uniform", "Normal", "Young"]
button_change_dist = RadioButtonGroup(labels=LABELS, active=0)

# Block/open slits
LABELS = ["Slit 1", "Slit 2"]
button_choose_slit = CheckboxButtonGroup(labels=LABELS, active=[0, 1])

# Detector
button_detector = Button(label="Detector", button_type="default")

# Sliders
slider_slitdist = Slider(start=0, end=1000, value=slitdist, step=10, title=None)
slider_slitwidth = Slider(start=0, end=5, value=0.1, step=0.1, title="Slit width")
slider_wavelength = Slider(start=100, end=1000, value=wavelength, step=10, title=None)

# Text -- slider labels
label_slitdist = Div(text="Slit distance: <b>{}</b> [nm]".format(slitdist), width=200, height=10)
label_wavelength = Div(text="Wavelength: <b>{}</b> [nm]".format(wavelength), width=200, height=10)
label_energy = Div(text="Energy: <b>{}</b> [eV]".format(round((6.626*3)/(wavelength)/1.602*100, 1)), width=200, height=20)

label_screen = Div(text="", height=20, align='center') # x [cm]

# Choose particle colour
LABELS_color = ["Red", "Blue", "Green"]
color_button = RadioGroup(labels=LABELS_color, active=0)

# Keep previous measurements after set up is changed
button_keep = CheckboxGroup(labels=["Keep previous measurements"], active=[])
text_keep_previous = Div(text="Keeping previous measurements is only possible if the sreen scale does not change. \n When using a detector, the screen is zoomed to nanometers scale. ", 
                         height=20, align='center', visible=False)

# Visibility of panels, theoretical PDF
check_visible_pdf = CheckboxGroup(labels=["Show theoretical probability density"], active=[])
check_visible_pd_panel = CheckboxGroup(labels=["Show histogram panel"], active=[0])

# -----------------------------------------------------------------------------
# JS functions definition

# Def function of normal distribution with parameters mean, var (to be read from sliders)
js_normalPDF = """
    function randn(mean, variance) { // Box Muller Transform
      let u = 0, v = 0;
      
      while(u === 0) u = Math.random(); //Converting [0,1) to (0,1)
      while(v === 0) v = Math.random();
      
      let num = Math.sqrt( -2.0 * Math.log( u ) ) * Math.cos( 2 * Math.PI * v );
      
      num = num * variance + mean
      
      if (num < -source.tags[0]['max_x']) {return randn()}  // resample between -max_x and max_x (abandon values out of screen)
      else if (source.tags[0]['max_x'] < num) {return randn();}
      else {return num;}
    }
 """

# Def function of Young double slit interference distribution
js_youngPDF = """
    function youngPDF(x) { // PDF of Young interferention
        const freq = Math.PI*slider_slitdist.value/(slider_wavelength.value*source.tags[0].L)
        return Math.cos(freq*x)**2
    }
    
    let minI = -source.tags[0].max_x;
    let N = 100;
    
    function youngCDF(maxI) { // CDF of Young interferention
      let integral = 0;
      let delta = (maxI - minI)/N
      for (let i = 0; i < N; i++) {
          integral += youngPDF(delta/2 + i*delta + minI)*delta;
      } 
      return integral
    }
    
    const sample = [];    // array of divistion points of the x-axis (up to max_x, step 0.01)
    for (let i = -source.tags[0].max_x; i < source.tags[0].max_x; i += 0.01) {
      sample.push(i);
    }
    const sample_y = sample.map(youngCDF);  // youngCDF function values = the division points of the y-axis
    
    // Normalization of PDF to the length of shade (2*max_x)
    const h = sample_y[sample_y.length-1] - sample_y[0];

    let index;
    function youngInverseCDF(t) { // inverse CDF of Young interferention
        for (let i = 0; i < sample_y.length; i++) {
            if (sample_y[i] > t) {
            index = i;
            break;
            }
        }
        const x2 = sample[index];
        const x1 = sample[index - 1];
        const y1 = sample_y[index - 1]
        const c = (x2 - x1)/(sample_y[index] + y1);
        return c*t - c*y1 + x1;
    }
"""

# Function to send one particle 
js_send = """
    function send(u, renormalize=true) {
        if (u!=false) { //nezapo49t8v8 částice při zavřených štěrbinách
            source.data.x.push(u)
            source.data.y.push(Math.random())
            source.data.colors.push(source.tags[0].color) // Read particle colour from source
            source.change.emit()  // update the data source with local changes

            // add to histogram
            let j = Math.floor((u + source.tags[0].max_x)/s_hist.data.widths[0]);  // index of bin where a particle with coord. new_x is added
            s_hist.data.counts[j]++;
            if (renormalize) {normalizehist()}; 
        }
        
    }
    
    // normalize histogram
    function normalizehist() {
       // normalization; number of particles = source.data.x.length
       let totalarea = source.data.x.length*s_hist.data.widths[0]; // total number of particles: source.data.x.length
        for (let i = 0; i < s_hist.data.hist.length; i++) {
            s_hist.data.hist[i] = s_hist.data.counts[i]/totalarea;
        }
        s_hist.change.emit();
    }
"""

# Generator of a random number from a given distribution (uniform / normal(max_x/2, var_x) / young)
js_generator = """
    function generator() {

        const slitdist_adjusted = slider_slitdist.value/source.tags[0].slitdist_factor
        // variance adjusted to approximate the one slit diffraction peak by gaussian distribution
        let var_gaussdiff = (slider_wavelength.value*source.tags[0].L)/(Math.PI*source.tags[0].slitwidth)
        
        if (slit.active.length==2){
           if (detector.button_type == "primary") {

               if (Math.random()<0.5) {
                   return randn(-slitdist_adjusted, var_gaussdiff) // source_pd.tags[0].var_x
               }
               else {
                   return randn(slitdist_adjusted, var_gaussdiff)
               }
           }
           else {
               return youngInverseCDF(h*Math.random())
           }
        } 
        else if (slit.active.length==1){
            if (slit.active[0]==0 ) {
                return randn(-slitdist_adjusted, var_gaussdiff)
            }
            else {
                return randn(slitdist_adjusted, var_gaussdiff)
            }
        }
        else {
            return false
        }
    }
"""

# Function for reseting the screen
js_reset = """
    function reset(resetpdf=true) { // no parameter => reset all
        // reset data in source
        source.data.x = []
        source.data.y = []
        source.data.colors = []
        source.change.emit()

        // reset histogram
        for (let i = 0; i < s_hist.data.hist.length; i++) {
                s_hist.data.hist[i] = 0;
                s_hist.data.counts[i] = 0;
                }
        s_hist.change.emit()
        
        if (resetpdf) { // parameter resetpdf=false => reset all except pdf
        // reset pdf
            source_pd.data.x = []
            source_pd.data.y = []
            source_pd.change.emit()
        }
        
        source_pd.tags[0].previous_num = 0
    }
"""


#-------------------------------------------------------------------------
# JS CALLBACKS
js_functions = js_normalPDF + js_youngPDF + js_send + js_generator + js_reset + js_update_scale

js_cb_objects = dict(source=source, source_pd=source_pd, s_hist=source_hist, 
                     streamtoggle=streamtoggle, 
                     selected=button_change_dist, slit=button_choose_slit, detector=button_detector, 
                     slider_slitdist=slider_slitdist, slider_slitwidth=slider_slitwidth, 
                     slider_wavelength=slider_wavelength, 
                     label_wavelength=label_wavelength, label_energy=label_energy, label_slitdist=label_slitdist,  
                     color_button=color_button, button_keep=button_keep, text_keep_previous=text_keep_previous, 
                     check_visible_pd_panel=check_visible_pd_panel, check_visible_pdf=check_visible_pdf, 
                     screen=p, pd=pd, pd_line=pd_line, 
                     screen_xaxis=p.xaxis, pd_xaxis=p.xaxis, label_screen=label_screen
                    )

# # # Theoretical pdf -- switch according to experiment set up
callback_Teorpdf = CustomJS(args=js_cb_objects, code="""
    
    const max_x = source.tags[0].max_x
    const slitdist = slider_slitdist.value
    
    // slit distance to be displayed on the screen (not realistic)
    const slitdist_adjusted = slider_slitdist.value/source.tags[0].slitdist_factor 
    
    // variance adjusted to approximate the one slit diffraction peak by gaussian distribution 
    // původně source_pd.tags[0].var_x
    let var_gaussdiff = (slider_wavelength.value*source.tags[0].L)/(Math.PI*source.tags[0].slitwidth)

    const wavelength = slider_wavelength.value
    const envfreq = Math.PI*slitdist/(wavelength*source.tags[0].L)
    
    // Normalization constant for the overall hist+pdf when button_keep active (not a pure distribution)
    let keep_norm;
    if (button_keep.active.length==0) {
        keep_norm = 1
    }
    else { // ratio of particles from the actual distribution and all particles
        keep_norm = (source.data.x.length - source_pd.tags[0].previous_num)/source.data.x.length
        if (keep_norm==0) { // set nonzero normalization constatnt when no new particles added
            keep_norm = 1
        }
    }    
    console.log("keep_norm:", keep_norm)
    
    // Normalization  constant for theoretical pdf: Young interference on the interval (-max_x; max_x)
     const Norm = 2*envfreq/(envfreq*2*max_x + Math.sin(envfreq*2*max_x))
    
    
    // Creating x.linspace(-max_x, max_x, 500)
    if (source_pd.data.x.length==0) {
        let step = 2*max_x/500;
        for (let i = 0; i <= 500; i++) {
             source_pd.data.x.push(i*step - max_x)}
    }
    
    const x = source_pd.data.x
        
    let y;
    
    if (slit.active.length==2) {
        if (detector.button_type == "primary") { // Sum of both slits normal distributions  
            y = Array.from(x, (x) => keep_norm * (1/(Math.sqrt(2*Math.PI)*var_gaussdiff))/2 * 
                                      Math.exp(-0.5*((x + slitdist_adjusted)/var_gaussdiff)**2) + 
                                     keep_norm * (1/(Math.sqrt(2*Math.PI)*var_gaussdiff))/2 * 
                                      Math.exp(-0.5*((x - slitdist_adjusted)/var_gaussdiff)**2)   )
        }
        else{ // Young interference
              y = Array.from(x, (x) => keep_norm * Norm * (Math.cos(envfreq*x)**2) )
              console.log(envfreq, y)
        }
    }
    else if (slit.active.length==1) { 
    
        if (slit.active[0]==0) { // Left slit normal distribution
            y = Array.from(x, (x) =>  keep_norm * (1/(Math.sqrt(2*Math.PI)*var_gaussdiff)) * 
                                      Math.exp(-0.5*((x + slitdist_adjusted)/var_gaussdiff)**2) )
        }
        else { // Right slit normal distribution
            y = Array.from(x, (x) =>  keep_norm * (1/(Math.sqrt(2*Math.PI)*var_gaussdiff)) * 
                                      Math.exp(-0.5*((x - slitdist_adjusted)/var_gaussdiff)**2) )
        }
    }
    else { // Both slits blocked, nothing
        y = Array.from(x, (x) => 0 )
    }
           

    source_pd.data = { x, y }
    source_pd.change.emit()
""")

callback = CustomJS(args=js_cb_objects, code=js_functions+"""
    send(generator())
    """)

callback100 = CustomJS(args=js_cb_objects, code=js_functions+"""   
    for (let i = 0; i < 100; i++) {
        send(generator(), false)
    }
    normalizehist()
    // console.log(selected.active)
    console.log(source_pd.tags[0].previous_num)
    """)

callback_reset = CustomJS(args=js_cb_objects, code=js_functions+"""  
    reset()
    """)

callback_stream = CustomJS(args=js_cb_objects, code=js_functions+"""
    function sendparticle() {
        send(generator())
    }
      
    if (streamtoggle.button_type == "success") {
        const interval = setInterval(sendparticle, 1000);
        source.interval = interval;
        streamtoggle.button_type = "danger";
        streamtoggle.label = "Stop stream";
        
    } else {
        clearInterval(source.interval);
        streamtoggle.button_type = "success";
        streamtoggle.label = "Send continuous stream";
    }
""")

# Block/open slits, change button color
callback_choose_slit = CustomJS(args=js_cb_objects, code=js_functions+ """     
    
   // update_scale()
    
    // If Keep not checked, reset after changing slit set up
    if (button_keep.active.length==0) {
        reset()
    }
    else { // save actual number of particles sent from this distribution
        source_pd.tags[0].previous_num = source.data.x.length
        console.log(source_pd.tags[0].previous_num)
    }
""")

# Switch detector active/no detector, change button color
callback_detector = CustomJS(args=js_cb_objects, code=js_functions+"""
    
    // Change detector colour  
    if (detector.button_type == "default") {
        detector.button_type = "primary";
        detector.label = "Detector active";
        
    } else {
        detector.button_type = "default";
        detector.label = "No detector present";
    }
    
  //  update_scale()
    
    // If Keep not checked, reset after switching on detector
    if (button_keep.active.length==0) {
        reset()
    }
    else { // save actual number of particles sent from this distribution
        source_pd.tags[0].previous_num = source.data.x.length
        console.log(source_pd.tags[0].previous_num)
    }
""")

callback_change_slider = CustomJS(args=js_cb_objects, code=js_functions+ """     
    
    // Slider labels
    label_wavelength.text = "Wavelenght: <b>" + slider_wavelength.value + "</b> [nm]"
    label_energy.text = "Energy: <b>" + ((6.626 * 3)/(slider_wavelength.value)/1.602*100).toFixed(1) + "</b> [eV]"
    label_slitdist.text = "Slit distance: <b>" + slider_slitdist.value + "</b> [nm]"
    
    // If Keep not checked, reset after changing slider set up
    if (button_keep.active.length==0) {
        reset(false)
    }
    
""")

# Choose partice colour
callback_choose_color = CustomJS(args=js_cb_objects, code=js_functions+"""
    if (color_button.active==0) {
       source.tags[0].color = "red"
       console.log("red")
    }
    else if (color_button.active==1) {
        source.tags[0].color = "#0073e6" // "#ffc857"
        console.log("blue")
    }
    else {
        source.tags[0].color = "#2eb82e" //"#2b9720"
        console.log("green")
    }
    console.log(color_button.active)
   // source.change.emit()
""")

callback_keep_text = CustomJS(args=js_cb_objects, code=js_functions+"""
    
    if (button_keep.active.length==0) {
        text_keep_previous.visible = false
    }
    else {
       text_keep_previous.visible = true
    }
""")

callback_visible_pd_panel = CustomJS(args=js_cb_objects, code=js_functions+"""
    
    if (check_visible_pd_panel.active.length==0) {
       pd.visible = false
    }
    else {
       pd.visible = true
    }
""")

callback_visible_pdf = CustomJS(args=js_cb_objects, code=js_functions+"""
    
    if (check_visible_pdf.active.length==0) {
       pd_line.visible = false
    }
    else {
       pd_line.visible = true
    }
""")
# -------------------------------------------------------------------------------
# Callbacks to objects

button.js_on_event('button_click', callback, callback_Teorpdf)
button100.js_on_event('button_click', callback100, callback_Teorpdf)
resetbutton.js_on_event('button_click', callback_reset)
streamtoggle.js_on_event('button_click', callback_stream, callback_Teorpdf)

button_choose_slit.js_on_click(callback_choose_slit)
button_choose_slit.js_on_click(callback_Teorpdf)
button_detector.js_on_event('button_click', callback_detector, callback_Teorpdf)

slider_slitdist.js_on_change('value', callback_Teorpdf, callback_change_slider)
slider_slitwidth.js_on_change('value', callback_Teorpdf, callback_change_slider)
slider_wavelength.js_on_change('value', callback_Teorpdf, callback_change_slider)
color_button.js_on_click(callback_choose_color)

check_visible_pd_panel.js_on_click(callback_visible_pd_panel)
check_visible_pdf.js_on_click(callback_visible_pdf)

layout = row(column(pd, p, label_screen), 
         column(button, button100, resetbutton, streamtoggle, 
               button_choose_slit, button_detector, #  button_change_dist, 
               label_slitdist, slider_slitdist, # slider_slitwidth
               label_wavelength, slider_wavelength, label_energy,  
               color_button,  
               check_visible_pd_panel, check_visible_pdf, button_keep, text_keep_previous))

output_file(filename="Applet_doubleslit88888.html", title="Applet double slit experiment")
show(layout)