#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.providers.aer.noise import NoiseModel
from qiskit.providers.aer.noise.errors.standard_errors import depolarizing_error, thermal_relaxation_error

#Import the RB Functions
import qiskit.ignis.verification.randomized_benchmarking as rb

import copy

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

from qiskit import IBMQ
IBMQ.load_account()
provider = IBMQ.get_provider(hub='ibm-q')
provider.backends()

from qiskit.tools.monitor import job_monitor

device = provider.get_backend('ibmq_lima')


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

def obtain_priors_and_data_from_fitter(printout = True):
    m_gates = copy.deepcopy(nCliffs)
    # We choose the count matrix corresponding to 2 Qubit RB
    Y = (np.array(rbfit._raw_data[0])*shots).astype(int)
    
    # alpha prior and bounds 
    alpha_ref = rbfit._fit[0]['params'][1]    
    alpha_lower = alpha_ref - 2*rbfit._fit[0]['params_err'][1] # modified for real
    alpha_upper = alpha_ref + 2*rbfit._fit[0]['params_err'][1] # modified for real
    
    # priors for A anbd B
    mu_AB = np.delete(rbfit._fit[0]['params'],1)
    cov_AB=np.delete(rbfit._fit[0]['params_err'],1)**2
    
    # prior for sigmatheta:
    sigma_theta = 0.004    
    if printout:
        print("priors:\nalpha_ref",alpha_ref)
        print("alpha_lower", alpha_lower, "alpha_upper", alpha_upper)
        print("A,B", mu_AB, "\ncov A,B", cov_AB)
        print("sigma_theta", sigma_theta)
    
    return m_gates, Y, alpha_ref, alpha_lower, alpha_upper, mu_AB, cov_AB, sigma_theta

def get_bayesian_model(model_type):
# Bayesian model
# from https://iopscience.iop.org/article/10.1088/1367-2630/17/1/013042/pdf 
# see https://docs.pymc.io/api/model.html
    
    RB_model = pm.Model()
    with RB_model:
        #Priors for unknown model parameters
        alpha = pm.Uniform("alpha",lower=alpha_lower,
                           upper=alpha_upper, testval = alpha_ref)
        
        BoundedMvNormal = pm.Bound(pm.MvNormal, lower=0.0)
        
        AB = BoundedMvNormal("AB", mu=mu_AB,testval = mu_AB,
                         cov= np.diag(cov_AB),
                         shape = (2))

        # Expected value of outcome
        GSP = AB[0]*alpha**m_gates + AB[1]
        
        if model_type == "pooled":
            total_shots = np.full(Y.shape, shots)
            theta = GSP
        
        elif model_type == "hierarchical":
            total_shots = np.full(Y.shape, shots)
            theta = pm.Beta("GSP",
                         mu=GSP,
                         sigma = sigma_theta,
                         shape = Y.shape[1])
        
        # Likelihood (sampling distribution) of observations    
        p = pm.Binomial("Counts", p=theta, observed=Y,
                            n = total_shots) 

    return RB_model

def get_trace(RB_model):
    # Gradient-based sampling methods
    # see also: https://docs.pymc.io/notebooks/sampler-stats.html
    # and https://docs.pymc.io/notebooks/api_quickstart.html
    with RB_model:   
        trace= pm.sample(draws = 2000, tune= 10000, target_accept=0.9, return_inferencedata=True)    

    with RB_model:
        az.plot_trace(trace);
        
    return trace

def get_summary(RB_model, trace, hdi_prob=.94, kind='all'):
    with RB_model:
        #  (hdi_prob=.94 is default)
        az_summary = az.summary(trace, round_to=4,  hdi_prob=hdi_prob, kind=kind )  
        
    return az_summary

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

def get_EPC_and_legends(azs):
    EPC_Bayes = alpha_to_EPC(azs['mean']['alpha'])
    EPC_Bayes_err = EPC_Bayes - alpha_to_EPC(azs['mean']['alpha']+azs['sd']['alpha'])
    Bayes_legend ="EPC Bayes {0:.5f} ({1:.5f})".format(EPC_Bayes, EPC_Bayes_err)
    Fitter_legend ="EPC Fitter {0:.5f} ({1:.5f})".format(rbfit.fit[0]['epc']\
                                                        ,rbfit._fit[0]['epc_err'])
    pred_epc_legend = "EPC predicted {0:.5f}".format(pred_epc)
    return EPC_Bayes, EPC_Bayes_err, Bayes_legend,Fitter_legend, pred_epc_legend
    
