#Import general libraries (needed for functions)
import numpy as np
import matplotlib.pyplot as plt
#from IPython import display

#Import Qiskit classes
import qiskit
from qiskit.tools.monitor import job_monitor
from qiskit import Aer
from qiskit.providers.aer.noise import NoiseModel
from qiskit import  QuantumRegister, QuantumCircuit
#Import the RB Functions
import qiskit.ignis.verification.randomized_benchmarking as rb

# import the bayesian packages
import pymc3 as pm
import arviz as az
from scipy.optimize import curve_fit
import bayesian_fitter as bf

# initialize the Bayesian extension
%config InlineBackend.figure_format = 'retina'
# Initialize random number generator
RANDOM_SEED = 8927
np.random.seed(RANDOM_SEED)
az.style.use("arviz-darkgrid")

RB_process = "2_Q RB"

if RB_process in ["3_Q RB","2-3_Q RB"] :
    #Number of qubits
    nQ = 3
    #There are 3 qubits: Q0,Q1,Q2.
    #2Q RB on Q0,Q2 and 1Q RB on Q1
    rb_pattern = [[0,1],[2]] # because 3 qubits
    #Do three times as many 1Q Cliffords    
    length_multiplier = [1,3]
    #Interleaved Clifford gates (2-qubits and 1-qubit)
    interleaved_gates = [['cx 0 1'],['x 2']]

else:
    #Number of qubits
    nQ = 2
    #There are 2 qubits: Q0,Q1.
    #2Q RB Q0,Q1
    rb_pattern = [[0,1]]
    length_multiplier = 1
    interleaved_gates = [['cx 0 1']]    

#Number of Cliffords in the sequence (start, stop, steps)
nCliffs = [1, 20, 40, 60, 80, 100, 150, 200, 250, 300]
#Number of seeds (random sequences)
nseeds=8

scale = (2 ** len(rb_pattern[0]) - 1) / (2 ** len(rb_pattern[0]))

qregs_02 = QuantumRegister(2)
circ_02 = QuantumCircuit(qregs_02, name='circ_02')
#circ_02.h(qregs_02[0]) # booptrap! WIP!
circ_02.cx(qregs_02[0], qregs_02[1])
circ_02.draw()

qregs_1 = QuantumRegister(1)
circ_1 = QuantumCircuit(qregs_1, name='circ_1')
circ_1.x(qregs_1[0]) # booptrap! WIP!
circ_1.draw()

rb_opts = {}
rb_opts['rand_seed'] = 61946
rb_opts['length_vector'] = nCliffs
rb_opts['nseeds'] = nseeds
rb_opts['rb_pattern'] = rb_pattern
rb_opts['length_multiplier'] = length_multiplier
#rb_opts['align_cliffs'] = True
if RB_process in ["3_Q RB","2-3_Q RB"]:
    rb_opts['interleaved_elem'] = [circ_02, circ_1]
if RB_process == "2_Q RB":
     rb_opts['interleaved_elem'] = [circ_02]
rb_original_circs, xdata, rb_interleaved_circs = rb.randomized_benchmarking_seq(**rb_opts)

#Original RB circuits
print (rb_original_circs[0][0])

#Interleaved RB circuits
print (rb_interleaved_circs[0][0])

from qiskit import IBMQ
IBMQ.load_account()
provider = IBMQ.get_provider(hub='ibm-q')
device = provider.get_backend('ibmq_lima')
properties = device.properties()
coupling_map = device.configuration().coupling_map
backend = device
hardware = device.name()
shots = 2**9
basis_gates = ['id', 'rz', 'sx', 'x', 'cx', 'reset']
noise_model = None

epc_calib = 6.401E-3 # as noted at experiment time

retrieve_list = []

original_result_list, original_transpile_list = bf.get_and_run_seeds(rb_circs=rb_original_circs,
                                               shots=shots,
                                               backend = backend,
                                               coupling_map = coupling_map,
                                               basis_gates = basis_gates,
                                               noise_model = noise_model,
                                               retrieve_list=retrieve_list)                                                

retrieve_list = []

interleaved_result_list, interleaved_transpile_list = bf.get_and_run_seeds(rb_circs=rb_interleaved_circs,
                                               shots=shots,
                                               backend = backend,
                                               coupling_map = coupling_map,
                                               basis_gates = basis_gates,
                                               noise_model = noise_model,
                                               retrieve_list=retrieve_list)                                                

Y1 = bf.get_count_data(original_result_list, nCliffs=nCliffs )
Y2 = bf.get_count_data(interleaved_result_list, nCliffs=nCliffs)

