from qiskit.circuit import QuantumCircuit
from qiskit.primitives import Estimator, Sampler
from qiskit.quantum_info import SparsePauliOp
from qiskit.visualization import plot_histogram

import numpy as np
import matplotlib.pyplot as plt
plt.style.use('dark_background') # optional

# create excited |1> state
qc_1 = QuantumCircuit(1)
qc_1.x(0)
qc_1.draw('mpl')




# create superposition |+> state
qc_plus = QuantumCircuit(1)
qc_plus.h(0)
qc_plus.draw('mpl')

qc_1.measure_all()
qc_plus.measure_all()

sampler = Sampler()
job_1 = sampler.run(qc_1)
job_plus = sampler.run(qc_plus)

job_1.result().quasi_dists

job_plus.result().quasi_dists

legend = ["Excited State", "Plus State"] # TODO: Excited State does not appear
plot_histogram([job_1.result().quasi_dists[0], job_plus.result().quasi_dists[0]], legend=legend)

qc_1.remove_final_measurements()
qc_plus.remove_final_measurements()

# rotate into the X-basis
qc_1.h(0)
qc_plus.h(0)

qc_1.measure_all()
qc_plus.measure_all()

sampler = Sampler()
job_1 = sampler.run(qc_1)
job_plus = sampler.run(qc_plus)

plot_histogram([job_1.result().quasi_dists[0], job_plus.result().quasi_dists[0]], legend=legend)

qc2_1 = QuantumCircuit(1)
qc2_1.x(0)

qc2_plus = QuantumCircuit(1)
qc2_plus.h(0)

obsvs = list(SparsePauliOp(['Z', 'X']))

estimator = Estimator()
job2_1 = estimator.run([qc2_1]*len(obsvs), observables=obsvs)
job2_plus = estimator.run([qc2_plus]*len(obsvs), observables=obsvs)

job2_1.result()

# TODO: make this into module that outputs a nice table
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]}')

A, a, B, b = SparsePauliOp(["Z", "X", "Z", "X"])

obsv = A.tensor(B) - A.tensor(b)+ a.tensor(B) + a.tensor(b)  # create operator for chsh witness

from qc_grader.challenges.qgss_2023 import grade_lab2_ex1

grade_lab2_ex1(obsv)

from qiskit.circuit import Parameter

theta = Parameter('θ')

qc = QuantumCircuit(2)
qc.h(0)
qc.cx(0, 1)
qc.ry(theta, 0)

qc.draw('mpl')

angles = [[float(v)] for v in np.linspace(0, 2*np.pi, 10).astype(np.float32)]
# create a parameterization of angles that will violate the inequality

estimator = Estimator()
job = estimator.run([qc]*len(angles), observables=[obsv]*len(angles), parameter_values=angles)
exps = job.result().values

plt.plot(angles, exps, marker='x', ls='-', color='green')
plt.plot(angles, [2]*len(angles), ls='--', color='red', label='Classical Bound')
plt.plot(angles, [-2]*len(angles), ls='--', color='red')
plt.xlabel('angle (rad)')
plt.ylabel('CHSH Witness')
plt.legend(loc=4)

from qc_grader.challenges.qgss_2023 import grade_lab2_ex2

grade_lab2_ex2(obsv, angles)

from qiskit.circuit import ClassicalRegister, QuantumRegister

theta = Parameter('θ')

qr = QuantumRegister(1, 'q')
qc = QuantumCircuit(qr)
qc.ry(theta, 0)
qc.draw('mpl')

tele_qc = qc.copy()
bell = QuantumRegister(2, 'Bell')
alice = ClassicalRegister(2, 'Alice')
bob = ClassicalRegister(1, 'Bob')
tele_qc.add_register(bell, alice, bob)
tele_qc.draw('mpl')

# create Bell state with other two qubits
tele_qc.barrier()
tele_qc.h(1)
tele_qc.cx(1, 2)
tele_qc.barrier()
tele_qc.draw('mpl')

# alice operates on her qubits
tele_qc.cx(0, 1)
tele_qc.h(0)
tele_qc.barrier()
tele_qc.draw('mpl')

tele_qc.measure([qr[0], bell[0]], alice)
tele_qc.draw('mpl')

graded_qc = tele_qc.copy()

##############################
# add gates to graded_qc here
with graded_qc.if_test((alice[1], 1)) as _:
    graded_qc.x(2)

with graded_qc.if_test((alice[0], 1)) as _:
    graded_qc.z(2)

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

graded_qc.draw('mpl')

graded_qc.barrier()
graded_qc.measure(bell[1], bob)
graded_qc.draw('mpl')

from qc_grader.challenges.qgss_2023 import grade_lab2_ex3

grade_lab2_ex3(graded_qc, theta, 5*np.pi/7)

from qiskit_aer.primitives import Sampler

angle = 5*np.pi/7

sampler = Sampler()
qc.measure_all()
job_static = sampler.run(qc.bind_parameters({theta: angle}))
job_dynamic = sampler.run(graded_qc.bind_parameters({theta: angle}))

print(f"Original Dists: {job_static.result().quasi_dists[0].binary_probabilities()}")
print(f"Teleported Dists: {job_dynamic.result().quasi_dists[0].binary_probabilities()}")

from qiskit.result import marginal_counts

tele_counts = marginal_counts(
    job_dynamic.result().quasi_dists[0].binary_probabilities(),
    indices=[2],
) # marginalize counts

legend = ['Original State', 'Teleported State']
plot_histogram([job_static.result().quasi_dists[0].binary_probabilities(), tele_counts], legend=legend)

from qc_grader.challenges.qgss_2023 import grade_lab2_ex4

grade_lab2_ex4(tele_counts, job_dynamic.result())

import qiskit.tools.jupyter
%qiskit_version_table

