import numpy as np
import time
from qiskit_experiments.library import StandardRB, InterleavedRB
from qiskit_experiments.framework import ParallelExperiment
from qiskit_experiments.library.randomized_benchmarking import RBUtils
import qiskit.circuit.library as circuits

# for retrieving gate calibration
from datetime import datetime
import qiskit.providers.aer.noise.device as dv

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

simulation = True # make your choice here
if simulation:
    from qiskit.providers.aer import AerSimulator
    from qiskit.test.mock import FakeParis
    backend = AerSimulator.from_backend(FakeParis())
else:
    from qiskit import IBMQ
    IBMQ.load_account()
    provider = IBMQ.get_provider(hub='ibm-q')
    backend = provider.get_backend('ibmq_bogota') # type here hardware backend  

# describe RB experiment
interleaved_gate = "cx"
qubits = [1,4]
interleaved_circuit = circuits.CXGate()
lengths = np.arange(1, 200, 15) 
num_samples = 10
seed = 1010

# get the backend's calibration value
t = None # enter t in datetime format if necessary
# use properties(datetime=t) if t is defined
e_list = dv.gate_error_values(backend.properties()) 
epc_calib = np.nan
for tuple_e in e_list:
    if tuple_e[0] == 'cx' and tuple_e[1] == qubits:
        epc_calib = tuple_e[2]
print('EPC calibration: {0:1.4e}'.format(epc_calib)) 

#prepare circuits
int_exp2 = InterleavedRB(interleaved_circuit, qubits,
                            lengths, num_samples=num_samples, seed=seed)
#run
print("start experiments",time.strftime('%d/%m/%Y %H:%M:%S'))
int_expdata2 = int_exp2.run(backend).block_for_results()
print("  end experiments",time.strftime('%d/%m/%Y %H:%M:%S'))

#analyse
print("   start analysis",time.strftime('%d/%m/%Y %H:%M:%S'))
int_results2 = int_expdata2.analysis_results()
print("     end analysis",time.strftime('%d/%m/%Y %H:%M:%S'))

# look at result data
for result in int_results2:
    print(result)

def get_GSP_counts(data, x_length, data_range):
#obtain the observed counts used in the bayesian model
#corrected for accomodation pooled data from 1Q, 2Q and 3Q interleave processes
    list_bitstring = ['0','00', '000', '100'] # all valid bistrings
    Y_list = []
    for i_samples in data_range:
        row_list = []
        for c_index in range(x_length) :  
            total_counts = 0
            i_data = i_samples*x_length + c_index
            for key,val in data[i_data]['counts'].items():
                if  key in list_bitstring:
                    total_counts += val
            row_list.append(total_counts)
        Y_list.append(row_list)    
    return np.array(Y_list)

# get count data and other values from int_expdata2
Y1 = get_GSP_counts(int_expdata2._data, len(lengths),
                    range(0,2*num_samples-1,2))
Y2 = get_GSP_counts(int_expdata2._data, len(lengths),
                    range(1,2*num_samples,2))
experiment_type = int_expdata2._data[0]['metadata']['experiment_type']
physical_qubits = int_expdata2._data[0]['metadata']['physical_qubits']
shots = int_expdata2._data[0]['shots']
nQ = len(qubits)
scale = (2 ** nQ - 1) / 2 ** nQ
# to compare ultimately: EPC and sigma(EPC) by LSF
epc_est_fm = int_expdata2.analysis_results()[3].value.value
epc_est_fm_err = int_expdata2.analysis_results()[3].value.stderr

# use 2m length array
Y = np.hstack((Y1,Y2))
RvsI_h = np.ravel(np.vstack((np.ones_like(lengths),np.zeros_like(lengths))))
IvsR_h = np.ravel(np.vstack((np.zeros_like(lengths),np.ones_like(lengths))))
X0 = np.tile(lengths,2)
X = np.vstack((X0,RvsI_h,IvsR_h))

y_mean = np.mean(Y, axis = 0)/shots
sigma_y = np.std(Y, axis = 0)/shots

model = "hierarchical_model"

# priors for unknown model parameters
T_priors = int_expdata2.analysis_results()[0].value.value
print(T_priors)

