import matplotlib.pyplot as plt

from neqr import NEQR
from qiskit import execute
from qiskit.providers.aer.backends import AerSimulator
from qiskit.quantum_info import Statevector
from skimage import data
from skimage.color import rgb2gray
from skimage.transform import resize

neqr_class = NEQR()
backend = AerSimulator()
shots = 2**14

astro_pic = resize(data.astronaut(), (32,32))
gray_astro_pic = rgb2gray(astro_pic)

plt.imshow(gray_astro_pic, cmap="gray")

qc_gray = neqr_class.image_quantum_circuit(image=gray_astro_pic, measurements=True)
counts = backend.run(qc_gray, shots=shots).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=gray_astro_pic.shape)

plt.imshow(image, cmap="gray")

qc_gray = neqr_class.image_quantum_circuit(image=gray_astro_pic)
qc_gray.x(qubit=qc_gray.qregs[0])
qc_gray.measure(qubit=qc_gray.qregs[0], cbit=qc_gray.cregs[0])
qc_gray.measure(qubit=qc_gray.qregs[1], cbit=qc_gray.cregs[1])
counts = backend.run(qc_gray, shots=shots).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=gray_astro_pic.shape)

plt.imshow(image, cmap="gray")

qc_gray = neqr_class.image_quantum_circuit(image=gray_astro_pic)
qc_gray.x(qubit=qc_gray.qregs[1][:int(qc_gray.qregs[1].size/2)])
qc_gray.measure(qubit=qc_gray.qregs[0], cbit=qc_gray.cregs[0])
qc_gray.measure(qubit=qc_gray.qregs[1], cbit=qc_gray.cregs[1])
counts = backend.run(qc_gray, shots=shots).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=gray_astro_pic.shape)

plt.imshow(image, cmap="gray")

qc_gray = neqr_class.image_quantum_circuit(image=gray_astro_pic)
qc_gray.x(qubit=qc_gray.qregs[1][int(qc_gray.qregs[1].size/2):])
qc_gray.measure(qubit=qc_gray.qregs[0], cbit=qc_gray.cregs[0])
qc_gray.measure(qubit=qc_gray.qregs[1], cbit=qc_gray.cregs[1])
counts = backend.run(qc_gray, shots=shots).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=gray_astro_pic.shape)

plt.imshow(image, cmap="gray")

qc_gray = neqr_class.image_quantum_circuit(image=gray_astro_pic)
qc_gray.x(qubit=qc_gray.qregs[1])
qc_gray.measure(qubit=qc_gray.qregs[0], cbit=qc_gray.cregs[0])
qc_gray.measure(qubit=qc_gray.qregs[1], cbit=qc_gray.cregs[1])
counts = backend.run(qc_gray, shots=shots).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=gray_astro_pic.shape)

plt.imshow(image, cmap="gray")

qc_gray = neqr_class.image_quantum_circuit(image=gray_astro_pic)
qc_gray.x(qubit=qc_gray.qregs[0])
qc_gray.x(qubit=qc_gray.qregs[1])
qc_gray.measure(qubit=qc_gray.qregs[0], cbit=qc_gray.cregs[0])
qc_gray.measure(qubit=qc_gray.qregs[1], cbit=qc_gray.cregs[1])
counts = backend.run(qc_gray, shots=shots).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=gray_astro_pic.shape)

plt.imshow(image, cmap="gray")

gray_astro_pic_full_size = rgb2gray(data.astronaut())
qc_gray = neqr_class.image_quantum_circuit(image=gray_astro_pic_full_size, measurements=True)

print(f"Number of qubits: {qc_gray.qregs[0].size + qc_gray.qregs[1].size}")
print(f"Circuit size: {qc_gray.size()}")
print(f"Circuit depth: {qc_gray.depth()}")
print(f"Circuit ops: {qc_gray.count_ops()}")

plt.imshow(gray_astro_pic_full_size, cmap="gray")

plt.imshow(astro_pic)

qc_rgb = neqr_class.image_quantum_circuit(image=astro_pic, measurements=True)
counts = backend.run(qc_rgb, shots=shots*4).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=astro_pic.shape)

plt.imshow(image)

qc_rgb = neqr_class.image_quantum_circuit(image=astro_pic)
qc_rgb.x(qubit=qc_rgb.qregs[1][:int(qc_rgb.qregs[1].size/2)])
qc_rgb.measure(qubit=qc_rgb.qregs[0], cbit=qc_rgb.cregs[0])
qc_rgb.measure(qubit=qc_rgb.qregs[1], cbit=qc_rgb.cregs[1])
qc_rgb.measure(qubit=qc_rgb.qregs[2], cbit=qc_rgb.cregs[2])
counts = backend.run(qc_rgb, shots=shots*4).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=astro_pic.shape)

plt.imshow(image)

qc_rgb = neqr_class.image_quantum_circuit(image=astro_pic)
qc_rgb.x(qubit=qc_rgb.qregs[1][int(qc_rgb.qregs[1].size/2):])
qc_rgb.measure(qubit=qc_rgb.qregs[0], cbit=qc_rgb.cregs[0])
qc_rgb.measure(qubit=qc_rgb.qregs[1], cbit=qc_rgb.cregs[1])
qc_rgb.measure(qubit=qc_rgb.qregs[2], cbit=qc_rgb.cregs[2])
counts = backend.run(qc_rgb, shots=shots*4).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=astro_pic.shape)

