%matplotlib inline
from qiskit import QuantumCircuit, QuantumRegister, transpile, ClassicalRegister
from qiskit.quantum_info import Statevector
from qiskit.visualization import array_to_latex, plot_histogram, plot_bloch_multivector
from qiskit.circuit.library import CHGate
from qiskit_aer import AerSimulator
from functools import partial
import numpy as np
from utils import get_image_path

show_unitary = partial(array_to_latex, max_size=100000)
show_state = lambda qc: Statevector(qc).draw('latex')

def get_unitary(qc, sim=AerSimulator()):
    qc_copy = qc.copy()
    qc_copy.save_unitary()
    unitary = sim.run(transpile(qc_copy, sim), shots=1000).result().get_unitary()
    return unitary

def measure(qc, sim=AerSimulator(), qubits=None):
    qc_copy = qc.copy()
    if(not qubits):
        qc_copy.measure_all()
    else:
        qc_copy.measure(qubits, list(range(len(qubits))))
    result = sim.run(transpile(qc_copy, sim), shots=1000).result().get_counts()
    return result

mem = QuantumCircuit(2, name="QROM")
mem.cz(0,1)
display(mem.draw('mpl'))
show_unitary(get_unitary(mem))

qc = QuantumCircuit(2)

qc.x(range(2))
qc.append(mem, [0,1])
qc.x(range(2))

display(qc.draw('mpl'))
display(show_unitary(get_unitary(qc)))
show_state(qc)

mem = QuantumCircuit(2, name="QROM")
mem.cp(np.pi/3,0,1)
display(mem.draw('mpl'))
show_unitary(get_unitary(mem))

qc = QuantumCircuit(2)

qc.x(range(2))
qc.append(mem, [0,1])
qc.x(range(2))

display(qc.draw('mpl'))
display(show_unitary(get_unitary(qc)))
show_state(qc)

plot_bloch_multivector(qc)

qc = QuantumCircuit(3)

qc.x(range(2))
qc.h(2)

qc.barrier()
qc.append(mem, [0,1])
qc.barrier()
qc.cx(2,1)
qc.x(range(2))
qc.h(2)



display(qc.draw('mpl'))
display(show_unitary(get_unitary(qc)))
display(show_state(qc))
display(plot_histogram(measure(qc)))


qc = QuantumCircuit(1)
#qc.h(0)
qc.p(np.pi/2, 0)
#qc.barrier()
#qc.cx(1,0)
qc.h(0)

display(qc.draw('mpl'))
display(show_unitary(get_unitary(qc)))
display(show_state(qc))
display(plot_histogram(measure(qc)))

mem = QuantumCircuit(2, name="QROM")
mem.cp(np.pi/3,0,1)
mem.x(range(2))
mem.cp(np.pi/7,0,1)
mem.x(range(2))
display(mem.draw('mpl'))
show_unitary(get_unitary(mem))

qc = QuantumCircuit(2)

#qc.x(range(2))
qc.append(mem, [0,1])
#qc.x(range(2))

display(qc.draw('mpl'))
display(show_unitary(get_unitary(qc)))
show_state(qc)

mem = QuantumCircuit(2, name="qrom")
mem.ch(0,1)
mem.x(0)
mem.cry(np.pi/3, 0, 1)
mem.x(0)
display(mem.draw('mpl', filename=get_image_path("qrom_1.png")))
display(show_unitary(get_unitary(mem)))
show_state(mem)

addr = QuantumRegister(1, name="addr")
out = QuantumRegister(1, name="out")
q = QuantumRegister(1,name="q")
result = ClassicalRegister(1, "m")

qc = QuantumCircuit(addr, out, q, result)

qc.x(addr)

qc.barrier(addr)
qc.append(mem, [addr, out])

qc.cx(out, q)

display(qc.draw('mpl'))
display(show_unitary(get_unitary(qc)))
display(show_state(qc))
display(plot_histogram(measure(qc, qubits=[2])))

addr = QuantumRegister(1, name="addr")
out = QuantumRegister(1, name="out")
q = QuantumRegister(1,name="q")
result = ClassicalRegister(1, "m")