def EPC_compare_fitter_to_bayes(RB_model, azs, trace):
    EPC_Bayes, EPC_Bayes_err, Bayes_legend,Fitter_legend, pred_epc_legend = get_EPC_and_legends(azs)
    with RB_model:
        az.plot_posterior(trace,  var_names=['alpha'], round_to=4,
                          transform = alpha_to_EPC, point_estimate=None)
        plt.title("Error per Clifford")
        plt.axvline(x=alpha_to_EPC(alpha_ref),color='red')
        #plt.axvline(x=pred_epc,color='green') # WIP
        #plt.legend((Bayes_legend, "Higher density interval",Fitter_legend, pred_epc_legend), fontsize=10 )# WIP
        plt.legend((Bayes_legend, "Higher density interval",Fitter_legend), fontsize=10 )
        plt.show()

def GSP_compare_fitter_to_bayes(RB_model, azs):
    EPC_Bayes, EPC_Bayes_err, Bayes_legend,Fitter_legend,_ = get_EPC_and_legends(azs)
    # plot ground state population ~ Clifford length
    fig, axes = plt.subplots(1, 1, sharex=True, figsize=(10, 6))

    axes.set_ylabel("Ground State Population")
    axes.set_xlabel("Clifford Length")
    axes.plot(m_gates, np.mean(Y/shots,axis=0), 'r.')
    axes.plot(m_gates,azs['mean']['AB[0]']*azs['mean']['alpha']**m_gates+azs['mean']['AB[1]'],'--')
    #axes.plot(m_gates,azs['mean']['GSP'],'--') # WIP
    #axes.errorbar(m_gates, azs['mean']['GSP'], azs['sd']['GSP'], linestyle='None', marker='^') # WIP
    axes.plot(m_gates,mu_AB[0]*np.power(alpha_ref,m_gates)+mu_AB[1],':') 
    for i_seed in range(nseeds):
        plt.scatter(m_gates-0.25, Y[i_seed,:]/shots, label = "data", marker="x")
    axes.legend(["Mean Observed Frequencies",
                 "Bayesian Model\n"+Bayes_legend,
                 "Fitter Model\n"+Fitter_legend],fontsize=12)
    #axes.set_title('2 Qubit RB with T1/T2 Noise', fontsize=18) # WIP
    

def get_predicted_EPC(error_source):

    #Count the number of single and 2Q gates in the 2Q Cliffords
    gates_per_cliff = rb.rb_utils.gates_per_clifford(transpile_list,xdata[0],basis_gates,rb_opts['rb_pattern'][0])
    for basis_gate in basis_gates:
        print("Number of %s gates per Clifford: %f "%(basis_gate ,
                                                      np.mean([gates_per_cliff[rb_pattern[0][0]][basis_gate],
                                                               gates_per_cliff[rb_pattern[0][1]][basis_gate]])))
    # Calculate the predicted epc
    # from the known depolarizing errors on the simulation
    if error_source == "depolarization":  
        # Error per gate from noise model
        epgs_1q = {'u1': 0, 'u2': p1Q/2, 'u3': 2*p1Q/2}
        epg_2q = p2Q*3/4
        pred_epc = rb.rb_utils.calculate_2q_epc(
            gate_per_cliff=gates_per_cliff,
            epg_2q=epg_2q,
            qubit_pair=[0, 2],
            list_epgs_1q=[epgs_1q, epgs_1q])

    # using the predicted primitive gate errors from the coherence limit
    if error_source == "from_T1_T2": 
        # Predicted primitive gate errors from the coherence limit
        u2_error = rb.rb_utils.coherence_limit(1,[t1],[t2],gate1Q)
        u3_error = rb.rb_utils.coherence_limit(1,[t1],[t2],2*gate1Q)
        epg_2q = rb.rb_utils.coherence_limit(2,[t1,t1],[t2,t2],gate2Q)
        epgs_1q = {'u1': 0, 'u2': u2_error, 'u3': u3_error}
        pred_epc = rb.rb_utils.calculate_2q_epc(
            gate_per_cliff=gates_per_cliff,
            epg_2q=epg_2q,
            qubit_pair=[0, 1],
            list_epgs_1q=[epgs_1q, epgs_1q])
    return pred_epc

