# instantiate runtime service
from qiskit_ibm_runtime import QiskitRuntimeService

QiskitRuntimeService.save_account(channel="ibm_quantum", token="MY_IBM_QUANTUM_TOKEN")
service = QiskitRuntimeService(channel="ibm_quantum")

# import required libraries and instantiate runtime service
from qiskit_ibm_runtime import Session, Estimator, Options
from qiskit.circuit.library import RealAmplitudes
from qiskit.quantum_info import SparsePauliOp

# create circuits and observables
psi1 = RealAmplitudes(num_qubits=2, reps=2)
psi2 = RealAmplitudes(num_qubits=2, reps=3)

H1 = SparsePauliOp.from_list([("II", 1), ("IZ", 2), ("XI", 3)])
H2 = SparsePauliOp.from_list([("IZ", 1)])
H3 = SparsePauliOp.from_list([("ZI", 1), ("ZZ", 1)])

psi1.decompose().draw(output='mpl')

psi2.decompose().draw(output='mpl')

import numpy as np

# generate random theta params
np.random.seed(0)
theta1 = np.random.rand(6) * np.pi
theta2 = np.random.rand(8) * np.pi
theta3 = np.random.rand(6) * np.pi
print(theta1)
print(theta2)
print(theta3)

with Session(service=service, backend="ibmq_qasm_simulator") as session:
    estimator = Estimator(session=session)


from qiskit_ibm_runtime import Options
options = Options(optimization_level=1)
options.execution.shots = 1024

with Session(service=service, backend="ibmq_qasm_simulator") as session:
    estimator = Estimator(session=session, options=options)

    # calculate [ <psi1(theta1)|H1|psi1(theta1)> ]
    psi1_H1 = estimator.run(circuits=[psi1], observables=[H1],
                            parameter_values=[theta1])
    print(psi1_H1.result())

    # You can invoke run() multiple times!
    
    # calculate [ <psi1(theta1)|H2|psi1(theta1)>, <psi1(theta1)|H3|psi1(theta1)> ]
    psi1_H23 = estimator.run(circuits=[psi1, psi1], observables=[H2, H3], 
                             parameter_values=[theta1]*2)
    print(psi1_H23.result())

    # Ex. 1: Calculate [ <psi2(theta2)|H2|psi2(theta2)> ]
    # psi2_H2 = ...


    # Ex. 2 calculate [ <psi1(theta1)|H1|psi1(theta1)>,
    #                   <psi2(theta2)|H2|psi2(theta2)>,
    #                   <psi1(theta3)|H3|psi1(theta3)> ]
    # psi12_H23 = ...

# Ex. 3: use optimization_level 3
options = Options()
# options... # hint: use autocomplete
with Session(service=service, backend="ibmq_qasm_simulator") as session:
    estimator = Estimator(session=session, options=options)
    theta4 = [0, 1, 1, 2, 3, 5]

    # calculate [ <psi1(theta1)|H1|psi1(theta1)> ]
    psi1_H1 = estimator.run(circuits=[psi1], observables=[H1], parameter_values=[theta4])
    print(psi1_H1.result())    

# Ex. 4: Apply REM to the following example
# Remeber REM is enabled when resilience level is 1
options = Options()
# options... # hint: use autocomplete

with Session(service=service, backend="ibmq_qasm_simulator",) as session:
    estimator = Estimator(session=session, options=options)
    theta4 = [0, 1, 1, 2, 3, 5]
    psi1_H1 = estimator.run(circuits=[psi1], observables=[H1], parameter_values=[theta4])
    print(psi1_H1.result())

from qiskit import QuantumCircuit
from qiskit_ibm_runtime import Sampler
bell = QuantumCircuit(2)
bell.h(0)
bell.cx(0, 1)
bell.measure_all()

with Session(service=service, backend="ibmq_qasm_simulator") as session:
    sampler = Sampler(session=session, options=options)
    job = sampler.run(circuits=bell)
    print(job.result())

    # You can invoke run() multiple times.
    job = sampler.run(circuits=bell)
    print(job.result())

## REM sampler

from qiskit import QuantumCircuit
from qiskit_ibm_runtime import Sampler, Options, Session
bell = QuantumCircuit(2)
bell.h(0)
bell.cx(0, 1)
bell.measure_all()

options_rem = Options(resilience_level=1)

with Session(service=service, backend="ibmq_qasm_simulator") as session:
    sampler = Sampler(session=session)
    job = sampler.run(circuits=bell)

    sampler_rem = Sampler(session=session, options=options_rem)
    job_rem = sampler_rem.run(circuits=bell)

    print(f"Quasi_dist without EM {job.result().quasi_dists}")
    print(f"Quasi_dist with EM {job_rem.result().quasi_dists}")