# required import of libraries and modules

# import quantum circuit from IBM's Qiskit circuit module
from qiskit.circuit import QuantumCircuit

# import estimator and sample
# from IBM's Qiskit primitives module
from qiskit.primitives import Estimator, Sampler

# import Sparse Pauli operator
# from IBM's Qiskit quantum information module
from qiskit.quantum_info import SparsePauliOp

# import histogram plotting
# from IBM's Qiskit visualization module
from qiskit.visualization import plot_histogram

# import numpy
import numpy as np

# import plotting from Matplotlib's pyplot module
import matplotlib.pyplot as plt

# OPTIONAL: use the dark background style
# for the plotting from Matplotlib's pyplot module
plt.style.use("dark_background")

# create an excited state |1>

# create a quantum circuit for one qubit
# and to implement an excited state |1>
qc_1 = QuantumCircuit(1)

# apply a Pauli-X (bit flip / not) gate to
# the quantum circuit, on the 1st qubit
qc_1.x(0)

# draw the quantum circuit implementing
# the excited state |1>
qc_1.draw("mpl")

# create quantum superposition state |+>

# create a quantum circuit for one qubit
# and to implement a quantum superposition state |+>
qc_plus = QuantumCircuit(1)

# apply a Hadamard gate to
# the quantum circuit, on the 1st qubit
qc_plus.h(0)

# draw the quantum circuit implementing
# the quantum superposition state |+>
qc_plus.draw("mpl")

# measure all the qubits of
# the quantum circuit implementing
# an excited state |1>
qc_1.measure_all()

# measure all the qubits of
# the quantum circuit implementing
# a quantum superposition state |+>
qc_plus.measure_all()

# create a Sampler object
sampler = Sampler()

# create a job for the quantum circuit implementing
# an excited state |1>, running the Sampler
job_1 = sampler.run(qc_1)

# create a job for the quantum circuit implementing
# a quantum superposition state |+>, running the Sampler
job_plus = sampler.run(qc_plus)

# retrieve the quasi-distribution of
# the job for the quantum circuit
# implementing an excited state |1>
quasi_dists_ket_1 = job_1.result().quasi_dists

# fix the legend issue for
# the quasi-distribution of the outcome 0
quasi_dists_ket_1[0][0] = 1e-20

# show the quasi-distribution of
# the job for the quantum circuit
# implementing an excited state |1>
quasi_dists_ket_1

# retrieve the quasi-distribution of
# the job for the quantum circuit implementing
# a quantum superposition state |+>
quasi_dists_ket_plus = job_plus.result().quasi_dists

# show the quasi-distribution of
# the job for the quantum circuit
# implementing a quantum superposition state |+>
quasi_dists_ket_plus

# create the legend for the two quantum states
# created before, namely the excited state |1> and
# the quantum superposition state |+>
legend = ["Excited State", "Plus State"]

# plot the histogram for the quasi-distributions
# for the two quantum states created before,
# namely the excited state |1> and
# the quantum superposition state |+>
plot_histogram( [ quasi_dists_ket_1[0], quasi_dists_ket_plus[0] ],
                  legend=legend, title="Measurement Results\n" )

# remove all the qubits' measurements of
# the quantum circuit implementing
# an excited state |1>
qc_1.remove_final_measurements()

# remove all the qubits' measurements of
# the quantum circuit implementing
# a quantum superposition state |+>
qc_plus.remove_final_measurements()


# prepare the 1st qubit of the quantum circuit
# implementing an excited state |1>,
# for the X-Basis measurement,
# applying a Hadamard gate
qc_1.h(0)

# prepare the 1st qubit of the quantum circuit
# implementing a quantum superposition state |+>,
# for the X-Basis measurement,
# applying a Hadamard gate
qc_plus.h(0)


# measure all the qubits of
# the quantum circuit implementing
# an excited state |1>
qc_1.measure_all()

# measure all the qubits of
# the quantum circuit implementing
# a quantum superposition state |+>
qc_plus.measure_all()

