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

# Shows a window in the upper left part with real-time information 
# on the status of the jobs running on the IBM Q device
#%qiskit_job_watcher

# We use ibmq_vigo
IBMQ.load_account()
backend = IBMQ.get_provider(hub='ibm-q', group='open', project='main').get_backend('ibmq_vigo')

# Local simulator and vector simulator
simulator = Aer.get_backend('qasm_simulator')

from amplitude_damping import *

SHOTS = 200 # 8192

# The values for R and corresponding times
R_values = [0.2, 100.0, 200.0, 400.0]
#R_values = [r*.05 for r in range(1,int(1/.05))] + [r * .4 for r in range(10,25)] + [r * 4. for r in range(10,100+1)]
R_values = np.array(R_values)

npoints = 400

t_values = {}
for idx, R in enumerate(R_values):
    t_values[str(idx)] = np.linspace(0.0, 6.0 * np.pi / np.sqrt(abs(2.0 * R - 1.0)), npoints)
#t_values = np.linspace(0.0,5.0,npoints)
    
# We create the quantum circuits
q = QuantumRegister(5, name="q")
c = ClassicalRegister(1, name="c")

## Indices of the system and environment qubits
sys = 1
env = 2

## For values of R and thirty values of t for each
circuits = {}
for idx, R in enumerate(R_values):
    circuits[R] = []
    for t in t_values[str(idx)]:
        circuits[R].append(initial_state(q, sys)
                           +amplitude_damping_channel(q, c, sys, env, R, t))

# Calibration circuits
cal_circuits, state_labels = complete_meas_cal([sys], q, c)

circuits[0.2][1].draw(output='mpl')

# Execute the circuits on the local simulator
jobs_sim = {}
for R in R_values:
    jobs_sim[R] = execute(circuits[R], backend = simulator, shots = SHOTS)


# Analyse the outcomes
populations_sim = {}
for idx, R in enumerate(R_values):
    populations_sim[str(idx)] = []  # changed to str(i) to comply with bokeh bullshit
    current_job_res = jobs_sim[R].result()  
    for i in range(npoints):
        counts = current_job_res.get_counts(i)
        sm = 0.0
        if '1' in counts:
            sm += counts['1']/float(SHOTS)
        populations_sim[str(idx)].append(sm) # changed to str(i) to comply with bokeh bullshit

# Plot the results
fig_idx = 221
plt.figure(figsize=(10,12))
for idx, R in enumerate(R_values):
    plt.subplot(fig_idx)
    plt.plot(t_values[str(idx)], populations_sim[str(idx)]) # changed to str to comply with bokeh bullshit
    plt.xlabel('t')
    plt.ylabel('Population')
    fig_idx += 1
    plt.grid()

from bokeh.layouts import row, column
from bokeh.models import ColumnDataSource, Slider, CustomJS, Text, DataRange1d, Title
from bokeh.plotting import Figure, show
from bokeh.io import output_notebook 

R_values_str = [str(i) for i in range(len(R_values))]

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

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


ys = populations_sim
rs = {R_values_str[i] : [R_values_str_R[i]] for i,_ in enumerate(R_values)} 
ts = t_values

initial_r = R_values_str[-1]

initial_r

rs

# 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)


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

# Add text
text_source = ColumnDataSource({'r_value': [rs[initial_r]]})
r_available = ColumnDataSource(data=rs)
text = Text(x=0, x_offset=315, y=.8, 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,
              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;
        
        // 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];
                
        // text
        text_source.data = {'r_value': [String(r_available.data[r_idx])]};
        
        // Update the plot
        source_visible.change.emit();
        
    """)

layout = column(plot,slider)

output_notebook()

show(layout)

# Run the calibration job
calibration_job = execute(cal_circuits, backend, shots=SHOTS)

# Run the circuits and save the jobs
jobs = {}
jobs_data = []
for R in R_values:
    jobs[R] = execute(circuits[R], backend = backend, shots = SHOTS)
    job_data = {'jobid': jobs[R].job_id(),
               'description': 'Amplitude damping channel for R = '+str(R),
               'metadata': {'t_values': list(t_values[R]), 'R': R}}
    jobs_data.append(job_data)

experiment_data = [{
    "backend": backend.name(),
    "calibration": calibration_job.job_id(),
    "description": "Circuits for the simulation of the amplitude damping channel",
    "jobs": jobs_data
    }]

filename = 'amplitude_damping_{}.json'.format(
            datetime.now().strftime(("%Y_%m_%d-%H_%M")))

with open(filename,'w') as file:
    json.dump(experiment_data, file)

# List the available experiment files
import glob

print("Available experiment files:")
for f in glob.glob('*.json'):
    print(f)

# Load the experiment file
filename = "amplitude_damping_2019_11_18-17_15.json"

with open(filename, 'r') as file:
    experiment_data = json.load(file)
print(experiment_data[0]['description'])
print("Run on", experiment_data[0]['backend'])

# Get the backend
backend = IBMQ.get_provider(hub='ibm-q-university', group='turku', project='main').get_backend(experiment_data[0]['backend'])

# Use the calibration job to implement the error mitigation
calibration_job = backend.retrieve_job(experiment_data[0]['calibration'])
meas_fitter = CompleteMeasFitter(calibration_job.result(), state_labels)
meas_filter = meas_fitter.filter

# Analyse the outcomes
t_values = {}

jobs = {}
for job in experiment_data[0]['jobs']:
    R = job['metadata']['R']
    t_values[R] = np.array(job['metadata']['t_values'])
    jobs[R] = backend.retrieve_job(job['jobid'])

populations = {}
for R in jobs:
    populations[R] = []
    current_job_res = jobs[R].result()
    for i in range(npoints):
        counts = meas_filter.apply(current_job_res).get_counts(i)
        sm = 0.0
        if '1' in counts:
            sm += counts['1']/float(SHOTS)
        populations[R].append(sm)

# Plot the results
fig_idx = 221
plt.figure(figsize=(10,12))
for R in R_values:
    plt.subplot(fig_idx)
    plt.plot(t_values[R], populations[R], label='Experiment')
    plt.plot(t_values[R], populations_sim[R], label='Simulation')
    plt.xlabel('t')
    plt.ylabel('Population')
    fig_idx += 1
    plt.grid()
plt.legend();

# Compute the channel capacity
Qa_exp = {}
Qa_sim = {}
t_sim = np.linspace(0, 0.7, 200)
for R in populations:
    Qa_exp[R] = []
    for i in range(len(t_values[R])):
        Qa_exp[R].append(Qa(populations[R][i]/populations[R][0]))
    Qa_sim[R] = []
    for t in t_sim:
        Qa_sim[R].append(Qa(c1(R, t)**2))

# Plot the results
plt.figure(figsize=(7,8))
for i, R in enumerate(R_values[1:]):
    c = ['r', 'g', 'b'][i]
    plt.plot(t_values[R], Qa_exp[R], marker = 's', lw = 0, c = c, label=r'R = ' + str(R))
    plt.plot(t_sim, Qa_sim[R], c = c, ls = '--')
plt.xlim(0., 0.7)
plt.ylim(0., 1.)
plt.xlabel('t')
plt.ylabel(r'$Q_a$')
plt.legend();
plt.grid()

jobs[0.2].properties().qubits

jobs[0.2].properties().gates

