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

#Import Qiskit classes
import qiskit

import qiskit_experiments as qe
rb = qe.randomized_benchmarking
from scipy.optimize import curve_fit

# import the bayesian packages
import pymc3 as pm
import arviz as az

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

# choice of "simulation, "real", "retrieve"
option = "retrieve"

# Determine the backend
if option == "simulation":
    from qiskit.test.mock import FakeAthens
    backend = FakeAthens()
    hardware = 'ibmq_athens' # hardware reference
else:
    from qiskit import IBMQ
    IBMQ.load_account()
    provider = IBMQ.get_provider(hub='ibm-q')
    device = provider.get_backend('ibmq_athens')
    backend = device
    hardware = device.name()  # hardware used  

# RB design
q_list = [0,1]
lengths = [1, 20, 40, 60, 80, 100, 150, 200, 250, 300, 350, 400, 450, 500]
num_samples = 5
seed = 1010 
num_qubits = len(q_list)
scale = (2 ** num_qubits - 1) / (2 ** num_qubits)
shots = 1024

if option != "retrieve":
    exp = rb.RBExperiment(q_list, lengths, num_samples=num_samples, seed=seed)
    eda= exp.run(backend)

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

# one or two qubits for retrieve and for the plots
list_bitstring = ['0', '00'] 

# obtain the current count values
Y_list = []

if option == "simulation":
    for i_sample in range(num_samples*len(lengths)):
        Y_list.append(eda.data[i_sample]['counts']\
                      [eda.data[i_sample]['metadata']['ylabel']])
        
else: # specify job (fresh job or run some time ago)
    job = backend.retrieve_job('6097aed0a4885edfb19508fa') # athens 01'

    for rbseed, result in enumerate(job.result().get_counts()):
            total_counts = 0
            for key,val in result.items():
                if  key in list_bitstring:
                    total_counts += val
            Y_list.append(total_counts)          
    
Y = np.array(Y_list).reshape(num_samples, len(lengths))

#get LSF EPC and priors
if option != "retrieve":
    popt = eda._analysis_results[0]['popt']
    pcov = eda._analysis_results[0]['pcov']

else: # manual entry (here for job ''6097aed0a4885edfb19508fa')
    popt = [0.7207075,  0.95899375, 0.2545933 ]
    pcov =  [[ 2.53455272e-04, -9.10001034e-06, -1.29839515e-05],
             [-9.10001034e-06,  9.55193998e-06, -7.07822420e-06],
             [-1.29839515e-05, -7.07822420e-06,  2.07452483e-05]]

alpha_ref=popt[1]     
mu_AB= [popt[0],popt[2]]
alpha_ref_err = np.sqrt(pcov[1][1])
EPC = scale*(1-alpha_ref)
EPC_err = scale*alpha_ref_err

cov_AB= [0.0001, 0.0001]
alpha_lower=0.8
alpha_upper=0.999                        
sigma_theta = .004

pooled_model = bf.get_bayesian_model(model_type="pooled",
                                        Y=Y,shots=1024,m_gates=lengths,
                                        alpha_ref =  alpha_ref,
                                        mu_AB=mu_AB,cov_AB=cov_AB)

pm.model_to_graphviz(pooled_model)

with pooled_model:   
    trace_p= pm.sample(draws = 2000, tune= 10000, target_accept=0.97,
                     return_inferencedata=True)    
    az.plot_trace(trace_p)

with pooled_model:
    azp_summary = az.summary(trace_p,  hdi_prob=.94, round_to=6, kind="all")
azp_summary

epc_p =scale*(1 - azp_summary['mean']['alpha'])
epc_p_err = scale* (azp_summary['sd']['alpha'])

