import time
import itertools
import numpy as np

import qiskit
from qiskit import *
from qiskit.quantum_info import Statevector
from qiskit.algorithms.optimizers import COBYLA, L_BFGS_B, SPSA
from qiskit.opflow.state_fns import StateFn, CircuitStateFn
from qiskit.providers.aer import StatevectorSimulator, AerSimulator

from qiskit.opflow import CircuitSampler

from qiskit.ignis.mitigation.measurement import CompleteMeasFitter # you will need to pip install qiskit-ignis
from qiskit.ignis.mitigation.measurement import complete_meas_cal

import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
from matplotlib.colors import BoundaryNorm
cmap = plt.get_cmap("plasma") #'viridis'

from modules.utils import *

from qae import *

import datetime

print(qiskit.__version__, np.__version__)

IBMQ.load_account() # this then automatically loads your saved account
provider = IBMQ.get_provider(hub='ibm-q-research')

device = provider.backend.ibmq_jakarta # 6 bogota ; 4 rome

### Real device execution:
backend_device = device 

### Simulation with noise profile from real device
backend_sim = qiskit.providers.aer.AerSimulator.from_backend(device)

### Simulation without noise
#backend = qiskit.providers.aer.AerSimulator()

### Preliminaries
L = 5
num_trash = 2
anti = -1 # 1 for ferromagnetic Ising model, -1 for antiferromagnet

run_VQE = True # whether or not you want to compute the VQE values; or load them
execute = True
real_device = True    # want to execute on real device

# dataset to be loaded (optimal VQE parameters)
# if execute_VQE is True this will overwrite the file
filename = f'data/jakarta_execute'

# Give a _unique_ name to this notebook execution, the results will be stored in the appropriate files with that name
name = "jakarta_execute" 

# uncomment in case you want to toy around with the ansatz circuit
rotation_blocks = "ry"
entanglement_blocks = "cz" # cx
entanglement = "sca" #"sca"
reps = 1

ansatz_config = dict(rotation_blocks = rotation_blocks, entanglement_blocks = entanglement_blocks, reps = reps, entanglement = entanglement)

ansatz = qiskit.circuit.library.TwoLocal(L, **ansatz_config)
ansatz.draw("mpl")

tansatz = qiskit.transpile(ansatz, backend_sim)
tansatz.draw("mpl")

coupling_map = device.configuration().coupling_map
noise_model = qiskit.providers.aer.noise.NoiseModel.from_backend(device)
basis_gates = noise_model.basis_gates

qi_device = qiskit.utils.QuantumInstance(backend=backend_device, # , seed_simulator=seed, seed_transpiler=seed
                                  shots = 1000,
                                  coupling_map=coupling_map,
                                  #noise_model=noise_model, # comment out on real device execution
                                  measurement_error_mitigation_cls= CompleteMeasFitter, 
                                  cals_matrix_refresh_period=30  #How often to refresh the calibration matrix in measurement mitigation. in minutes
                                 )

qi_sim = qiskit.utils.QuantumInstance(backend=backend_sim, # , seed_simulator=seed, seed_transpiler=seed
                                  shots = 1000,
                                  coupling_map=coupling_map,
                                  noise_model=noise_model, # comment out on real device execution
                                  measurement_error_mitigation_cls= CompleteMeasFitter, 
                                  cals_matrix_refresh_period=30  #How often to refresh the calibration matrix in measurement mitigation. in minutes
                                 )

if not real_device:
    backend_device = backend_sim 
    qi_device = qi_sim