# create a Sampler object
sampler = Sampler()

# create a job for the quantum circuit implementing
# an excited state |1>, running the Sampler
job_1 = sampler.run(qc_1)

# create a job for the quantum circuit implementing
# a quantum superposition state |+>, running the Sampler
job_plus = sampler.run(qc_plus)

# retrieve the quasi-distribution of
# the job for the quantum circuit
# implementing an excited state |1>,
# measured in the X-Basis
quasi_dists_ket_1 = job_1.result().quasi_dists

# show the quasi-distribution of
# the job for the quantum circuit
# implementing an excited state |1>,
# measured in the X-Basis
quasi_dists_ket_1

# retrieve the quasi-distribution of
# the job for the quantum circuit
# implementing a quantum superposition state |+>,
# measured in the X-Basis
quasi_dists_ket_plus = job_plus.result().quasi_dists

# fix the legend issue for
# the quasi-distribution of the outcome 1
quasi_dists_ket_plus[0][1] = 1e-20

# show the quasi-distribution of
# the job for the quantum circuit
# implementing a quantum superposition state |+>,
# measured in the X-Basis
quasi_dists_ket_plus

# plot the histogram for the quasi-distributions
# for the two quantum states created before,
# namely the excited state |1> and
# the quantum superposition state |+>,
# measured in the X-Basis
plot_histogram( [ quasi_dists_ket_1[0], quasi_dists_ket_plus[0] ],
                  legend=legend, title="Measurement Results\n" )

# create a quantum circuit for one qubit
# and to implement an excited state |1>
qc2_1 = QuantumCircuit(1)

# apply a Pauli-X (bit flip / not) gate to
# the quantum circuit, on the 1st qubit
qc2_1.x(0)


# create a quantum circuit for one qubit
# and to implement a quantum superposition state |+>
qc2_plus = QuantumCircuit(1)

# apply a Hadamard gate to
# the quantum circuit, on the 1st qubit
qc2_plus.h(0)


# create a list for the observables for
# the Sparse Pauli-Z and Pauli-X operators
obsvs = list( SparsePauliOp( [ "Z", "X" ] ) )

# crate an Estimator object
estimator = Estimator()

# create a job for the quantum circuit implementing
# an excited state |1>, running the Estimator
job2_1 = estimator.run( [ qc2_1 ] * len(obsvs), observables=obsvs )

# create a job for the quantum circuit implementing
# a quantum superposition state |+>, running the Estimator
job2_plus = estimator.run( [ qc2_plus ] * len(obsvs), observables=obsvs )

# show the results of the job ran by
# the Estimator created before,
# that was ran on the excited state |1>
job2_1.result()

# print a table with results (and eigenvalues) for
# the quantum states created before and measured in the bases set before
print(f'    |   <Z>   |   <X> ')
print(f'----|-------------------------------')
print(f'|1> | {job2_1.result().values[0]}    | {job2_1.result().values[1]}')
print(f'|+> |  {job2_plus.result().values[0]}    | {job2_plus.result().values[1]}')

# import the Sparse Pauli operator from
# the IBM's Qiskit quantum information module
from qiskit.quantum_info import SparsePauliOp


# For Alice: A = Z ; a = X
# For Bob: B = Z ; b = X

# create the Sparse Pauli operator ZZ, for <AB>
ZZ = SparsePauliOp.from_list( [ ( "ZZ", 1 ) ] )

# create the Sparse Pauli operator ZX, for <Ab>
ZX = SparsePauliOp.from_list( [ ( "ZX", 1 ) ] )

# create the Sparse Pauli operator XZ, for <aB>
XZ = SparsePauliOp.from_list( [ ( "XZ", 1 ) ] )

# create the Sparse Pauli operator XX, for <ab>
XX = SparsePauliOp.from_list( [ ( "XX", 1 ) ] )


# create the list of required
# operators ZZ, ZX, XZ, and XX
ops = [ ZZ, ZX, XZ, XX ]