# building the model
h_model = pm.Model()
with h_model:
    
    # Tying parameters
    BoundedUniform = pm.Bound(pm.Uniform,
                               lower=np.fmax(T_priors-0.1, np.full(T_priors.shape,1.e-9)),
                               upper=np.fmin(T_priors+0.1, np.full(T_priors.shape,1.-1e-9)))
    pi = BoundedUniform("Tying_Parameters",testval = T_priors, shape = T_priors.shape)
    EPC = pm.Deterministic('EPC', scale*(1-pi[2])) 
    
    # sigma of Beta functions
    sigma_t = pm.Uniform("σ_Beta",  testval = 0.005,
                            upper = 0.05, lower = 0.0005)    
    
    # Tying function
    GSP = pi[0] * ( X[1]*pi[1]**X[0] +\
                    X[2]*(pi[1]*pi[2])**X[0] ) +  pi[3]
                 
    theta = pm.Beta('θ', mu=GSP, sigma = sigma_t,
                     shape = ((2*len(lengths, ))) )                                 

    # Likelihood (sampling distribution) of observations    
    p = pm.Binomial("Counts", p = theta, observed = Y, n = shots) 

# model graph                                
pm.model_to_graphviz(h_model)

# sample
with h_model:
    trace_h = pm.sample(draws = 4000, tune= 1000, target_accept=.99,
                         return_inferencedata=True)    

with h_model:
        az.plot_trace(trace_h);

with h_model:
    az.plot_posterior(trace_h, var_names = ["Tying_Parameters","σ_Beta","EPC"], round_to = 4, figsize = [16, 8]);

# look at the posterior values of the hyperparameters:
with h_model:
    #  (hdi_prob=.94 is default)
    azt_summary = az.summary(trace_h, round_to=12,
                            var_names = ["Tying_Parameters", "σ_Beta","EPC"],
                            kind="stats") 
azt_summary

# mean and sigma of EPC
epc_est_a = azt_summary['mean']['EPC']
epc_est_a_err = azt_summary['sd']['EPC']

# plot 
import matplotlib.pyplot as plt # if not yet imported 

with h_model:
    az.plot_posterior(trace_h, var_names = ["EPC"],
                      round_to = 4, figsize = [10,6],
                      textsize = 12) 

Bayes_legend =  "EPC  SMC: {0:1.3e} ± {1:1.3e}"\
                    .format(epc_est_a, epc_est_a_err)
LSF_legend =    "EPC  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)
plt.axvline(x=epc_est_fm,color='cyan',ls="-")    
if epc_calib != np.nan:
    plt.axvline(x=epc_calib,color='r',ls=":")
plt.axvline(x=epc_est_a,color='blue',ls=":")
if epc_calib > 0.0:   
    plt.legend((Bayes_legend, "$Highest\; density\; interval$ HDI",
                LSF_legend,
                Cal_legend), fontsize=12 )
else: 
    plt.legend((Bayes_legend, "$Highest\; density\; interval$ HDI",
                LSF_legend))
plt.title(experiment_type +', ' + interleaved_gate  + " qubit(s):" + str(physical_qubits)\
                     +', backend: '+backend.name() + "\n Bayesian "+model,
                     fontsize=16);

# compare LSF and SMC

print("Model:    Frequentist         Bayesian        Calibration")
print("__________________________________________________________")
print("EPC         {0:1.3e}        {1:1.3e}         {2:1.3e}"
      .format(epc_est_fm,epc_est_a,epc_calib ))  
print("± sigma    ± {0:1.3e}      ± {1:1.3e} "
      .format(epc_est_fm_err, epc_est_a_err))

def calc_chisquare(ydata, sigma, ycalc): 
    r = ydata - ycalc
    chisq = np.sum((r / sigma) ** 2)
    return chisq

# prepare box for GSP plot

# perform reduced χ² value calculation for Bayes hierarchical 
mean_h = trace_h.posterior.mean(dim=['chain', 'draw'])
theta_stacked = mean_h.θ.values
NDF_h = len(lengths)*2 - 4 - 1 # (-1 is for σ_Beta)
chisq_h = calc_chisquare(y_mean, sigma_y, theta_stacked)/NDF_h