qc = QuantumCircuit(addr, out, q)

qc.h(addr)

qc.barrier(addr)
qc.append(mem, [addr, out])

qc.cx(out, q)

display(qc.draw('mpl', filename=get_image_path("qrom_1_usage.png")))
display(show_unitary(get_unitary(qc)))
display(show_state(qc))
#display(plot_histogram(measure(qc, qubits=[2])))

CCH = CHGate().control(1)

addr = QuantumRegister(2, name="addr")
data = QuantumRegister(2, name="data")

mem = QuantumCircuit(addr, data, name="qrom")

mem.x(addr[0])
mem.ccx(addr[0], addr[1], data[0])
mem.x(addr[0])
mem.barrier()
mem.append(CCH, [addr[0], addr[1], data[0]])
mem.ccx(addr[0], addr[1], data[1])

display(mem.draw('mpl'))
display(show_unitary(get_unitary(mem)))
show_state(mem)

addr = QuantumRegister(2, name="addr")
data = QuantumRegister(2, name="data")
q = QuantumRegister(1, name="q")
out = ClassicalRegister(1, name="c")

qc = QuantumCircuit(addr, data, q, out)

qc.x(range(2))
qc.barrier(addr)
qc.append(mem, [*addr, *data])

qc.cx(data[1], q[0])
qc.cx(data[0], q[0])

qc.append(mem, [*addr, *data])
qc.barrier(addr)
qc.x(range(2))



display(qc.draw('mpl'))
#display(show_unitary(get_unitary(qc)))
display(show_state(qc))
display(plot_histogram(measure(qc, qubits=[4])))

#memory with teleportation

addr = QuantumRegister(1, name="addr")
data = QuantumRegister(1, name="data")

mem = QuantumCircuit(addr, data, name="qrom")
mem.cry(np.pi/5, 0, 1)

display(mem.draw('mpl'))
display(show_unitary(get_unitary(mem)))
show_state(mem)

addr = QuantumRegister(1, name="addr")
data = QuantumRegister(2, name="data")
q = QuantumRegister(1, name="q")
out = ClassicalRegister(1, name="c")

qc = QuantumCircuit(addr, data, q, out)

qc.x(addr)
qc.barrier(addr)
qc.append(mem, [addr, data[0]])

qc.cx(data[0], q)

qc.barrier(addr)
qc.append(mem, [addr, data[0]])
qc.x(addr)



display(qc.draw('mpl'))
#display(show_unitary(get_unitary(qc)))
display(show_state(qc))
display(plot_histogram(measure(qc, qubits=[3])))

addr = QuantumRegister(1, name="addr")
data = QuantumRegister(2, name="data")
q = QuantumRegister(1, name="q")
out = ClassicalRegister(1, name="c")

qc = QuantumCircuit(addr, data, q, out)

qc.x(addr)
qc.barrier(addr)
qc.append(mem, [addr, data[0]])

qc.h(data[1])
qc.cx(data[1], q)

qc.cx(data[0], data[1])
qc.h(data[0])

qc.cx(data[1], q)
qc.cz(data[0], q)


display(qc.draw('mpl'))
#display(show_unitary(get_unitary(qc)))
display(show_state(qc))
display(plot_histogram(measure(qc, qubits=[3])))
display(plot_histogram(measure(qc, qubits=[1])))


(np.cos(np.pi/5/2))**2, (np.sin(np.pi/5/2))**2

qc = QuantumCircuit(3,3)
qc.ry(np.pi/5, 0)
qc.barrier()

qc.h(1)
qc.cx(1,2)

qc.barrier()

qc.cx(0,1)
qc.h(0)

qc.cx(1,2)
qc.cz(0,2)
qc.draw('mpl')

display(qc.draw('mpl'))
#display(show_unitary(get_unitary(qc)))
display(show_state(qc))
display(plot_histogram(measure(qc, qubits=[2])))

data = QuantumRegister(1, name="data")
teleportation_provider = QuantumRegister(1, name="t")
q = QuantumRegister(1, name="q")
meas = ClassicalRegister(1, name="out")