if run_VQE:
    maxiter = [200,200,200,200,200,
               200,200,200,150,150,
               150,150,100,100,100,
               60,30,30,30,30
              ]
    maxiter = maxiter[::2]
    maxiter = [150] + [50] * 9
    logspace_size = len(maxiter)
    gx_vals = np.logspace(-2,2,logspace_size)
    #gz_vals = [0.] # np.logspace(-2,2,logspace_size)

    counts = []
    values = []
    params = []

    def store_intermediate_result(eval_count, parameters, mean, std):
        counts.append(eval_count)
        values.append(mean)
        params.append(parameters)

    opt_params = []
    gx_list = []
    gz_list = []
    countss, valuess, paramss = [], [], []
    Qmag_sim, QZ_sim, Qen_sim= [],[], []
    Smag, Sen = [], []

    #maxiter = [200,200,200,200, 150, 60, 60, 60, 60, 60]
    for j,gx in enumerate(gx_vals):
        gz = 0
        t0 = datetime.datetime.now()
        counts = []
        values = []
        params = []

        optimizer = SPSA(maxiter=maxiter[j], blocking=True) # ,learning_rate=0.05, perturbation=0.05,; setting that speeds things up sometimes

        vqe = qiskit.algorithms.VQE(ansatz=ansatz,
                                    initial_point = opt_params[j-1] if j != 0 else None,
                                    optimizer=optimizer,
                                    callback=store_intermediate_result, # comment out on real device
                                    quantum_instance=qi_sim
                                   )

        H = QHIsing(L,anti,np.float32(gx),np.float32(gz))
        result = vqe.compute_minimum_eigenvalue(H, aux_operators = [QNKron(L,Z,I,i) for i in range(L)]) #ED with Qiskit VQE
        Qen_sim.append(result.eigenvalue)
        QZ_sim.append(result.aux_operator_eigenvalues[:,0]) # real-space resolved Z expectation value
        Qmag_sim.append(np.mean(anti**np.arange(L) * result.aux_operator_eigenvalues[:,0]))
        countss.append(counts)
        valuess.append(values)
        paramss.append(params)

        # uncomment if you want to recompute ED results (typically provided with guess parameters)
        ED_state, ED_E, ham = ising_groundstate(L, anti, np.float32(gx), np.float32(gz))
        Sen.append(ED_E)
        Smag.append(ED_state.T.conj()@Mag(L,anti)@ED_state)

        print(f"ED energy: {Sen[j]} ;; VQE energy: {Qen_sim[j]} ;; diff {Qen_sim[j] - Sen[j]}")
        print(f"ED mag: {Smag[j]} ;; VQE mag: {Qmag_sim[j]} ;; diff {Qmag_sim[j] - Smag[j]}")


        gx_list.append(gx)
        gz_list.append(gz)
        opt_params.append(sort_params(result.optimal_parameters))
        print(f"{j+1} / {len(gx_vals)}, gx = {gx:.2f}, gz = {gz:.2f}, time : {(datetime.datetime.now() - t0)}")
        fig, axs = plt.subplots(ncols=2,figsize=(6,3))
        axs[0].plot(countss[j],valuess[j])
        axs[1].plot(result.aux_operator_eigenvalues[:,0])
        plt.show()
        np.savez(filename, gx_list=gx_list, gz_list=gz_list, opt_params=opt_params,
                 Qmag_sim=Qmag_sim, Qen_sim=Qen_sim, QZ_sim = QZ_sim, Sen=Sen, Smag=Smag,
                 ansatz_config=ansatz_config,
                 paramss = paramss, valuess = valuess, countss = countss
                )

temp = np.load(filename + ".npz",allow_pickle=True)
gx_list = temp["gx_list"]
gz = 0
gz_list = [0]
Qmag_sim = temp["Qmag_sim"]
QZ_sim = temp["QZ_sim"]
Qen_sim = temp["Qen_sim"]
Sen = temp["Sen"]
Smag = temp["Smag"]
opt_params = temp["opt_params"]
paramss = temp["paramss"]; valuess = temp["valuess"]; countss = temp["countss"]