texto_0 = "  alpha = {0:7.4f} ± {1:1.4e}"\
            .format(azt_summary['mean']['Tying_Parameters[1]'],
            azt_summary['sd']['Tying_Parameters[1]'])                                              

texto_1 =" alpha_c = {0:7.4f} ± {1:1.4e}"\
            .format(azt_summary['mean']['Tying_Parameters[2]'],
            azt_summary['sd']['Tying_Parameters[2]'])                                              

texto_2 = "   EPC = {0:7.4f} ± {1:1.4e}"\
            .format(azt_summary['mean']['EPC'],
            azt_summary['sd']['EPC'])

texto_3 = "             Fit χ² = {0:7.4f} "\
            .format(chisq_h)

texto = texto_0 + "\n" + texto_1 + "\n" + texto_2 + "\n" + texto_3

# prepare data for GSP plot
# get the calculated GSP values
with h_model:
    hdi_prob = .94
    #  (hdi_prob=.94 is default, roughly coreresponding to 2σ)
    theta_summary = az.summary(trace_h, round_to=12, hdi_prob = hdi_prob,
                            var_names = ["θ"], kind="stats")
y1 = theta_summary.values[:,0][0:len(lengths)]
y2 = theta_summary.values[:,0][len(lengths):len(lengths)*2]  

HDI = False # make your choice here
if HDI:
    # HDI values as bounds
    bounds_rmk = "(shown bounds are "+ str(int(100*hdi_prob)) + "% HDI)"
    y1_min = theta_summary.values[:,2][0:len(lengths)]
    y2_min = theta_summary.values[:,2][len(lengths):len(lengths)*2]
    y1_max = theta_summary.values[:,3][0:len(lengths)]
    y2_max = theta_summary.values[:,3][len(lengths):len(lengths)*2]
else:    
    # two sigma bounds for plot
    bounds_rmk = "(shown bounds are ± two σ)"
    sy = theta_summary.values[:,1]
    y1_min = y1 - sy[0:len(lengths)]*2
    y1_max = y1 + sy[0:len(lengths)]*2
    y2_min = y2 - sy[len(lengths):len(lengths)*2]*2
    y2_max = y2 + sy[len(lengths):len(lengths)*2]*2    

# GSP plot
import matplotlib.pyplot as plt # if not yet imported
font = {'family' : 'DejaVu Sans',
        'weight' : 'normal',
        'size'   : 14}
plt.rc('font', **font)
fig, plt = plt.subplots(1, 1, figsize = [8,5])

plt.set_ylabel("P(0)")
plt.set_xlabel("Clifford Length")

plt.legend(("Standard" , "Interleaved" ),
          loc = 'center right', fontsize=10)

plt.plot(lengths,y1,color="purple", marker="o", lw = 0.75)
#plt.errorbar(lengths,y1,2*sy[0:len(lengths)],
             #color="purple", marker='o')
plt.fill_between(lengths, y1_min, y1_max,
                alpha=.2, edgecolor='purple', facecolor= 'r')
                
plt.plot(lengths,y2,color="cyan", marker='^', lw = 0.75)
#plt.errorbar(lengths,y2,2*sy[len(lengths):2*len(lengths)],
             #color="cyan", marker='^')
plt.fill_between(lengths, y2_min, y2_max,
                alpha=.2, edgecolor='cyan', facecolor= 'cyan')
                

for i_seed in range(num_samples):
    plt.scatter(lengths, Y1[i_seed,:]/shots,
               label = "data", marker="x",color="grey")
    plt.scatter(lengths, Y2[i_seed,:]/shots,
               label = "data", marker="+",color="grey")

plt.legend(("Standard" , "Interleaved" ),
          loc = 'center right', fontsize=10)


plt.text(lengths[-1]*0.3,0.75, texto, bbox=dict(facecolor='white'))
plt.grid()
plt.set_title(experiment_type +', ' + interleaved_gate\
              + str(physical_qubits)\
              +', backend: '+backend.name()+\
              "\n Bayesian "+model+" "+ bounds_rmk,
              fontsize=14);

# View result data for frequentist model
display(int_expdata2.figure(0))