# there are two possible CHSH witnesses

# create operator for CHSH witness <CHSH_1>
#     <CHSH_1> = <ZZ> - <ZX> + <XZ> + <XX> (=)
# (=) <CHSH_1> = <AB> - <Ab> + <aB> + <ab>
#obsv = ops[0] - ops[1] + ops[2] + ops[3]

# create operator for CHSH witness <CHSH_2>
#     <CHSH_2> = <ZZ> + <ZX> - <XZ> + <XX> (=)
# (=) <CHSH_2> = <AB> + <Ab> - <aB> + <ab>
obsv = ( ops[0] + ops[1] - ops[2] + ops[3] )

# import the grader for the exercise 1 of the lab 2
from qc_grader.challenges.qgss_2023 import grade_lab2_ex1

# grade the exercise 1 of the lab 2
grade_lab2_ex1( obsv )

# import the parameter from
# the IBM's Qiskit circuit module
from qiskit.circuit import Parameter

# create a parameter for the theta angle
theta = Parameter("θ")


# create a quantum circuit for two qubits
# and to implement a parametrized quantum entangled pair
qc = QuantumCircuit(2)


# apply a Hadamard gate to
# the quantum circuit, on the 1st qubit
qc.h(0)

# apply a CX/CNOT gate to
# the quantum circuit,
# with control on 1st qubit
# and target on 2nd qubit
qc.cx(0, 1)

# apply a RY (Rotation-Y) gate to
# the quantum circuit, on the 1st qubit,
# with the theta angle defined before
qc.ry(theta, 0)


# draw the quantum circuit implementing
# a parametrized quantum entangled pair
qc.draw("mpl")

# create a parameterization of theta angles that
# will violate the CHSH inequality

# number of theta angles to violate the CHSH inequality
number_of_phases = 20


# for the 1st CHSH witness
#     <CHSH_1> = <ZZ> - <ZX> + <XZ> + <XX> (=)
# (=) <CHSH_1> = <AB> - <Ab> + <aB> + <ab>
#phases_chsh_1_1 = np.linspace(0.5 * np.pi, np.pi, number_of_phases)
#phases_chsh_1_2 = np.linspace(1.5 * np.pi, 2 * np.pi, number_of_phases)


# create the two lists of theta angles that
# violate the CHSH inequality
#angles_chsh_1_1 = [[ph] for ph in phases_chsh_1_1]
#angles_chsh_1_2 = [[ph] for ph in phases_chsh_1_2]


# merge the two lists of theta angles that
# violate the CHSH inequality, for the final result
#angles = angles_chsh_1_1 + angles_chsh_1_2



# for the 2nd CHSH witness
#     <CHSH_2> = <ZZ> + <ZX> - <XZ> + <XX> (=)
# (=) <CHSH_2> = <AB> + <Ab> - <aB> + <ab>
phases_chsh_2_1 = np.linspace( 0, ( 0.5 * np.pi ), number_of_phases )
phases_chsh_2_2 = np.linspace( ( 1 * np.pi ) , ( 1.5 * np.pi ), number_of_phases )


# create the two lists of theta angles that
# violate the CHSH inequality
angles_chsh_2_1 = [ [ph] for ph in phases_chsh_2_1 ]
angles_chsh_2_2 = [ [ph] for ph in phases_chsh_2_2 ]


# merge the two lists of theta angles that
# violate the CHSH inequality, for the final result
angles = angles_chsh_2_1 + angles_chsh_2_2

# crate an Estimator object
estimator = Estimator()

# create a job for the quantum circuit implementing
# a parametrized quantum entangled pair,
# varying the elements in the list of
# theta angles defined before
job = estimator.run( [qc] * len(angles),
                     observables=[obsv] * len(angles),
                     parameter_values=angles )


# retrieve the expectation values
# from the results of the job ran by
# the Estimator created before,
# that was ran on the parametrized
# quantum entangled pair
exps = job.result().values


# plot the angles theta parametrized from the list built before
plt.plot( angles, exps, marker='x', ls='-', color="green" )