# This just recomputes the magnetization with the parameters provided by guess_params
if execute:
    QZ_executed=np.zeros((len(opt_params),L), dtype="complex")
    Qmag_executed=np.zeros((len(opt_params),), dtype="complex")
    for j in range(len(opt_params)):
        t0 = datetime.datetime.now()
        gx = gx_list[j]
        #H = QHIsing(L, anti, np.float32(gx), np.float32(gz)) # build Hamiltonian Op
        state = ansatz.assign_parameters(opt_params[j])
        for i in range(L):
            meas_outcome = ~StateFn(QNKron(L,Z,I,i)) @ StateFn(state)
            QZ_executed[j,i] = CircuitSampler(qi_device).convert(meas_outcome).eval() #https://quantumcomputing.stackexchange.com/questions/12080/evaluating-expectation-values-of-operators-in-qiskit
        Qmag_executed[j] = np.mean(anti**np.arange(L) * QZ_executed[j])
        #e_outcome = ~StateFn(H) @ StateFn(state)
        #Qen_executed[j] = CircuitSampler(qi).convert(e_outcome).eval()
        print(f"{j+1} / {len(opt_params)}, QZ: {Qmag_executed[j]}, gx = {gx:.2f}, gz = {gz:.2f}, time : {(datetime.datetime.now() - t0)}")
        fig, ax = plt.subplots(ncols=1,figsize=(4,3))
        ax.plot(QZ_sim[j].real)
        ax.plot(QZ_executed[j].real,"x--")
        ax.set_title("S exec = {:.3f}, S sim = {:.3f}".format(Qmag_executed[j].real, Qmag_sim[j].real))
        plt.show()

np.savez(filename + f"_executed_-real-device-{real_device}", gx_list=gx_list, gz_list=gz_list, opt_params=opt_params,
         Qmag_sim=Qmag_sim, Qen_sim=Qen_sim, QZ_sim = QZ_sim, Sen=Sen, Smag=Smag,
         Qmag_device= Qmag_executed, QZ_device = QZ_executed,
         ansatz_config=ansatz_config,
         paramss = paramss, valuess = valuess, countss = countss
        )

fig, axs = plt.subplots(ncols=2, figsize=(10,5))
ax = axs[0]
ax.plot(gx_list, np.abs(Smag),"x--", label="ED")
ax.plot(gx_list, np.abs(Qmag_sim),"x--", label="VQE sim")
ax.plot(gx_list, np.abs(Qmag_executed),"x--", label="VQE ibmq_rome")
ax.set_xscale("log")
ax.legend()

ax = axs[1]
ax.plot(gx_list, Smag,"x--", label="ED")
ax.plot(gx_list, Qmag_sim,"x--", label="VQE sim")
ax.plot(gx_list, Qmag_executed,"x--", label="VQE ibmq_rome")
ax.set_xscale("log")
ax.legend()

##############################################################################
### II - Anomaly Detection  ##################################################
##############################################################################

trash_qubits_idxs = [1,3]

# Using our QAEAnsatz class (described in qae.py)
circ = QAEAnsatz(num_qubits = L, num_trash_qubits= num_trash, trash_qubits_idxs = trash_qubits_idxs, measure_trash=True)

circ.draw("mpl")

# an example with assigned parameters
circ = circ.assign_parameters(np.arange(L*2+2))
circ.draw("mpl")

def prepare_circuit(thetas, L=6, num_trash=2, init_state=None, measurement=True, vqe=True):
    QAE_circ = QAEAnsatz(num_qubits = L, num_trash_qubits= num_trash, trash_qubits_idxs = trash_qubits_idxs, measure_trash=measurement).assign_parameters(thetas)
    if vqe:
        VQE_circ = ansatz.assign_parameters(init_state)
        fullcirc = VQE_circ + QAE_circ
    else:
        fullcirc = QAE_circ.initialize(init_state, qreg)
    return fullcirc

# example of the full circuit with random parameters assigned
fullcirc = prepare_circuit(np.random.rand(L*2+2),L,num_trash, vqe=True, init_state=opt_params[0])
fullcirc.draw("mpl")

backend = backend_device