# function to optimize
def lsf(x, a, alpha, b):
    return a * alpha ** x + b

# curve fit
popt_s,pcov_s = curve_fit(lsf, np.array(nseeds*list(nCliffs)),
                          np.ravel(Y1)/shots,
                          bounds = ([.65,.9,0.15],[.85,.999,0.35]))
                                
perr_s= np.sqrt(np.diag(pcov_s))
# get EPC and EPC sigma for LSF accelerated
alpha_f = popt_s[1] 
alpha_f_err = perr_s[1] 
popt_s,perr_s

# curve fit
popt_i,pcov_i = curve_fit(lsf, np.array(nseeds*list(nCliffs)),
                          np.ravel(Y2)/shots,
                          bounds = ([.65,.9,0.15],[.85,.999,0.35]))
                                
perr_i= np.sqrt(np.diag(pcov_i))
# get EPC and EPC sigma for LSF accelerated
alphC_f = popt_i[1] 
alphC_f_err = perr_i[1] 
popt_i,perr_i

epc_est_f = scale*(1 - alphC_f/alpha_f)
epc_est_f_err = scale*(alphC_f/alpha_f)*(np.sqrt(alpha_f_err**2 + alphC_f_err**2))

# function to optimize
def lsmf(x, a, alpha, p_tilde_m, b):
    return x[1]*(a * alpha ** x[0] + b) + x[2]*(a * (alpha*p_tilde_m) ** x[0] + b)

# obtain the data
m_len = len(nCliffs)*nseeds
x0_lsmf = np.array(nseeds*2*list(nCliffs))
x1_lsmf = np.hstack((np.ones(m_len),np.zeros(m_len)))
x2_lsmf = np.hstack((np.zeros(m_len),np.ones(m_len)))
x_lsmf = np.vstack((x0_lsmf,x1_lsmf,x2_lsmf))
y_lsmf=np.hstack((np.ravel(Y1),np.ravel(Y2)))/shots

# curve fit
popt_m,pcov_m = curve_fit(lsmf, x_lsmf, y_lsmf,
                    bounds = ([.65,.9,.9,0.15],
                                [.85,.999,.999,0.35]))
perr_m = np.sqrt(np.diag(pcov_m))
# get EPC and EPC sigma for LSF accelerated
alpha_fm = popt_m[1] 
p_tilde_m =  popt_m[2] 
alpha_fm_err = perr_m[1] 
p_tilde_m_err =  perr_m[2] 
popt_m,perr_m

epc_est_fm = scale*(1 - p_tilde_m)
epc_est_fm_err = scale*p_tilde_m_err

original_model = bf.get_bayesian_model(model_type="pooled",Y=Y1,shots=shots,m_gates=nCliffs,
                                       mu_AB=[popt_s[0],popt_s[2]],cov_AB=[perr_s[0],perr_s[2]],
                                       alpha_ref=alpha_f)

pm.model_to_graphviz(original_model)

trace_o = bf.get_trace(original_model)

azo_summary = bf.get_summary(original_model, trace_o)
azo_summary

alpha_original_p = azo_summary['mean']['alpha']
alpha_original_p_err = azo_summary['sd']['alpha']

interleaved_model = bf.get_bayesian_model(model_type="pooled",Y=Y2,shots=shots,m_gates=nCliffs,
                                          mu_AB=[popt_i[0],popt_i[2]],cov_AB=[perr_i[0],perr_i[2]],
                                          alpha_ref=alphC_f)
pm.model_to_graphviz(interleaved_model)

trace_i = bf.get_trace(interleaved_model)

azi_summary = bf.get_summary(interleaved_model, trace_i)
azi_summary

alpha_c_p = azi_summary['mean']['alpha']
alpha_c_p_err = azi_summary['sd']['alpha']

epc_est_p = scale*(1 - alpha_c_p/alpha_original_p)
epc_est_p_err = scale*(alpha_c_p/alpha_original_p)*(np.sqrt(alpha_original_p_err**2 + alpha_c_p_err**2))

Y = np.vstack((Y1,Y2))
RvsI = np.vstack((np.ones_like(Y1),np.zeros_like(Y2)))
IvsR = np.vstack((np.zeros_like(Y1),np.ones_like(Y2)))

tilde =bf.get_bayesian_model("tilde",Y=Y,shots=shots, m_gates=nCliffs,
                             alpha_ref=alpha_fm, p_testval= p_tilde_m,
                             mu_AB=[popt_m[0],popt_m[3]],cov_AB=[perr_m[0],perr_m[3]],
                             RvsI=RvsI,IvsR=IvsR)
