import numpy as np
from numpy import pi
import matplotlib.pyplot as plt
import matplotlib.image  as mpimg

filename = './schrodin_yang.png'
im = mpimg.imread(filename)
fig, ax = plt.subplots()
ax.imshow(im)

from skimage.transform import resize

n_pixels = 2**5
im = resize(im, (n_pixels, n_pixels))
fig, ax = plt.subplots()
ax.imshow(im)

data = im[:,:,0].ravel()

from qiskit_aqua.components.initial_states import Custom

n_qubits   = np.int_(np.log2(len(data)))
init_state = Custom(n_qubits, state_vector=data)

circ = init_state.construct_circuit('circuit')
qr   = circ.qregs

from copy import deepcopy

circ_init = deepcopy(circ)

circ.h(qr[0][0])

from qiskit import BasicAer, execute

simulator = BasicAer.get_backend('statevector_simulator')
result = execute(circ, simulator).result()
final_state_vector = result.get_statevector(circ)

edge_even = np.real(final_state_vector)
n_rows    = int(np.sqrt(len(edge_even)))
edge_even = edge_even.reshape(n_rows, -1)

edge_even[:,::2] = 0

fig, ax = plt.subplots(1,2)
ax[0].imshow(edge_even)
ax[1].imshow(im)

def qft(c, q, n):
    """n-qubit QFT on q in c."""
    for j in range(n):
        for k in range(j):
            c.cu1(pi/float(2**(j-k)), q[j], q[k])
        c.h(q[j])
        
def iqft(c, q, n):
    """n-qubit IQFT on q in c."""
    for j in range(n):
        for k in range(j):
            c.cu1(-pi/float(2**(j-k)), q[int(n-j-1)], q[int(n-k-1)])
        c.h(q[int(n-j-1)])
        
def shiftBases(c, q, n, p):
    """Shift the register q by p positions"""
    iqft(c,q,n)
    for k in range(n):
        c.u1(-p*2*pi/float(2**(n-k)), q[int(n-k-1)])
    qft(c,q,n)

from qiskit import QuantumRegister, QuantumCircuit

n = 2
q = QuantumRegister(n,'q')
c = QuantumCircuit(q)
iqft(c,q,n)
shiftBases(c,q,n,1)
qft(c,q,n)
c.draw(output='mpl')

circ = deepcopy(circ_init)
qr   = circ.qregs

shiftBases(circ,qr[0],n_qubits,1)
circ.h(qr[0][0])
shiftBases(circ,qr[0],n_qubits,-1)

simulator = BasicAer.get_backend('statevector_simulator')
result = execute(circ, simulator).result()
final_state_vector = result.get_statevector(circ)

edge_odd = np.real(final_state_vector)
n_rows = np.int_(np.sqrt(len(edge_odd)))
edge_odd = edge_odd.reshape(n_rows, -1)

edge_odd[:,1::2] = 0

fig, ax = plt.subplots(1,2)
ax[0].imshow(edge_odd)
ax[1].imshow(im)

edge = edge_even + edge_odd

fig, ax = plt.subplots(1,2)
ax[0].imshow(edge)
ax[1].imshow(im)