def calibrate_circuit(L, nums_trash=trash_qubits_idxs,shots=1000):
    # outputs a CompleteMeasFitter object for alter use
    num_trash = len(nums_trash)
    qreg = QuantumRegister(L, 'q')
    # obtain calibration matrix
    qubit_list = [i for i in nums_trash] # only need to calibrate the trash qubits
    circlabel = f'mcal_{datetime.datetime.now()}'
    meas_calibs, state_labels = complete_meas_cal(qubit_list=qubit_list, qr=qreg, circlabel=circlabel)
    cal_job = backend.run(meas_calibs, shots=shots) #, noise_model=noise_model)
    cal_results = cal_job.result()
    meas_fitter = CompleteMeasFitter(cal_results, state_labels, circlabel=circlabel)
    print(circlabel, meas_fitter.cal_matrix)
    return meas_fitter

def run_circuit(thetas, L, num_trash, init_state, vqe=True, shots=100, meas_fitter = None):
    circ = prepare_circuit(thetas, L, num_trash, init_state, vqe=vqe)
    tcirc = qiskit.transpile(circ, backend)
    # Execute the circuit 
    job_sim = backend.run(tcirc, shots=shots) # , seed_simulator=123, seed_transpiler=234 fix seed to make it reproducible
    result = job_sim.result()
    # Results without mitigation
    counts = result.get_counts()
    if meas_fitter != None:
        # Get the filter object
        meas_filter = meas_fitter.filter

        # Results with mitigation
        mitigated_results = meas_filter.apply(result)
        counts = mitigated_results.get_counts(0)
    return counts

phis = opt_params # the parameters for the VQE initialization

def cost_function_single(thetas, L, num_trash, p, shots=1000, vqe=True, param_encoding=False, x=0, meas_fitter=None):
    """ Optimizes circuit """
    if vqe:
        init_state = phis[p]
    else:
        J, gx, gz = p
        init_state, _ = ising_groundstate(L, J, gx, gz)
    if param_encoding: thetas = feature_encoding(thetas, x) 
    out = run_circuit(thetas, L, num_trash, init_state, vqe=vqe, shots=shots, meas_fitter=meas_fitter)
    cost = out.get('11', 0)*2 + out.get('01', 0) + out.get('10', 0)
    return cost/shots

def cost_function(thetas, L, num_trash, ising_params, shots=1000, vqe=True, param_encoding=False, x=0, meas_fitter=None):
    """ Optimizes circuit """
    cost = 0.
    n_samples = len(ising_params)
    for i, p in enumerate(ising_params):
        if param_encoding: 
            cost += cost_function_single(thetas, L, num_trash, p, shots, vqe, param_encoding, x[i], meas_fitter=meas_fitter)
        else:
            cost += cost_function_single(thetas, L, num_trash, p, shots, vqe, param_encoding, meas_fitter=meas_fitter)
    return cost/n_samples

def optimize(ising_params, L=6, num_trash=2, thetas=None, shots=1000, max_iter=400, vqe=True, param_encoding=False, x=0, pick_optimizer = None,
            meas_fitter=None, callback=True):
    if thetas is None:
        n_params = (2*L+2)*2 if param_encoding else (2*L+2)
        thetas = np.random.uniform(0, 2*np.pi, n_params) # initial parameters without feature encoding
        
    print("Initial cost: {:.3f}".format(cost_function(thetas, L, num_trash, ising_params, shots, vqe, param_encoding, x)))
    
    counts, values, accepted = [], [], []
    def store_intermediate_result(eval_count, parameters, mean, std, ac):
        # counts.append(eval_count)
        values.append(mean)
        accepted.append(ac)

    # Initialize optimizer
    if pick_optimizer == "cobyla":
        optimizer = COBYLA(maxiter=max_iter, tol=0.0001)
    if pick_optimizer == "adam" or pick_optimizer == "ADAM":
        optimizer = qiskit.algorithms.optimizers.ADAM(maxiter=max_iter)
    # optimizer = L_BFGS_B(maxfun=300, maxiter=max_iter)#, factr=10, iprint=- 1, epsilon=1e-08)
    if pick_optimizer == "spsa" or pick_optimizer == None:
        optimizer = SPSA(maxiter=max_iter,
                         blocking=True,
                         callback=store_intermediate_result if callback else None,
                         learning_rate=0.1,
                         perturbation=0.1
                         ) # recommended from qiskit (first iteraction takes quite long)
                           # to reduce time figure out optimal learning rate and perturbation in advance

    start_time = time.time()
    ret = optimizer.optimize(
                            num_vars=len(thetas),
                            objective_function=(lambda thetas: cost_function(thetas, L, num_trash, ising_params, shots, vqe, param_encoding, x, meas_fitter=meas_fitter)),
                            initial_point=thetas
                            )
    print("Time: {:.5f} sec".format(time.time()-start_time))
    print(ret)
    return ret[0], values, accepted



