import numpy as np
import matplotlib.pyplot as plt

from datetime import datetime
import json
import copy

# Main qiskit imports
from qiskit import QuantumRegister, ClassicalRegister
from qiskit import QuantumCircuit, execute, Aer, IBMQ

# Error mitigation
from qiskit.ignis.mitigation.measurement import (complete_meas_cal, 
                                                CompleteMeasFitter, 
                                                MeasurementFilter)

# Utility functions
from qiskit.tools.jupyter import *
from qiskit.tools.monitor import job_monitor
from qiskit.providers.jobstatus import JobStatus

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 

import pickle

#pickle.dump([ys,rs,ts,initial_r],open("witness_pop.p", "wb"))
ys,rs,ts,initial_r = pickle.load(open( "witness_pop.p", "rb" ))
ys_witness,rs_witness,ts_witness,initial_r = pickle.load(open( "nonmark_witness.p", "rb" ))

initial_r = str(len(rs)//2)

# 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[initial_r], y = ys[initial_r]))

source_available_x = ColumnDataSource(data=ts)
source_available_y = ColumnDataSource(data=ys)


source_visible2 = ColumnDataSource(data=dict(
    x = ts[initial_r], y = ys_witness[initial_r]))
source_available2_x = ColumnDataSource(data=ts)
source_available2_y = ColumnDataSource(data=ys_witness)



# Define plot elements
plot = Figure(plot_width=300, plot_height=300, x_range = DataRange1d(), y_range=(-.01, 1.01))
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 = DataRange1d(), y_range=(-.01, 1.01))
plot2.line('x', 'y', source=source_visible2, legend_label="Non-Mark. Witness", 
           line_width=3, line_alpha=0.6, line_color="#f01001")


# Add text
text_source = ColumnDataSource({'r_value': [rs[initial_r]]})
r_available = ColumnDataSource(data=rs)
text = Text(x=0, x_offset=215, 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_x = source_available_x,
              source_available_y = source_available_y,
              source_visible2 = source_visible2,
              source_available2_x = source_available2_x,
              source_available2_y = source_available2_y,
              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_x = source_available_x.data;
        var data_available_y = source_available_y.data;
        var data_visible2 = source_visible2.data;
        var data_available2_x = source_available_x.data;
        var data_available2_y = source_available2_y.data;
        
        // Change y-axis data according to the selected value
        data_visible.x = data_available_x[r_idx];
        data_visible.y = data_available_y[r_idx];
        data_visible2.x = data_available_x[r_idx];
        data_visible2.y = data_available2_y[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("nonmark_witness.html", title="Non-Markovianity Witness")
output_notebook()
show(layout)

