import numpy as np
import copy
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
import unif_bayesian_fitter as bf

simulation = False # 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  

# for WIP
import importlib
importlib.reload(bf)

lengths = np.arange(1, 2225, 225)
num_samples = 10
seed = 1010
qubits = [4]

# last backend's recorded EPG before experiments
print(RBUtils.get_error_dict_from_backend(backend, qubits))

# backend's recorded EPG after experiments
print(RBUtils.get_error_dict_from_backend(backend, qubits))

# Correct reference for shift
# example of interpolated EPC_cal for hardware experiments
# EPC0 + (t_exp - tO) * (EPC1 - EPC0) / (t1 - t0)
# code here:
epc_calib = 0.0001600 + (18+25/60-7)*(0.0001772 - 0.0001600)/24
print('EPC calibration: {0:1.4e}'.format(epc_calib))

# Run an RB experiment on qubit 0
exp1 = StandardRB(qubits, lengths, num_samples=num_samples, seed=seed)
expdata1 = exp1.run(backend).block_for_results()
results1 = expdata1.analysis_results()

# View result data
display(expdata1.figure(0))
for result in results1:
    print(result)

popt = expdata1.analysis_results()[0].value.value
pcov = expdata1.analysis_results()[0].extra['covariance_mat']
epc_est_fm = expdata1.analysis_results()[2].value.value
epc_est_fm_err = expdata1.analysis_results()[2].value.stderr

EPG_dic = {}
for i in range(3,6):
    EPG_key = expdata1.analysis_results()[i].name
    EPG_dic[EPG_key] = expdata1.analysis_results()[i].value.value 

nQ = len(qubits)
scale = (2 ** nQ - 1) / 2 ** nQ
interleaved_gate =''

# get count data
Y = bf.get_GSP_counts(expdata1._data, len(lengths),range(num_samples))
expdata1._data[1]
experiment_type = expdata1._data[0]['metadata']['experiment_type']
physical_qubits = expdata1._data[0]['metadata']['physical_qubits']
shots = expdata1._data[0]['shots']

#build model
pooled_model = bf.build_bayesian_model(model_type="pooled",Y=Y,
                                       shots=shots,m_gates=lengths,
                                       popt = popt,
                                       pcov = pcov)
                                           
pm.model_to_graphviz(pooled_model)

trace_p = bf.get_trace(pooled_model, target_accept = 0.95)

bf.RB_bayesian_results(pooled_model, trace_p, lengths,
                       epc_est_fm,
                       epc_est_fm_err,
                       experiment_type,
                       scale,                       
                       num_samples, Y, shots, physical_qubits, interleaved_gate, backend,
                       EPG_dic = EPG_dic,
                       routine = 'build_bayesian_model')

#build model
hierarchical_model = bf.build_bayesian_model(model_type="h_sigma",Y=Y,
                                           shots=shots,m_gates=lengths,
                                           popt = popt,
                                           pcov = pcov,
                                           sigma_theta=0.001,sigma_theta_l=0.0005,sigma_theta_u=0.0015)
pm.model_to_graphviz(hierarchical_model)

trace_h = bf.get_trace(hierarchical_model, target_accept = 0.99)

bf.RB_bayesian_results(hierarchical_model, trace_h, lengths,
                       epc_est_fm, epc_est_fm_err, experiment_type, scale,                       
                       num_samples, Y, shots, physical_qubits, interleaved_gate, backend,
                       EPG_dic = EPG_dic, routine = 'build_bayesian_model')

# describe RB experiment
interleaved_gate = "x"
qubits = [4]
interleaved_circuit = circuits.XGate()

lengths = np.arange(1, 2250, 225)
num_samples = 10
seed = 1010

# Run an interleaved RB experiment

int_exp1 = InterleavedRB(interleaved_circuit, qubits,
                            lengths, num_samples=num_samples, seed=seed)

# Run
int_expdata1 = int_exp1.run(backend).block_for_results()
int_results1 = int_expdata1.analysis_results()

# View result data
display(int_expdata1.figure(0))
for result in int_results1:
    print(result)

popt = int_expdata1.analysis_results()[0].value.value
pcov = int_expdata1.analysis_results()[0].extra['covariance_mat']
popt[2] = popt[1]/popt[2] # replace alpha_C by p_tilde
# WIP rigorously the covariance matrix could be modified too if used
epc_est_fm = int_expdata1.analysis_results()[3].value.value
epc_est_fm_err = int_expdata1.analysis_results()[3].value.stderr

nQ = len(qubits)
scale = (2 ** nQ - 1) / 2 ** nQ
interleaved_gate ='x'

# get count data
Y1 = bf.get_GSP_counts(int_expdata1._data, len(lengths),
                    range(0,2*num_samples-1,2))

Y2 = bf.get_GSP_counts(int_expdata1._data, len(lengths),
                    range(1,2*num_samples,2))
int_expdata1._data[1]
experiment_type = int_expdata1._data[0]['metadata']['experiment_type']
physical_qubits = int_expdata1._data[0]['metadata']['physical_qubits']
shots = int_expdata1._data[0]['shots']

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

tilde1 =  bf.build_bayesian_model("tilde",Y=Y,shots=shots,
                                  m_gates=lengths,
                                  popt = popt,
                                  pcov = pcov,
                                  RvsI=RvsI,IvsR=IvsR)
pm.model_to_graphviz(tilde1)

trace_t = bf.get_trace(tilde1)

bf.RB_bayesian_results(tilde1, trace_t, lengths,
                   epc_est_fm, epc_est_fm_err, experiment_type, scale,                      
                   num_samples, Y, shots, physical_qubits, interleaved_gate, backend,
                   epc_calib = epc_calib, Y1 = Y1, Y2 = Y2,
                   routine = 'build_bayesian_model')

# Correct reference for shift
# example of interpolated EPC_cal for hardware experiments
# EPC0 + (t_exp - tO) * (EPC1 - EPC0) / (t1 - t0)
# code here:
epc_calib = 0.0001600 + (17+36/60-7)*(0.0001772 - 0.0001600)/24
print('EPC calibration: {0:.6f}'.format(epc_calib))

import importlib
importlib.reload(bf)

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

tilde2 =  bf.build_bayesian_model("h_tilde",Y=Y,shots=shots, m_gates=lengths,
                                 popt = popt,
                                 pcov = pcov,
                                 RvsI = RvsI_h, IvsR = IvsR_h,
                                 sigma_theta=0.001,sigma_theta_l=0.0005,sigma_theta_u=0.0015)  
                                            

pm.model_to_graphviz(tilde2)

trace_t3 = bf.get_trace(tilde2, target_accept = .95)

bf.RB_bayesian_results(tilde2, trace_t3, lengths,
                       epc_est_fm, epc_est_fm_err, experiment_type, scale,                      
                       num_samples, Y, shots, physical_qubits, interleaved_gate, backend,
                       epc_calib = epc_calib, Y1 = Y1, Y2 = Y2,
                       routine = 'build_bayesian_model')

# Correct reference for shift
# example of interpolated EPC_cal for hardware experiments
# EPC0 + (t_exp - tO) * (EPC1 - EPC0) / (t1 - t0)
# code here:
epc_calib = 0.0001600 + (17+36/60-7)*(0.0001772 - 0.0001600)/24
print('EPC calibration: {0:.6f}'.format(epc_calib))