backend = backend_sim

#calibration
meas_fitter = calibrate_circuit(L, nums_trash=trash_qubits_idxs)

# Training
phys_params = [0]
thetas_opt, loss, accepted = optimize(phys_params, max_iter=90, L=5, meas_fitter=meas_fitter) #, pick_optimizer="adam")
plt.plot(loss)

np.savez(filename + "run1_thetas_opt_sim", thetas_opt=thetas_opt, loss=loss, accepted=accepted)

# Inference; note that on the real device each inference points takes about ~30 seconds
cost = np.zeros((len(gx_vals)))
thetas_opt = np.load(filename + "run1_thetas_opt_sim.npz",allow_pickle=True)["thetas_opt"]
shots = 1000
for i,p in enumerate(list(gx_list)):
    t0 = datetime.datetime.now()
    cost[i] = cost_function_single(thetas_opt, L, num_trash, i, shots=shots, meas_fitter=meas_fitter) # np.random.uniform(0, 2*np.pi, 2*L+2) random parameters to check if training "does" something - result was: with random paremters just get noise, so yes, it "does" something
    print(f"{i} / {len(gx_list)} - execution time {datetime.datetime.now()-t0}")

np.savez(filename + "run1_sim",cost=cost, thetas_opt=thetas_opt, loss=loss, accepted = accepted)

# Training
phys_params = [0]
thetas_opt, loss, accepted = optimize(phys_params, max_iter=90, L=5, meas_fitter=meas_fitter) #, pick_optimizer="adam")
plt.plot(loss)

np.savez(filename + "run1_thetas_opt_sim2", thetas_opt=thetas_opt, loss=loss, accepted=accepted)

# Inference; note that on the real device each inference points takes about ~30 seconds
cost = np.zeros((len(gx_vals)))
thetas_opt = np.load(filename + "run1_thetas_opt_sim2.npz",allow_pickle=True)["thetas_opt"]
shots = 1000
for i,p in enumerate(list(gx_list)):
    t0 = datetime.datetime.now()
    cost[i] = cost_function_single(thetas_opt, L, num_trash, i, shots=shots, meas_fitter=meas_fitter) # np.random.uniform(0, 2*np.pi, 2*L+2) random parameters to check if training "does" something - result was: with random paremters just get noise, so yes, it "does" something
    print(f"{i} / {len(gx_list)} - execution time {datetime.datetime.now()-t0}")

np.savez(filename + "run1_sim2",cost=cost, thetas_opt=thetas_opt, loss=loss, accepted = accepted)

fig, axs = plt.subplots(ncols=2, figsize=(10,5))
temp0 = np.load(filename + "run1_sim2.npz",allow_pickle=True)
cost = temp0["cost"] ; 
ax = axs[0]
ax.plot(gx_list, cost,"X--", label="cost")
ax.plot(gx_list, np.abs(Qmag_sim),"x--", label="VQE res")
ax.plot(gx_list, np.abs(Qmag_executed),"x--", label="VQE resample")
ax.set_xscale("log")
ax.legend()

ax = axs[1]
ax.plot(gx_list, cost,"X--", label="cost")
ax.plot(gx_list, Qmag_sim,"x--", label="VQE res")
ax.plot(gx_list, Qmag_executed,"x--", label="VQE resample")
ax.set_xscale("log")
ax.legend()