plt.imshow(image)

qc_rgb = neqr_class.image_quantum_circuit(image=astro_pic)
qc_rgb.x(qubit=qc_rgb.qregs[1])
qc_rgb.measure(qubit=qc_rgb.qregs[0], cbit=qc_rgb.cregs[0])
qc_rgb.measure(qubit=qc_rgb.qregs[1], cbit=qc_rgb.cregs[1])
qc_rgb.measure(qubit=qc_rgb.qregs[2], cbit=qc_rgb.cregs[2])
counts = backend.run(qc_rgb, shots=shots*4).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=astro_pic.shape)

plt.imshow(image)

qc_rgb = neqr_class.image_quantum_circuit(image=astro_pic)
qc_rgb.x(qubit=qc_rgb.qregs[0])
qc_rgb.measure(qubit=qc_rgb.qregs[0], cbit=qc_rgb.cregs[0])
qc_rgb.measure(qubit=qc_rgb.qregs[1], cbit=qc_rgb.cregs[1])
qc_rgb.measure(qubit=qc_rgb.qregs[2], cbit=qc_rgb.cregs[2])
counts = backend.run(qc_rgb, shots=shots*4).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=astro_pic.shape)

plt.imshow(image)

qc_rgb = neqr_class.image_quantum_circuit(image=astro_pic)
qc_rgb.x(qubit=qc_rgb.qregs[2])
for i in range(8):
    qc_rgb.mct(control_qubits=qc_rgb.qregs[2], target_qubit=qc_rgb.qregs[0][i])
qc_rgb.x(qubit=qc_rgb.qregs[2])
qc_rgb.measure(qubit=qc_rgb.qregs[0], cbit=qc_rgb.cregs[0])
qc_rgb.measure(qubit=qc_rgb.qregs[1], cbit=qc_rgb.cregs[1])
qc_rgb.measure(qubit=qc_rgb.qregs[2], cbit=qc_rgb.cregs[2])
counts = backend.run(qc_rgb, shots=shots*4).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=astro_pic.shape)

plt.imshow(image)

qc_rgb = neqr_class.image_quantum_circuit(image=astro_pic)
qc_rgb.x(qubit=qc_rgb.qregs[2][1])
for i in range(8):
    qc_rgb.mct(control_qubits=qc_rgb.qregs[2], target_qubit=qc_rgb.qregs[0][i])
qc_rgb.x(qubit=qc_rgb.qregs[2][1])
qc_rgb.measure(qubit=qc_rgb.qregs[0], cbit=qc_rgb.cregs[0])
qc_rgb.measure(qubit=qc_rgb.qregs[1], cbit=qc_rgb.cregs[1])
qc_rgb.measure(qubit=qc_rgb.qregs[2], cbit=qc_rgb.cregs[2])
counts = backend.run(qc_rgb, shots=shots*4).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=astro_pic.shape)

plt.imshow(image)

qc_rgb = neqr_class.image_quantum_circuit(image=astro_pic)
qc_rgb.x(qubit=qc_rgb.qregs[2][0])
for i in range(8):
    qc_rgb.mct(control_qubits=qc_rgb.qregs[2], target_qubit=qc_rgb.qregs[0][i])
qc_rgb.x(qubit=qc_rgb.qregs[2][0])
qc_rgb.measure(qubit=qc_rgb.qregs[0], cbit=qc_rgb.cregs[0])
qc_rgb.measure(qubit=qc_rgb.qregs[1], cbit=qc_rgb.cregs[1])
qc_rgb.measure(qubit=qc_rgb.qregs[2], cbit=qc_rgb.cregs[2])
counts = backend.run(qc_rgb, shots=shots*4).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=astro_pic.shape)

plt.imshow(image)

qc_rgb = neqr_class.image_quantum_circuit(image=astro_pic)
qc_rgb.x(qubit=qc_rgb.qregs[0])
qc_rgb.x(qubit=qc_rgb.qregs[1])
qc_rgb.measure(qubit=qc_rgb.qregs[0], cbit=qc_rgb.cregs[0])
qc_rgb.measure(qubit=qc_rgb.qregs[1], cbit=qc_rgb.cregs[1])
qc_rgb.measure(qubit=qc_rgb.qregs[2], cbit=qc_rgb.cregs[2])
counts = backend.run(qc_rgb, shots=shots*4).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=astro_pic.shape)

plt.imshow(image)

cat = data.cat()
cat_gray_resized = resize(rgb2gray(cat), (int(cat.shape[0]/10), int(cat.shape[1]/10)))

plt.imshow(cat_gray_resized, cmap="gray")

qc_gray = neqr_class.image_quantum_circuit(image=cat_gray_resized, measurements=True)
counts = backend.run(qc_gray, shots=shots).result().get_counts()
image = neqr_class.reconstruct_image_from_neqr_result(counts=counts, image_shape=cat_gray_resized.shape)

plt.imshow(image, cmap="gray")

print(f"Number of qubits: {qc_gray.qregs[0].size + qc_gray.qregs[1].size}")
print(f"Circuit size: {qc_gray.size()}")
print(f"Circuit depth: {qc_gray.depth()}")
print(f"Circuit ops: {qc_gray.count_ops()}")