def get_count_data(result_list):
### another way to obtain the observed counts 

    Y_list = []
    for rbseed, result in enumerate(result_list):
        row_list = []
        for c_index, c_value in enumerate(nCliffs):
            if nQ == 2: 
                list_bitstring = ['00']
            elif nQ == 3:
                list_bitstring = ['000', '100'] # because q2 measured in c1
            total_counts = 0
            for bitstring in list_bitstring:
                    total_counts += result.get_counts()[c_index][bitstring]
            row_list.append(total_counts)
        Y_list.append(row_list)    
    return np.array(Y_list)    

#Number of qubits
nQ = 2
#There are 2 qubits: Q0,Q1.
#Number of seeds (random sequences)
nseeds = 10 # more data for the Rev. Mr. Bayes
#Number of Cliffords in the sequence (start, stop, steps)
nCliffs = np.arange(1,200,20)
#2Q RB Q0,Q1
rb_pattern = [[0,1]]
length_multiplier = 1

rb_opts  = {}
rb_opts ['length_vector'] = nCliffs
rb_opts ['nseeds'] = nseeds
rb_opts ['rb_pattern'] = rb_pattern
rb_opts ['length_multiplier'] = length_multiplier
rb_circs , xdata  = rb.randomized_benchmarking_seq(**rb_opts )

backend = device
basis_gates = ['u1','u2','u3','cx'] # use U,CX for now
shots = 1024
result_list = []
transpile_list = []
import time
for rb_seed,rb_circ_seed in enumerate(rb_circs):
    print('Compiling seed %d'%rb_seed)
    rb_circ_transpile = qiskit.transpile(rb_circ_seed,
                                         optimization_level=0,
                                         basis_gates=basis_gates)
    print('Runing seed %d'%rb_seed)
    job = qiskit.execute(rb_circ_transpile, 
                         shots=shots,
                         backend=backend)
    job_monitor(job)
    result_list.append(job.result())
    transpile_list.append(rb_circ_transpile)    
    
print("Finished Real Jobs")

print(rb_circs[0][0])

#Create an RBFitter object 
rbfit = rb.RBFitter(result_list, xdata, rb_opts['rb_pattern'])

m_gates, Y, alpha_ref, alpha_lower, alpha_upper, mu_AB, cov_AB, sigma_theta =\
    obtain_priors_and_data_from_fitter(printout = True)

### a check of the count matrix
np.sum((Y == (get_count_data(result_list)))*1) == Y.size

pooled = get_bayesian_model("pooled")
pm.model_to_graphviz(pooled)

trace_p = get_trace(pooled)

azp_summary = get_summary(pooled, trace_p)
azp_summary

hierarchical = get_bayesian_model("hierarchical")
pm.model_to_graphviz(hierarchical)

trace_h = get_trace(hierarchical)

azh_summary = get_summary(hierarchical, trace_h)
azh_summary

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

# predict EPC from the noisy model
#pred_epc = get_predicted_EPC(error_source = 'from_T1_T2') # this was for a noise model
pred_epc = 0.0165 # will not appear on graphs for real device but at this point functions need value (WIP)
print("Fake 2Q Error per Clifford: %e"%pred_epc)

EPC_compare_fitter_to_bayes(pooled, azp_summary, trace_p)

EPC_compare_fitter_to_bayes(hierarchical, azh_summary, trace_h)

GSP_compare_fitter_to_bayes(pooled, azp_summary)

GSP_compare_fitter_to_bayes(hierarchical, azh_summary)

import qiskit.tools.jupyter
%qiskit_version_table
%qiskit_copyright

%load_ext watermark
%watermark -n -u -v -iv -w