# Training
phys_params = [-1]
thetas_opt, loss, accepted = optimize(phys_params, max_iter=40, L=5, meas_fitter=meas_fitter) #, pick_optimizer="adam")
plt.plot(loss)

np.savez(filename + "run2_thetas_opt_sim2", thetas_opt=thetas_opt, loss=loss, accepted=accepted)

# Inference; note that on the real device each inference points takes about ~30 seconds
cost = np.zeros((len(gx_vals)))
thetas_opt = np.load(filename + "run2_thetas_opt_sim2.npz",allow_pickle=True)["thetas_opt"]
shots = 1000
for i,p in enumerate(list(gx_list)):
    t0 = datetime.datetime.now()
    cost[i] = cost_function_single(thetas_opt, L, num_trash, i, shots=shots, meas_fitter=meas_fitter) # np.random.uniform(0, 2*np.pi, 2*L+2) random parameters to check if training "does" something - result was: with random paremters just get noise, so yes, it "does" something
    print(f"{i} / {len(gx_list)} - execution time {datetime.datetime.now()-t0}")

np.savez(filename + "run2_sim2",cost=cost, thetas_opt=thetas_opt, loss=loss, accepted = accepted)

fig, axs = plt.subplots(ncols=2, figsize=(10,5))
temp0 = np.load(filename + "run2_sim2.npz",allow_pickle=True)
cost = temp0["cost"] ; 
ax = axs[0]
ax.plot(gx_list, cost,"X--", label="cost")
ax.plot(gx_list, np.abs(Qmag_sim),"x--", label="VQE res")
ax.plot(gx_list, np.abs(Qmag_executed),"x--", label="VQE resample")
ax.set_xscale("log")
ax.legend()

ax = axs[1]
ax.plot(gx_list, cost,"X--", label="cost")
ax.plot(gx_list, Qmag_sim,"x--", label="VQE res")
ax.plot(gx_list, Qmag_executed,"x--", label="VQE resample")
ax.set_xscale("log")
ax.legend()

backend = backend_device

meas_fitter = calibrate_circuit(L, nums_trash=trash_qubits_idxs)

# Training
phys_params = [0]
thetas_guess = np.load(filename + "run1_thetas_opt_sim.npz",allow_pickle=True)["thetas_opt"]
thetas_opt, loss, accepted = optimize(phys_params, thetas= thetas_guess, callback=False, max_iter=15, L=5, meas_fitter=meas_fitter) #, pick_optimizer="adam")
plt.plot(loss)

np.savez(filename + "run1_thetas_opt_device", thetas_opt=thetas_opt, loss=loss, accepted=accepted)

# Inference; note that on the real device each inference points takes about ~30 seconds
cost = np.zeros((len(gx_vals)))
thetas_opt = np.load(filename + "run1_thetas_opt_device.npz",allow_pickle=True)["thetas_opt"]
shots = 1000
for i,p in enumerate(list(gx_list)):
    t0 = datetime.datetime.now()
    cost[i] = cost_function_single(thetas_opt, L, num_trash, i, shots=shots, meas_fitter=meas_fitter) # np.random.uniform(0, 2*np.pi, 2*L+2) random parameters to check if training "does" something - result was: with random paremters just get noise, so yes, it "does" something
    print(f"{i} / {len(gx_list)} - execution time {datetime.datetime.now()-t0}")

np.savez(filename + "run1_device",cost=cost, thetas_opt=thetas_opt, loss=loss, accepted = accepted)

fig, axs = plt.subplots(ncols=2, figsize=(10,5))
cost_device = np.load(filename + "run1_device.npz",allow_pickle=True)["cost"]
cost_sim = np.load(filename + "run1_sim.npz",allow_pickle=True)["cost"]
ax = axs[0]
ax.plot(gx_list, cost_sim,"X--", label="cost sim.")
ax.plot(gx_list, cost_device,"X--", label="cost ibmq_jakarta")
ax.plot(gx_list, np.abs(Qmag_sim),"x--", label="$\hat{S}$ sim")
ax.plot(gx_list, np.abs(Qmag_executed),"x--", label="$\hat{S}$ ibmq_jakarta")
ax.set_xscale("log")
ax.legend()

