import numpy as np
import matplotlib.pyplot as plt
import imageio.v2 as imageio

from qiskit import execute, transpile
from qiskit.circuit import ClassicalRegister, QuantumCircuit, QuantumRegister
from qiskit.providers.aer.backends import AerSimulator
from qiskit.visualization import plot_histogram
from qiskit.quantum_info import DensityMatrix, state_fidelity, partial_trace
from skimage import data
from skimage.color import rgb2gray
from skimage.transform import resize

from neqr import NEQR

image_neqr = NEQR()
shots = 8192
backend = AerSimulator()

astronaut = data.astronaut()

astronaut.shape

plt.imshow(astronaut)
plt.show()

gray_astro = rgb2gray(astronaut)
plt.imshow(gray_astro, cmap="gray")
plt.show()

resized_gray_astro = resize(gray_astro, (2,2))
plt.imshow(resized_gray_astro, cmap="gray")

qubits_idx = QuantumRegister(size=2, name="qubits_idx")
intensity = QuantumRegister(size=8, name="intensity")
bits_idx = ClassicalRegister(size=2, name="bits_idx")
bits_intensity = ClassicalRegister(size=8, name="bits_intensity")

qc = QuantumCircuit(intensity, qubits_idx, bits_intensity, bits_idx)
qc.draw(output="mpl")

qc.h(qubit=qubits_idx)
qc.barrier()
qc.draw(output="mpl")

pixel_pos_intensity = []

for row in resized_gray_astro:
    for item in row:
        pixel_pos_intensity.append(int(np.round(255*item)))
        
bin_list = [bin(num)[2:] for num in pixel_pos_intensity]

for idx, bnum in enumerate(bin_list):
    
    if idx == 0:
        qc.x(qubit=qubits_idx)
        for pos, item in enumerate(bnum[::-1]):
            if item == "1":
                qc.mct(control_qubits=qubits_idx, target_qubit=intensity[pos])
        qc.x(qubit=qubits_idx)
        qc.barrier()
    elif idx == 1:
        qc.x(qubit=qubits_idx[1])
        for pos, item in enumerate(bnum[::-1]):
            if item == "1":
                qc.mct(control_qubits=qubits_idx, target_qubit=intensity[pos])
        qc.x(qubit=qubits_idx[1])
        qc.barrier()
    elif idx == 2:
        qc.x(qubit=qubits_idx[0])
        for pos, item in enumerate(bnum[::-1]):
            if item == "1":
                qc.mct(control_qubits=qubits_idx, target_qubit=intensity[pos])
        qc.x(qubit=qubits_idx[0])
        qc.barrier()
    else:
        for pos, item in enumerate(bnum[::-1]):
            if item == "1":
                qc.mct(control_qubits=qubits_idx, target_qubit=intensity[pos])
        qc.barrier()
        
qc.draw(output="mpl")

qc.measure(qubit=intensity, cbit=bits_intensity)
qc.barrier()
qc.measure(qubit=qubits_idx, cbit=bits_idx)
qc.draw(output="mpl")

print("Circuit dimensions")
print(f"Circuit depth: {qc.depth()}")
print(f"Circuit size: {qc.size()}")
print(f"Circuit operations: {qc.count_ops()}")

transpiled_qc = transpile(circuits=qc, basis_gates=["id", "x", "sx", "cx", "rz"])
transpiled_qc.draw(output="mpl")

print("Circuit dimensions")
print(f"Circuit depth: {transpiled_qc.depth()}")
print(f"Circuit size: {transpiled_qc.size()}")
print(f"Circuit operations: {transpiled_qc.count_ops()}")

qcircuit = image_neqr.image_quantum_circuit(image=resized_gray_astro, measurements=True)

qcircuit.draw(output="mpl")

print("Circuit dimensions")
print(f"Circuit depth: {qcircuit.depth()}")
print(f"Circuit size: {qcircuit.size()}")
print(f"Circuit operations: {qcircuit.count_ops()}")

counts = execute(experiments=qcircuit, backend=backend, shots=shots).result().get_counts()
plot_histogram(counts)

resized_gray_astro_33 = resize(gray_astro, (3,3))
plt.imshow(resized_gray_astro_33, cmap="gray")

qcircuit = image_neqr.image_quantum_circuit(image=resized_gray_astro_33, measurements=True)

qcircuit.draw(output="mpl")

print("Circuit dimensions")
print(f"Circuit depth: {qcircuit.depth()}")
print(f"Circuit size: {qcircuit.size()}")
print(f"Circuit operations: {qcircuit.count_ops()}")

counts = execute(experiments=qcircuit, backend=backend, shots=shots).result().get_counts()

keys_list = [key for key, _ in sorted(counts.items())][:9]
processed_counts = {key: counts[key] for key in keys_list}
plot_histogram(processed_counts)

test_matrix = np.array([[0,0,0],[0,1,0],[0,1,0],[0,0,0]])

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

qcircuit = image_neqr.image_quantum_circuit(image=test_matrix, measurements=True)

qcircuit.draw(output="mpl")

print("Circuit dimensions")
print(f"Circuit depth: {qcircuit.depth()}")
print(f"Circuit size: {qcircuit.size()}")
print(f"Circuit operations: {qcircuit.count_ops()}")

counts = execute(experiments=qcircuit, backend=backend, shots=shots).result().get_counts()
plot_histogram(counts)

test_matrix = np.array([[0,1,0]])

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

qcircuit = image_neqr.image_quantum_circuit(image=test_matrix, measurements=True)

qcircuit.draw(output="mpl")

print("Circuit dimensions")
print(f"Circuit depth: {qcircuit.depth()}")
print(f"Circuit size: {qcircuit.size()}")
print(f"Circuit operations: {qcircuit.count_ops()}")

counts = execute(experiments=qcircuit, backend=backend, shots=shots).result().get_counts()

keys_list = [key for key, _ in sorted(counts.items())][:3]
processed_counts = {key: counts[key] for key in keys_list}
plot_histogram(processed_counts)

test_image = np.array([[0,0],[0,0]])
plt.imshow(test_image, cmap="gray")

qc = image_neqr.image_quantum_circuit(image=test_image)
qc.draw(output="mpl")

dm = DensityMatrix(data=qc)

dm

dm_reduced = partial_trace(state=dm, qargs=[8,9])

dm_reduced

state_fidelity(dm_reduced, dm_reduced)

test_image2 = np.array([[1,1],[1,1]])
plt.imshow(test_image2, cmap="gray")

qc2 = image_neqr.image_quantum_circuit(image=test_image2)
qc2.draw(output="mpl")

dm2 = DensityMatrix(data=qc2)

dm2

dm_reduced2 = partial_trace(state=dm2, qargs=[8,9])

dm_reduced2

state_fidelity(dm_reduced, dm_reduced2)

resized_rgb_astro = resize(astronaut, (2,2))
plt.imshow(resized_rgb_astro)

qc = image_neqr.image_quantum_circuit(image=resized_rgb_astro, measurements=True)
qc.draw(output="mpl")

print("Circuit dimensions")
print(f"Circuit depth: {qc.depth()}")
print(f"Circuit size: {qc.size()}")
print(f"Circuit operations: {qc.count_ops()}")

counts = execute(experiments=qc, backend=backend, shots=shots).result().get_counts()
keys_list = [key for key, _ in sorted(counts.items())][:12]
processed_counts = {key: counts[key] for key in keys_list}
plot_histogram(processed_counts)