# plot the two classical bounds that violates the CHSH inequality
plt.plot( angles, [2]*len(angles), ls="--", color="red", label="Classical Bound" )
plt.plot( angles, [-2]*len(angles), ls="--", color="red" )


# add the X-axis label to the plot built before
plt.xlabel("angle (rad)")

# add the Y-axis label to the plot built before
plt.ylabel("CHSH Witness")


# add the title to the plot built before
plt.title("CHSH Witnesses\n")


# add the legend to the plot built before
plt.legend(loc=4)

# import the grader for the exercise 2 of the lab 2
from qc_grader.challenges.qgss_2023 import grade_lab2_ex2

# grade the exercise 2 of the lab 2
grade_lab2_ex2( obsv, angles )

# import the classical and quantum registers
# from the IBM's Qiksit circuit module
from qiskit.circuit import ClassicalRegister, QuantumRegister

# create a parameter for the theta angle
theta = Parameter("θ")


# create a quantum register for one qubit
# and to implement a parametrized quantum operation
qr = QuantumRegister(1, "q")

# create a quantum circuit with
# the quantum register created before
qc = QuantumCircuit(qr)


# apply a RY (Rotation-Y) gate to
# the quantum circuit, on the 1st qubit,
# with the theta angle defined before
qc.ry(theta, 0)


# draw the quantum circuit implementing
# a parametrized quantum operation
qc.draw("mpl")

# create a copy of the quantum circuit defined before
# to initialize the quantum circuit implementing
# a parametrized quantum teleportation
tele_qc = qc.copy()


# create a quantum register with two qubits
# for the Bell State to be shared on
# the parametrized quantum teleportation
bell = QuantumRegister(2, "Bell")

# create a classical register with two qubits
# for Alice's part on the parametrized
# quantum teleportation
alice = ClassicalRegister(2, "Alice")

# create a classical register with one qubit
# for Bob's part on the parametrized
# quantum teleportation
bob = ClassicalRegister(1, "Bob")


# add all the quantum and classical registers
# created before to the quantum circuit implementing
# a parametrized quantum teleportation
tele_qc.add_register(bell, alice, bob)


# draw the quantum circuit implementing
# the 1st step of a parametrized
# quantum teleportation
tele_qc.draw("mpl")

# add a barrier to the quantum circuit implementing
# a parametrized quantum teleportation
tele_qc.barrier()


# create Bell state with the other two qubits,
# in the quantum register defined before
# for that purpose

# apply a Hadamard gate to
# the quantum circuit, on the 2nd qubit
tele_qc.h(1)

# apply a CX/CNOT gate to
# the quantum circuit,
# with control on 2nd qubit
# and target on 3rd qubit
tele_qc.cx(1, 2)


# add a barrier to the quantum circuit implementing
# a parametrized quantum teleportation
tele_qc.barrier()


# draw the quantum circuit implementing
# the 2nd step of a parametrized
# quantum teleportation
tele_qc.draw("mpl")

# operations for Alice perform on her qubits
# during the parametrized quantum teleportation

# apply a CX/CNOT gate to
# the quantum circuit,
# with control on 1sr qubit
# and target on 2nd qubit
tele_qc.cx(0, 1)

# apply a Hadamard gate to
# the quantum circuit, on the 1st qubit
tele_qc.h(0)


# add a barrier to the quantum circuit implementing
# a parametrized quantum teleportation
tele_qc.barrier()

# draw the quantum circuit implementing
# the 3rd step of a parametrized
# quantum teleportation
tele_qc.draw("mpl")

# measurements performed by Alice on
# the qubit q and on her part of the Bell State,
# performing the entanglement swapping
tele_qc.measure( [ qr[0], bell[0] ], alice )

# draw the quantum circuit implementing
# the 4th step of a parametrized
# quantum teleportation
tele_qc.draw("mpl")

