import numpy as np
import copy
from qiskit import QuantumRegister, QuantumCircuit, ClassicalRegister, Aer, execute, transpile, assemble
from qiskit.tools.visualization import *
from qiskit.ignis.mitigation.measurement import (complete_meas_cal, tensored_meas_cal,
                                                 CompleteMeasFitter, TensoredMeasFitter)
import json
import time
from qiskit.tools.monitor import job_monitor

from c_utils import new_cut # circuit building utilities

data_directory = "data2_files/" # this directory for 2d pilot project data
def json_dic_loader(dic_name):
    f = open(data_directory+dic_name+'.json')
    return json.load(f)
def json_dic_dumper(dic, dic_name):
    with open(data_directory+dic_name+'.json', 'w') as f:
        json.dump(dic,f)   

simulator = Aer.get_backend('qasm_simulator')

#specify the layout of the devices
used_qubits = 5
qubit_list = [0,1,2,3,4]

program_name="AL2" # This for a mix of W/Psi+ and W_bar/Phi+ separable states (2d pilot project)
Flag_char = "DS" # use the joint set
if len(Flag_char) >= 2:
    unique_char = "M" # for "mixed"
else:
    unique_char = Flag_char

# These dictionaries for the devices used in the study
QV_dic = {'ibmq_athens': 32.0, 'ibmq_valencia': 16.0, 'ibmq_ourense': 8.0, 
          "ibmqx2": 8.0, 'ibmq_santiago': 32.0, 'ibmq_vigo': 16.0, 'ideal_device': np.inf}
dev_dic =  {'ibmq_santiago': "San",'ibmq_athens': "Ath", 'ibmq_valencia': "Val", 'ibmq_vigo': 'Vig','ibmq_ourense': "Our",
               "ibmqx2": 'Yor',  'ideal_device': "Ide"}

# specify the device: here first the ideal noise-free device
project_device = 'ideal_device'
device_name = dev_dic[project_device]

# specify the nb of id gates between state creation and measurements
# zero for the ideal device of course
id_gates = 0
str_nb_id = str(id_gates)
zfilled = str_nb_id.zfill(4-len(str_nb_id))

# tail of the file names for RAM storage
mitig_name = program_name + "_" + device_name
project_name = mitig_name + "_" + unique_char + zfilled
print(mitig_name)
print(project_name)

# establish the result label list
# meas_calibs will be used for mitigation in the real device section
qr = QuantumRegister(used_qubits) # 
meas_calibs, label_list = complete_meas_cal(qubit_list=qubit_list, qr=qr, circlabel='mcal')
nb_labels=len(label_list)
print(nb_labels,label_list)

len(meas_calibs)

# permutation list
# here it is simple to write down the list, 
# but a version using itertools will be wellcome for >5 qubits projects
q_perm = [[0, 1, 2, 3, 4], [0, 1, 3, 2, 4], [0, 1, 4, 2, 3], [0, 2, 3, 1, 4], [0, 2, 4, 1, 3],
          [0, 3, 4, 1, 2], [1, 2, 3, 0, 4], [1, 2, 4, 0, 3], [1, 3, 4, 0, 2], [2, 3, 4, 0, 1]]

# define the two subsets of 10 separable states
if program_name == "QAD":
    state_1a = ["W","Phi+"]
    state_1b = ["GHZ","Psi+"]
elif program_name == "ALT" or "AL2":
    state_1a = ["W","Psi+"]
    state_1b = ["Wbar","Phi+"]
l_states = state_1a+state_1b

l_states

# version 20 circuits for demonstration
# (in the version run on real devices: two batches of 10 circuits)

# these circuits limited to state creation are ready to be saved 
# for  ultimately building circuits adapted to noisy simulator and real devices
# as option, these circuits will include a row of id gates between creation and measurements

circ_ori = []
for i_s in range(0,len(l_states),2):
    for perm in q_perm:
            mycircuit = QuantumCircuit(used_qubits, used_qubits)                                   
            mycircuit = new_cut.circuit_builder(mycircuit, perm, l_states[i_s],l_states[i_s+1])
            circ_ori.append(mycircuit)

# add measurement section to the circuit set newly created:
nb_states = len(circ_ori)
circ_ideal = copy.deepcopy(circ_ori)
for i_state in range(nb_states):
    new_cut.add_barrier_and_measure(circ_ideal[i_state],qubit_list)

# execute on noise free simulator  
s_sim = 12000
job_simul = execute(circ_ideal, backend=simulator, shots=s_sim)
tot_results_simul = job_simul.result()

# establish a dictionary of count results on noise free simulator:
# (this step is only useful if ram storage is performed)
void_counts = dict(zip(label_list, np.full(2**used_qubits,0.0))) #, dtype=int)))
tot_results_sim_dic = {}
ideal_dic = {}
for i_state in range(nb_states):
    counts_simul = copy.deepcopy(void_counts)
    counts_simul.update(tot_results_simul.get_counts(i_state))
    ideal_dic[str(i_state)]=counts_simul

i_state_test = 0
print(device_name, "circuit #",i_state_test)
circ_ideal[i_state_test].draw(output='mpl')

print(device_name, "circuit #",i_state_test)
plot_histogram(ideal_dic[str(i_state_test)],
               legend=['noise free simulation'],
               color = "b", figsize=(10.,5.))

i_state_test = 10
print(device_name, "circuit #",i_state_test)
circ_ideal[i_state_test].draw(output='mpl')