ax = axs[1]
ax.plot(gx_list, cost_sim,"X--", label="cost sim.")
ax.plot(gx_list, cost_device,"X--", label="cost ibmq_jakarta")
ax.plot(gx_list, Qmag_sim,"x--", label="$\hat{S}$ sim")
ax.plot(gx_list, Qmag_executed,"x--", label="$\hat{S}$ ibmq_jakarta")
ax.set_xscale("log")
ax.legend()

# Training
phys_params = [0]
thetas_guess = np.load(filename + "run1_thetas_opt_sim2.npz",allow_pickle=True)["thetas_opt"]
thetas_opt, loss, accepted = optimize(phys_params, thetas= thetas_guess, callback=False, max_iter=25, L=5, meas_fitter=meas_fitter) #, pick_optimizer="adam")
plt.plot(loss)

np.savez(filename + "run1_thetas_opt_device2", thetas_opt=thetas_opt, loss=loss, accepted=accepted)

# Inference; note that on the real device each inference points takes about ~30 seconds
cost = np.zeros((len(gx_vals)))
thetas_opt = np.load(filename + "run1_thetas_opt_device2.npz",allow_pickle=True)["thetas_opt"]
shots = 1000
for i,p in enumerate(list(gx_list)):
    t0 = datetime.datetime.now()
    cost[i] = cost_function_single(thetas_opt, L, num_trash, i, shots=shots, meas_fitter=meas_fitter) # np.random.uniform(0, 2*np.pi, 2*L+2) random parameters to check if training "does" something - result was: with random paremters just get noise, so yes, it "does" something
    print(f"{i} / {len(gx_list)} - execution time {datetime.datetime.now()-t0}")

np.savez(filename + "run1_device2",cost=cost, thetas_opt=thetas_opt, loss=loss, accepted = accepted)

filename

fig, axs = plt.subplots(ncols=2, figsize=(10,5))
cost_device = np.load(filename + "run1_device2.npz",allow_pickle=True)["cost"]
cost_sim = np.load(filename + "run1_sim2.npz",allow_pickle=True)["cost"]
ax = axs[0]
ax.plot(gx_list, cost_sim,"X--", label="cost sim.")
ax.plot(gx_list, cost_device,"X--", label="cost ibmq_jakarta")
ax.plot(gx_list, np.abs(Qmag_sim),"x--", label="$\hat{S}$ sim")
ax.plot(gx_list, np.abs(Qmag_executed),"x--", label="$\hat{S}$ ibmq_jakarta")
ax.set_xscale("log")
ax.legend()

ax = axs[1]
ax.plot(gx_list, cost_sim,"X--", label="cost sim.")
ax.plot(gx_list, cost_device,"X--", label="cost ibmq_jakarta")
ax.plot(gx_list, Qmag_sim,"x--", label="$\hat{S}$ sim")
ax.plot(gx_list, Qmag_executed,"x--", label="$\hat{S}$ ibmq_jakarta")
ax.set_xscale("log")
ax.legend()

meas_fitter = calibrate_circuit(L, nums_trash=trash_qubits_idxs)

# Inference; note that on the real device each inference points takes about ~30 seconds
cost = np.zeros((len(gx_vals)))
thetas_opt = np.load(filename + "run2_thetas_opt_sim.npz",allow_pickle=True)["thetas_opt"]
shots = 1000
for i,p in enumerate(list(gx_list)):
    t0 = datetime.datetime.now()
    cost[i] = cost_function_single(thetas_opt, L, num_trash, i, shots=shots, meas_fitter=meas_fitter) # np.random.uniform(0, 2*np.pi, 2*L+2) random parameters to check if training "does" something - result was: with random paremters just get noise, so yes, it "does" something
    print(f"{i} / {len(gx_list)} - execution time {datetime.datetime.now()-t0}")