# create a copy of the quantum circuit defined before
# to initialize the quantum circuit implementing
# a parametrized quantum teleportation,
# and to be graded with dynamic operations 
graded_qc = tele_qc.copy()


##############################

# add gates to graded_qc here

# Bob, who already has the qubit teleported,
# applies the following gates depending on
# the state of the classical bits:
# 00 -> Do nothing
# 01 -> Apply X gate
# 10 -> Apply Z gate
# 11 -> Apply ZX gate

# apply the Pauli-X (bit flip / not) gate on
# the 2nd qubit of the Bell state if
# the Alice's 2nd classical bit is 1
graded_qc.x( bell[1] ).c_if( alice[1], 1 )

# apply the Pauli-Z (phase flip) gate on
# the 2nd qubit of the Bell state if
# the Alice's 1st classical bit is 1
graded_qc.z( bell[1] ).c_if( alice[0], 1 )

##############################


# draw the quantum circuit implementing
# the parametrized quantum teleportation, to be graded
graded_qc.draw("mpl")

# add a barrier to the quantum circuit implementing
# a parametrized quantum teleportation, to be graded
graded_qc.barrier()


# measure the 2nd qubit of the Bell state to
# the Bob's classical register
graded_qc.measure( bell[1], bob )


# draw the quantum circuit implementing
# a parametrized quantum teleportation, to be graded
graded_qc.draw("mpl")

# import the grader for the exercise 3 of the lab 2
from qc_grader.challenges.qgss_2023 import grade_lab2_ex3

# grade the exercise 3 of the lab 2
grade_lab2_ex3( graded_qc, theta, 5 * ( np.pi / 7 ) )

# impor the Sampler from the IBM's Qiskit
# Aer primitives module
from qiskit_aer.primitives import Sampler


# define the pretended theta angle
angle = 5 * ( np.pi / 7 ) 


# create a Sampler object
sampler = Sampler()


# measure all the qubits on the quantum circuit
# implementing a parametrized quantum teleportation
qc.measure_all()


# create a job for the quantum circuit implementing
# a parametrized quantum teleportation,
# with a static quantum circuit and
# with the theta angle defined before
job_static = sampler.run( qc.bind_parameters({theta: angle}) )

# create a job for the quantum circuit implementing
# a parametrized quantum teleportation,
# with a dynamic quantum circuit and
# with the theta angle defined before
job_dynamic = sampler.run( graded_qc.bind_parameters({theta: angle}) )


# print the measurement distributions of the qubits from
# the parametrized quantum teleportation, with a static quantum circuit
print( f"Original Dists.: {job_static.result().quasi_dists[0].binary_probabilities()}" )

# print the measurement distributions of the qubits from
# the parametrized quantum teleportation, with a dynamic quantum circuit
print( f"Teleported Dists.: {job_dynamic.result().quasi_dists[0].binary_probabilities()}" )

# import the counts marginalization from
# the IBM's Qiskit result module
from qiskit.result import marginal_counts

# marginalize counts from the quasi-distributitions
# obtained from the parametrized quantum teleportation
tele_counts = marginal_counts( job_dynamic.result().quasi_dists[0]
                               .binary_probabilities(), [2] )

# create the legend for the two quantum states
# created before, namely the original quantum state to
# be teleported and the quantum state which was indeed teleported
legend = [ "Original State", "Teleported State" ]

# plot the histogram for the quasi-distributions
# for the two quantum states created before,
# namely the original quantum state to be teleported
# and the quantum state which was indeed teleported
plot_histogram( [ job_static.result().quasi_dists[0]
                  .binary_probabilities(), tele_counts ],
                legend=legend, title="Measurement Results\n" )

# import the grader for the exercise 4 of the lab 2
from qc_grader.challenges.qgss_2023 import grade_lab2_ex4

# grade the exercise 4 of the lab 2
grade_lab2_ex4( tele_counts, job_dynamic.result() )

# import the IBM's Qiskit Jupyter Tools
import qiskit.tools.jupyter

# show the table of the IBM's Qiskit version
%qiskit_version_table