qram = QuantumCircuit(data, teleportation_provider, q,meas)

qram.h(teleportation_provider)
qram.cx(teleportation_provider, data)

qram.barrier()

qram.ry(np.pi/7, q)

qram.barrier()

qram.cx(q, teleportation_provider)
qram.h(q)
qram.cx(teleportation_provider, data)
qram.cz(q, data)

display(qram.draw('mpl'))
#display(show_unitary(get_unitary(qram)))
display(show_state(qram))
display(plot_histogram(measure(qram, qubits=[0])))
display(plot_histogram(measure(qram, qubits=[2])))

np.cos(np.pi/7/2)**2, np.sin(np.pi/7/2)**2



data = QuantumRegister(2, name="data")
teleportation_provider = QuantumRegister(2, name="t")
q = QuantumRegister(2, name="q")
meas = ClassicalRegister(1, name="out")

qram = QuantumCircuit(data, teleportation_provider, q,meas)

qram.h(teleportation_provider)
qram.cx(teleportation_provider, data)

qram.barrier()

qram.ry(np.pi/3, q[0])
qram.ry(np.pi/7, q[1])

qram.barrier()

qram.cx(q, teleportation_provider)
qram.h(q)
qram.cx(teleportation_provider, data)
qram.cz(q, data)

display(qram.draw('mpl', filename=get_image_path('qram.png')))
#display(show_unitary(get_unitary(qram)))
display(show_state(qram))
display(plot_histogram(measure(qram, qubits=[0])))
display(plot_histogram(measure(qram, qubits=[4])))
display(plot_histogram(measure(qram, qubits=[1])))
display(plot_histogram(measure(qram, qubits=[5])))

np.cos(np.pi/3/2)**2, np.sin(np.pi/3/2)**2

data = QuantumRegister(2, name="data")
teleportation_provider = QuantumRegister(2, name="t")
q = QuantumRegister(2, name="q")
meas = ClassicalRegister(1, name="out")

qram = QuantumCircuit(data, teleportation_provider, q,meas)

qram.h(teleportation_provider)
qram.cx(teleportation_provider, data)

qram.barrier()

qram.h(q[0])
qram.h(q[1])

qram.barrier()

qram.cx(q, teleportation_provider)
qram.h(q)
qram.cx(teleportation_provider, data)
qram.cz(q, data)

qram.barrier(label='second')
qram.cz(q, data)
qram.cx(teleportation_provider, data)
qram.h(q)
qram.cx(q, teleportation_provider)

qram.barrier()
qram.h(q[0])
qram.barrier()
qram.cx(q, teleportation_provider)
qram.h(q)
qram.cx(teleportation_provider, data)
qram.cz(q, data)


display(qram.draw('mpl'))
#display(show_unitary(get_unitary(qram)))
display(show_state(qram))
display(plot_histogram(measure(qram, qubits=[0])))
display(plot_histogram(measure(qram, qubits=[4])))
display(plot_histogram(measure(qram, qubits=[1])))
display(plot_histogram(measure(qram, qubits=[5])))

data = QuantumRegister(2, name="data")
teleportation_provider = QuantumRegister(1, name="t")
q = QuantumRegister(2, name="q")
meas = ClassicalRegister(1, name="out")

qram = QuantumCircuit(data, teleportation_provider, q,meas)

qram.h(teleportation_provider)
qram.cx(teleportation_provider, data)

qram.barrier()

qram.ry(np.pi/3, q[0])
qram.ry(np.pi/7, q[1])

qram.barrier()

qram.cx(q, teleportation_provider)
qram.h(q)
qram.cx(teleportation_provider, data)
qram.cz(q, data)

display(qram.draw('mpl', filename=get_image_path('qram.png')))
#display(show_unitary(get_unitary(qram)))
display(show_state(qram))
display(plot_histogram(measure(qram, qubits=[0])))
display(plot_histogram(measure(qram, qubits=[4])))
display(plot_histogram(measure(qram, qubits=[1])))
display(plot_histogram(measure(qram, qubits=[5])))