import numpy as np
import matplotlib.pyplot as plt

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

SHOTS = 8192
BACKEND = AerSimulator()
image_frqi = FRQI()

mcry = RYGate(theta=2*Parameter(r"$\theta$")).control(num_ctrl_qubits=2)

ry_qc = QuantumCircuit(1)
ry_qc.ry(theta=2*Parameter(r"$\theta$"), qubit=0)
mcry_alt = ry_qc.to_gate(label="mcry").control(num_ctrl_qubits=2)

test = QuantumCircuit(3)
test.append(mcry, [0,1,2])
test.draw(output="mpl")

qubits = QuantumRegister(size=2, name="pixels")
color_qubit = QuantumRegister(size=1, name="color")
bits = ClassicalRegister(size=2, name="bits_pixels")
color_bit = ClassicalRegister(size=1)

qc = QuantumCircuit(qubits, color_qubit, bits, color_bit)

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

theta = np.pi/4

for i in range(4):
    if i == 0:
        qc.x(qubit=qubits)
    elif i == 1:
        qc.x(qubit=qubits[1])
    elif i == 2:
        qc.x(qubit=qubits[0])
    qc.cry(theta=theta, control_qubit=qubits[0], target_qubit=color_qubit)
    qc.cx(control_qubit=qubits[0], target_qubit=qubits[1])
    qc.cry(theta=-theta, control_qubit=qubits[1], target_qubit=color_qubit)
    qc.cx(control_qubit=qubits[0], target_qubit=qubits[1])
    qc.cry(theta=theta, control_qubit=qubits[1], target_qubit=color_qubit)
    if i == 0:
        qc.x(qubit=qubits)
    elif i == 1:
        qc.x(qubit=qubits[1])
    elif i == 2:
        qc.x(qubit=qubits[0])
    qc.barrier()
    
qc.measure(qubit=qubits, cbit=bits)
qc.measure(qubit=color_qubit, cbit=color_bit)
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"], optimization_level=0)
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()}")

transpiled_qc = transpile(circuits=qc, basis_gates=["id", "x", "sx", "cx", "rz"], optimization_level=1)
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()}")

transpiled_qc = transpile(circuits=qc, basis_gates=["id", "x", "sx", "cx", "rz"], optimization_level=2)
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()}")

transpiled_qc = transpile(circuits=qc, basis_gates=["id", "x", "sx", "cx", "rz"], optimization_level=3)
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()}")

counts = execute(experiments=qc, backend=BACKEND, shots=SHOTS).result().get_counts()
plot_histogram(counts)

qubits = QuantumRegister(size=2, name="pixels")
color_qubit = QuantumRegister(size=1, name="color")
bits = ClassicalRegister(size=2, name="bits_pixels")
color_bit = ClassicalRegister(size=1)

qc = QuantumCircuit(qubits, color_qubit, bits, color_bit)

qc.h(qubit=qubits)
qc.barrier()

theta = np.pi/4

for i in range(4):
    if i == 0:
        qc.x(qubit=qubits)
    elif i == 1:
        qc.x(qubit=qubits[1])
    elif i == 2:
        qc.x(qubit=qubits[0])
    mcry = RYGate(theta=2*theta).control(num_ctrl_qubits=2)
    qc.append(mcry, qargs=[qubits[0], qubits[1], color_qubit])
    if i == 0:
        qc.x(qubit=qubits)
    elif i == 1:
        qc.x(qubit=qubits[1])
    elif i == 2:
        qc.x(qubit=qubits[0])
    qc.barrier()
     
qc.measure(qubit=qubits, cbit=bits)
qc.measure(qubit=color_qubit, cbit=color_bit)
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"], optimization_level=3)
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()}")

counts = execute(experiments=qc, backend=BACKEND, shots=SHOTS).result().get_counts()
plot_histogram(counts)

from fractions import Fraction

a = Fraction(2*255,90)
print(a)

(((127*3)/17))/90

astronaut_gray = rgb2gray(data.astronaut())
plt.imshow(astronaut_gray, cmap="gray")

resized_astronaut_gray = resize(astronaut_gray, (2,2))

resized_astronaut_gray

pixels_intensity_list = list(resized_astronaut_gray[0]) + list(resized_astronaut_gray[1])

qubits = QuantumRegister(size=2, name="pixels")
color_qubit = QuantumRegister(size=1, name="color")
bits = ClassicalRegister(size=2, name="bits_pixels")
color_bit = ClassicalRegister(size=1)

qc = QuantumCircuit(qubits, color_qubit, bits, color_bit)

qc.h(qubit=qubits)
qc.barrier()

for i, pixel in enumerate(pixels_intensity_list):
    theta = ((pixel*255*12)/17)/180
    if i == 0:
        qc.x(qubit=qubits)
    elif i == 1:
        qc.x(qubit=qubits[1])
    elif i == 2:
        qc.x(qubit=qubits[0])
    qc.cry(theta=theta, control_qubit=qubits[0], target_qubit=color_qubit)
    qc.cx(control_qubit=qubits[0], target_qubit=qubits[1])
    qc.cry(theta=-theta, control_qubit=qubits[1], target_qubit=color_qubit)
    qc.cx(control_qubit=qubits[0], target_qubit=qubits[1])
    qc.cry(theta=theta, control_qubit=qubits[1], target_qubit=color_qubit)
    if i == 0:
        qc.x(qubit=qubits)
    elif i == 1:
        qc.x(qubit=qubits[1])
    elif i == 2:
        qc.x(qubit=qubits[0])
    qc.barrier()
    
qc.measure(qubit=qubits, cbit=bits)
qc.measure(qubit=color_qubit, cbit=color_bit)
qc.draw(output="mpl")

counts = execute(experiments=qc, backend=BACKEND, shots=SHOTS).result().get_counts()
plot_histogram(counts)

image = np.array([[0.5,0.5],[0.5,0.5]])
qc = image_frqi.image_quantum_circuit(image=image, measurements=True)

qc.draw(output="mpl")

counts = execute(experiments=qc, backend=BACKEND, shots=SHOTS).result().get_counts()
plot_histogram(counts)

image = np.ones((4,4))
qc = image_frqi.image_quantum_circuit(image=image, measurements=True)

qc.draw(output="mpl")

counts = execute(experiments=qc, backend=BACKEND, shots=SHOTS).result().get_counts()
counts

qc = image_frqi.image_quantum_circuit(image=resized_astronaut_gray, measurements=True)

qc.draw(output="mpl")

counts = execute(experiments=qc, backend=BACKEND, shots=SHOTS).result().get_counts()
plot_histogram(counts)

astronaut = data.astronaut()

plt.imshow(astronaut)

resized_astronaut = resize(astronaut, (2,2))

plt.imshow(resized_astronaut)

qc = image_frqi.image_quantum_circuit(image=resized_astronaut, measurements=True)

qc.draw(output="mpl")

counts = execute(experiments=qc, backend=BACKEND, shots=SHOTS).result().get_counts()
plot_histogram(counts)

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

image = np.ones((4,4,4))
qc = image_frqi.image_quantum_circuit(image=image, measurements=True)

qc.draw(output="mpl")