filename

np.savez(filename + "run2_device",cost=cost, thetas_opt=thetas_opt, loss=loss, accepted = accepted)

fig, axs = plt.subplots(ncols=2, figsize=(10,5))
cost_device = np.load(filename + "run2_device.npz",allow_pickle=True)["cost"]
cost_sim = np.load(filename + "run2_sim.npz",allow_pickle=True)["cost"]
ax = axs[0]
ax.plot(gx_list, cost_sim,"X--", label="cost sim.")
ax.plot(gx_list, cost_device,"X--", label="cost ibmq_jakarta")
ax.plot(gx_list, np.abs(Qmag_sim),"x--", label="$\hat{S}$ sim")
ax.plot(gx_list, np.abs(Qmag_executed),"x--", label="$\hat{S}$ ibmq_jakarta")
ax.set_xscale("log")
ax.legend()

ax = axs[1]
ax.plot(gx_list, cost_sim,"X--", label="cost sim.")
ax.plot(gx_list, cost_device,"X--", label="cost ibmq_jakarta")
ax.plot(gx_list, Qmag_sim,"x--", label="$\hat{S}$ sim")
ax.plot(gx_list, Qmag_executed,"x--", label="$\hat{S}$ ibmq_jakarta")
ax.set_xscale("log")
ax.legend()

meas_fitter = calibrate_circuit(L, nums_trash=trash_qubits_idxs)

# Training
phys_params = [-1]
thetas_guess = np.load(filename + "run2_thetas_opt_sim.npz",allow_pickle=True)["thetas_opt"]
thetas_opt, loss, accepted = optimize(phys_params, thetas= thetas_guess, callback=False, max_iter=25, L=5, meas_fitter=meas_fitter) #, pick_optimizer="adam")
plt.plot(loss)

np.savez(filename + "run2_thetas_opt_device_overnight", thetas_opt=thetas_opt, loss=loss, accepted=accepted)

# Inference; note that on the real device each inference points takes about ~30 seconds
cost = np.zeros((len(gx_vals)))
thetas_opt = np.load(filename + "run2_thetas_opt_device_overnight.npz",allow_pickle=True)["thetas_opt"]
shots = 1000
for i,p in enumerate(list(gx_list)):
    t0 = datetime.datetime.now()
    cost[i] = cost_function_single(thetas_opt, L, num_trash, i, shots=shots, meas_fitter=meas_fitter) # np.random.uniform(0, 2*np.pi, 2*L+2) random parameters to check if training "does" something - result was: with random paremters just get noise, so yes, it "does" something
    print(f"{i} / {len(gx_list)} - execution time {datetime.datetime.now()-t0}")

np.savez(filename + "run2_device_overnight",cost=cost, thetas_opt=thetas_opt, loss=loss, accepted = accepted)

fig, axs = plt.subplots(ncols=2, figsize=(10,5))
cost_device = np.load(filename + "run2_device_overnight.npz",allow_pickle=True)["cost"]
cost_sim = np.load(filename + "run2_sim2.npz",allow_pickle=True)["cost"]
ax = axs[0]
ax.plot(gx_list, cost_sim,"X--", label="cost sim.")
ax.plot(gx_list, cost_device,"X--", label="cost ibmq_jakarta")
ax.plot(gx_list, np.abs(Qmag_sim),"x--", label="$\hat{S}$ sim")
ax.plot(gx_list, np.abs(Qmag_executed),"x--", label="$\hat{S}$ ibmq_jakarta")
ax.set_xscale("log")
ax.legend()

ax = axs[1]
ax.plot(gx_list, cost_sim,"X--", label="cost sim.")
ax.plot(gx_list, cost_device,"X--", label="cost ibmq_jakarta")
ax.plot(gx_list, Qmag_sim,"x--", label="$\hat{S}$ sim")
ax.plot(gx_list, Qmag_executed,"x--", label="$\hat{S}$ ibmq_jakarta")
ax.set_xscale("log")
ax.legend()