with pooled_model:
        ax = az.plot_posterior(trace_p,  var_names=['alpha'], round_to=4, point_estimate=None, 
                          transform =  alpha_to_EPC, textsize = 10.0, color='b')
        ax.set_title("RB_process: standard "+str(q_list)+", backend: "+backend.name(),
               fontsize=12)
        Bayes_legend ="Bayesian Pooled Model:\n EPC  {0:1.3e} ± {1:1.3e}".format(epc_p, epc_p_err)
        ax.axvline(x=EPC,color='r',ls=":")
        Fitter_legend ="Frequentist model:\n EPC {0:1.3e} ± {1:1.3e}".format(EPC, EPC_err) 
        ax.legend((Bayes_legend, "$Highest\; density\; interval$ HDI",
                   Fitter_legend),fontsize=10 )                       

hierarchical_model = bf.get_bayesian_model(model_type="hierarchical",
                                        Y=Y,shots=1024,m_gates=lengths,
                                        alpha_ref =  alpha_ref,
                                        mu_AB=mu_AB,cov_AB=cov_AB)

pm.model_to_graphviz(hierarchical_model)

with hierarchical_model:   
    trace_h= pm.sample(draws = 2000, tune= 10000, target_accept=0.97,
                     return_inferencedata=True)    
    az.plot_trace(trace_h)

with hierarchical_model:
    azh_summary = az.summary(trace_h,  hdi_prob=.94, round_to=6,
                             kind="all", var_names=["~GSP"])
azh_summary

epc_h =scale*(1 - azh_summary['mean']['alpha'])
epc_h_err = scale* (azh_summary['sd']['alpha'])

with hierarchical_model:
        ax = az.plot_posterior(trace_h,  var_names=['alpha'], round_to=4, point_estimate=None, 
                          transform =  alpha_to_EPC, textsize = 10.0, color='b')
        ax.set_title("RB_process: standard "+str(q_list)+", backend: "+backend.name(),
               fontsize=12)
        Bayes_legend ="Bayesian Hierarchical Model:\n EPC  {0:1.3e} ± {1:1.3e}".format(epc_h, epc_h_err)
        ax.axvline(x=EPC,color='r',ls=":")
        Fitter_legend ="Frequentist model:\n EPC {0:1.3e} ± {1:1.3e}".format(EPC, EPC_err) 
        ax.legend((Bayes_legend, "$Highest\; density\; interval$ HDI",
                   Fitter_legend),fontsize=10 )                       

# compare LSF and SMC
print("Model:  Frequentist            Bayesian")
print("           LSF           pooled      hierarchical")
print("EPC      {0:.5f}         {1:.5f}       {2:.5f}"
      .format(EPC, epc_p, epc_h))
print("ERROR   ± {0:.5f}       ± {1:.5f}     ± {2:.5f}"
      .format(EPC_err, epc_p_err, epc_h_err))

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

#fig, plt = plt.subplots(1, 1, sharex=True, figsize=(10, 6))
fig, plt = plt.subplots(1, 1, sharex=True, figsize=(10, 6))

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

for i_seed in range(num_samples):
    plt.scatter(lengths, Y[i_seed,:]/1024, label = "data",
                marker="+",color="purple")
plt.plot(np.array(lengths)-0.5,azp_summary['mean']['AB[0]']\
         *azp_summary['mean']['alpha']**lengths+\
          azp_summary['mean']['AB[1]']-0.002,'o-',color="c")
plt.plot(np.array(lengths)+0.5,azh_summary['mean']['AB[0]']\
         *azh_summary['mean']['alpha']**\
          lengths+azh_summary['mean']['AB[1]']+0.002,'o-',color="orange")

plt.legend(("Pooled Model",
             "Hierarchical Model"))

plt.set_title("RB_process: standard "+str(q_list)+\
              ", device: "+hardware+', backend: '+backend.name(),
              fontsize=14);

# Leave-one-out Cross-validation (LOO) comparison
df_comp_loo = az.compare({"hierarchical": trace_h, "pooled": trace_p})
df_comp_loo

az.plot_compare(df_comp_loo, insample_dev=False);

import qiskit.tools.jupyter
%qiskit_version_table