pm.model_to_graphviz(tilde)

trace_t = bf.get_trace(tilde)

azt_summary = bf.get_summary(tilde, trace_t)
azt_summary

epc_est_a = scale*(1 - azt_summary['mean']['p_tilde'])
epc_est_a_err = scale* (azt_summary['sd']['p_tilde'])

# compare LSF and SMC
print("Model:           Frequentist                     Bayesian            Calibration")
print("            two-run      accelerated       two-run      accelerated ")
print("EPC         {0:.5f}       {1:.5f}          {2:.5f}       {3:.5f}      {4:.5f}  "
      .format(epc_est_f, epc_est_fm, epc_est_p, epc_est_a, epc_calib))  
print("± sigma   ± {0:.5f}       ± {1:.5f}        ± {2:.5f}     ± {3:.5f}     ------   "
      .format(epc_est_f_err, epc_est_fm_err, epc_est_p_err, epc_est_a_err))

# obtain EPC from alpha (used by plot_posterior)
def alpha_to_EPC(alpha):
        return scale*(1-alpha)

# for refering the interleaved gate in the title of the graphs
intl_g=str(interleaved_gates[0][0][0:2])+str(rb_pattern[0][0:2])
if RB_process in ["3_Q RB","2-3_Q RB"] :
    intl_g=intl_g+"<"+str(interleaved_gates[1][0][0:1]+str(rb_pattern[1][0:2]))


import matplotlib.pyplot as plt # seems we need to reimport for replot WIP

with tilde:
        ax = az.plot_posterior(trace_t,  var_names=['p_tilde'], round_to=4, point_estimate=None, 
                          transform =  alpha_to_EPC)
        ax.set_xlim(0.0, 0.010)
        plt.axvline(x=epc_est_fm,color='red',ls=":")
        plt.axvline(x=epc_est_p,color='orange',ls="-")
        plt.axvline(x=epc_est_f,color='cyan',ls="-")       
        if epc_calib > 0.0:
            plt.axvline(x=epc_calib,color='green',ls=":")
        plt.axvline(x=epc_est_a,color='blue',ls=":")
        plt.title(RB_process +' $accelerated$, gate: ' + intl_g\
                  +", "+hardware+', backend: '+backend.name(),
               fontsize=12)
        Bayes_legend =  "EPC Accelerated SMC: {0:1.3e} ({1:1.3e})".format(epc_est_a, epc_est_a_err)
        Bayes2_legend = "EPC SMC 2-runs: {0:1.3e} ({1:1.3e})".format(epc_est_p, epc_est_p_err)
        Fitter_legend = "EPC LSF 2-runs: {0:1.3e} ({1:1.3e})".format(epc_est_f, epc_est_f_err) 
        LSM_legend =    "EPC Accelerated LSF: {0:1.3e} ({1:1.3e})".format(epc_est_fm, epc_est_fm_err)  
        Cal_legend =    "EPC Calibration: {0:1.3e}".format(epc_calib)
        if epc_calib > 0.0:   
            plt.legend((Bayes_legend, "$Higher\; density\; interval$ HDI",
                        LSM_legend,Bayes2_legend,
                        Fitter_legend,Cal_legend), fontsize=10 )
        else: 
            plt.legend((Bayes_legend, "$Higher\; density\; interval$ HDI",
                        LSM_legend, Bayes2_legend,
                        Fitter_legend), fontsize=10 )

import matplotlib.pyplot as plt # seems we need to reimport for replot WIP

fig, plt = plt.subplots(1, 1)

plt.set_ylabel("Ground State Population")
plt.set_xlabel("Number of Cliffords")

for i_seed in range(nseeds):
    plt.scatter(nCliffs, Y1[i_seed,:]/shots, label = "data", marker="x",color="b")
    plt.scatter(nCliffs, Y2[i_seed,:]/shots, label = "data", marker="+",color="r")
plt.plot(nCliffs,azt_summary['mean']['AB[0]']*azt_summary['mean']['alpha']**nCliffs+\
          azt_summary['mean']['AB[1]'],'--',color="b")
plt.plot(nCliffs,azt_summary['mean']['AB[0]']*(azt_summary['mean']['alpha']*azt_summary['mean']['p_tilde'])**\
          nCliffs+azt_summary['mean']['AB[1]'],'--',color="r")

plt.legend(("Standard",
             "Interleaved"))

plt.set_title(RB_process +' SMC $accelerated$, gate: ' + intl_g\
                  +", "+hardware+', backend: '+backend.name(),
               fontsize=14);

import qiskit.tools.jupyter
%qiskit_version_table