print(device_name, "circuit #",i_state_test)
plot_histogram(ideal_dic[str(i_state_test)],
               legend=['noise free simulation'],
               color = "b", figsize=(10.,5.))

PD_ideal =  np.ndarray((nb_states,nb_labels))
for i_state in range(nb_states):
    PD_ideal[i_state, :] = list(ideal_dic[str(i_state)].values())
# now a little trick to get the ideal values from the simulator approximated values
with np.errstate(divide='ignore'): # ignore the divide by zero warning
    PD_ideal = 1/np.round(s_sim/(PD_ideal))
# have a look at the matrix head and tail:
print("first and last state probability distributions:")
print(np.round(np.vstack((PD_ideal[0:1,:],PD_ideal[-1:,:])),4)) 

from qiskit import IBMQ

IBMQ.load_account()

provider = IBMQ.get_provider(hub='ibm-q')

print(provider.backends())

project_device = 'ibmq_valencia'# you may choice here a different backend
device_name = dev_dic[project_device]
mitig_name = program_name + "_" + device_name
print(mitig_name)

# determine here the backend
device = provider.get_backend(project_device) # the backend names are listed here above
properties = device.properties()
coupling_map = device.configuration().coupling_map

id_gates = 0 # choice of 0 or 256 at this time
str_nb_id = str(id_gates)
zfilled = str_nb_id.zfill(4-len(str_nb_id))
project_name = mitig_name + "_" + unique_char + zfilled
print(project_name)

circuit_dic = json_dic_loader("circuit_"+ project_name)

real_circs = []
for i_state in list(range(nb_states)): 
    real_circs.append(QuantumCircuit().from_qasm_str(circuit_dic[str(i_state)]))  

i_state_test = 10
print(project_device, "circuit #",i_state_test,
      "circuit depth:",real_circs[i_state_test].depth())
print('gates = ',real_circs[i_state_test].count_ops())

# you may want to skip this if large nb of id gates before measurement
real_circs[i_state_test].draw(output='mpl')

job_simul = execute(real_circs[i_state_test], backend=simulator, shots=s_sim)
print(project_device, "circuit #",i_state_test, "on noise free simulator")
simul_results = job_simul.result().get_counts()
plot_histogram(simul_results,
               legend=['noise free simulation'],
               color = "b", figsize=(10.,5.))

# retrieve the corresponding measurement mitigation filter obtained at experimental time
# use a fake job because the calibration results were stored as dictionary
simulator = Aer.get_backend('qasm_simulator')
fake_job_cal = execute(meas_calibs, backend=simulator, shots=1)
fake_cal_results = fake_job_cal.result()
cal_results_dic = json_dic_loader("cal_results_dic_"+mitig_name)

if 'date' in cal_results_dic.keys():
    str(cal_results_dic['date'])

cal_results = fake_cal_results.from_dict(cal_results_dic)

meas_fitter = CompleteMeasFitter(cal_results, label_list, qubit_list=qubit_list, circlabel='mcal')
meas_filter = meas_fitter.filter
# have a look at the average measurement fidefily of this device:
print("Average Measurement Fidelity was: %f" % meas_fitter.readout_fidelity(), "for",project_device)

empirical_dic = json_dic_loader('experimental_'+project_name)
test_dic = json_dic_loader('test_'+project_name) 

def rectify_counts(tot_res, test_cqi,mitigation,m_filter) :       
        # IMPORTANT MODIFICATION
        try:
            counts_results_real_test = tot_res[str(test_cqi)]
        except KeyError as error:            
            counts_results_real_test = tot_res[test_cqi]

        raw_counts_test = copy.deepcopy(void_counts)
        raw_counts_test.update(counts_results_real_test)

        if mitigation:
            mitigated_results_test = meas_filter.apply(raw_counts_test, method = 'least_squares')
            returned_counts = copy.deepcopy(void_counts)
            returned_counts.update(mitigated_results_test)
                        
        else:
            returned_counts = copy.deepcopy(raw_counts_test)
            
        return returned_counts
def get_clean_matrix(dic, mitigation,m_filter):
    clean_matrix =  np.ndarray((nb_states,nb_labels)) 
    for i_state in range(nb_states):
        rectified_counts = rectify_counts(dic,i_state, mitigation,m_filter) # get a rectified counts dictionary       
        clean_matrix[i_state, :] = list(rectified_counts.values())
    clean_matrix = clean_matrix/clean_matrix.sum(axis=1, keepdims=True)
    return clean_matrix

def obtain_pooled_PDM(mitigation):
    PD_exper = get_clean_matrix(empirical_dic, mitigation=mitigation,
                                m_filter=meas_filter)
    PD_test = get_clean_matrix(test_dic, mitigation=mitigation,
                                m_filter=meas_filter)
    return PD_exper + PD_test

PD_tot = obtain_pooled_PDM(False)
PD_totm = obtain_pooled_PDM(True)

i_state_test = 19 # choose here the circuit
print(project_device, "circuit #",i_state_test,
      "circuit depth:",real_circs[i_state_test].depth())
print('gates = ',real_circs[i_state_test].count_ops())

ideal_results = dict(zip(label_list,PD_ideal[i_state_test]))
real_results = dict(zip(label_list,PD_tot[i_state_test]))
mit_results = dict(zip(label_list,PD_totm[i_state_test]))

plot_histogram([ideal_results, real_results, mit_results],
               legend=['ideal device','real results on\n '+ project_device, 'after measurement\n errror mitigation'],
               color =["b","r","g"],
               bar_labels=False,
               figsize=(10.,5.))

import qiskit.tools.jupyter
%qiskit_version_table

