import matplotlib.pyplot as plt
import numpy as np
from qiskit import QuantumCircuit
from qiskit import execute, Aer


def margolus(circ, t, c0, c1):
    circ.ry(np.pi / 4, t)
    circ.cx(c0, t)
    circ.ry(np.pi / 4, t)
    circ.cx(c1, t)
    circ.ry(-np.pi / 4, t)
    circ.cx(c0, t)
    circ.ry(-np.pi / 4, t)


def rccx(circ, t, c0, c1):
    circ.h(t)
    circ.t(t)
    circ.cx(c0, t)
    circ.tdg(t)
    circ.cx(c1, t)
    circ.t(t)
    circ.cx(c0, t)
    circ.tdg(t)
    circ.h(t)


def rcccx(circ, t, c0, c1, c2):
    circ.h(t)
    circ.t(t)
    circ.cx(c0, t)
    circ.tdg(t)
    circ.h(t)
    circ.cx(c1, t)
    circ.t(t)
    circ.cx(c2, t)
    circ.tdg(t)
    circ.cx(c1, t)
    circ.t(t)
    circ.cx(c2, t)
    circ.tdg(t)
    circ.h(t)
    circ.t(t)
    circ.cx(c0, t)
    circ.tdg(t)
    circ.h(t)


def ccry(circ, angle, t, c0, c1):
    circ.cu3(angle / 2, 0, 0, c1, t)
    circ.cx(c1, c0)
    circ.cu3(-angle / 2, 0, 0, c0, t)
    circ.cx(c1, c0)
    circ.cu3(angle / 2, 0, 0, c0, t)


def mary(circ, angle, t, c0, c1):
    circ.ry(angle / 4, t)
    circ.cx(c0, t)
    circ.ry(-angle / 4, t)
    circ.cx(c1, t)
    circ.ry(angle / 4, t)
    circ.cx(c0, t)
    circ.ry(-angle / 4, t)
    circ.cx(c1, t)


def cccry(circ, angle, t, a, c0, c1, c2):
    margolus(circ, a, c1, c2)
    mary(circ, angle, t, a, c0)
    margolus(circ, a, c1, c2)


def mary_4(circ, angle, t, c0, c1, c2):
    circ.h(t)
    circ.t(t)
    circ.cx(c0, t)
    circ.tdg(t)
    circ.h(t)
    circ.cx(c1, t)
    circ.rz(angle / 4, t)
    circ.cx(c2, t)
    circ.rz(-angle / 4, t)
    circ.cx(c1, t)
    circ.rz(angle / 4, t)
    circ.cx(c2, t)
    circ.rz(-angle / 4, t)
    circ.h(t)
    circ.t(t)
    circ.cx(c0, t)
    circ.tdg(t)
    circ.h(t)


def mary_8(circ, angle, t, c0, c1, c2, c3, c4, c5, c6):
    circ.h(t)
    circ.t(t)
    rccx(circ, t, c0, c1)
    circ.tdg(t)
    circ.h(t)
    rccx(circ, t, c2, c3)
    circ.rz(angle / 4, t)
    rcccx(circ, t, c4, c5, c6)
    circ.rz(-angle / 4, t)
    rccx(circ, t, c2, c3)
    circ.rz(angle / 4, t)
    rcccx(circ, t, c4, c5, c6)
    circ.rz(-angle / 4, t)
    circ.h(t)
    circ.t(t)
    rccx(circ, t, c0, c1)
    circ.tdg(t)
    circ.h(t)


def c10ry(circ, angle, bin, target, anc, controls):
    clist = []

    for i in bin:
        clist.append(int(i))

    for i in range(len(clist)):
        if clist[i] == 0:
            circ.x(controls[-i - 1])

    margolus(circ, anc, controls[0], controls[1])
    circ.x(controls[0])
    circ.x(controls[1])
    margolus(circ, controls[1], controls[2], controls[3])
    circ.x(controls[2])
    circ.x(controls[3])
    margolus(circ, controls[3], controls[4], controls[5])
    circ.x(controls[4])
    circ.x(controls[5])

    margolus(circ, controls[5], controls[8], controls[9])
    margolus(circ, controls[4], controls[6], controls[7])
    margolus(circ, controls[2], controls[4], controls[5])
    margolus(circ, controls[0], controls[2], controls[3])

    mary_4(circ, angle, target, anc, controls[0], controls[1])

    margolus(circ, controls[0], controls[2], controls[3])
    margolus(circ, controls[2], controls[4], controls[5])
    margolus(circ, controls[4], controls[6], controls[7])
    margolus(circ, controls[5], controls[8], controls[9])

    circ.x(controls[5])
    circ.x(controls[4])
    margolus(circ, controls[3], controls[4], controls[5])
    circ.x(controls[3])
    circ.x(controls[2])
    margolus(circ, controls[1], controls[2], controls[3])
    circ.x(controls[1])
    circ.x(controls[0])
    margolus(circ, anc, controls[0], controls[1])

    for i in range(len(clist)):
        if clist[i] == 0:
            circ.x(controls[-i - 1])


def c10mary(circ, angle, bin, target, anc, controls):
    clist = []

    for i in bin:
        clist.append(int(i))

    for i in range(len(clist)):
        if clist[i] == 0:
            circ.x(controls[-i - 1])

    rccx(circ, anc, controls[4], controls[5])
    circ.x(controls[4])
    circ.x(controls[5])
    rccx(circ, controls[4], controls[6], controls[7])
    rccx(circ, controls[5], controls[8], controls[9])

    mary_8(circ, angle, target, anc, controls[0], controls[1], controls[2], controls[3], controls[4], controls[5])

    rccx(circ, controls[5], controls[8], controls[9])
    rccx(circ, controls[4], controls[6], controls[7])
    circ.x(controls[5])
    circ.x(controls[4])
    rccx(circ, anc, controls[4], controls[5])

    for i in range(len(clist)):
        if clist[i] == 0:
            circ.x(controls[-i - 1])


if __name__ == '__main__':
    from keras.datasets import mnist

    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    img_num = 0

    # show original image
    plt.imshow(x_train[img_num], cmap='gray')
    # plt.savefig('mnistimg'+str(img_num)+'.png')
    plt.show()

    # 2-dimentional data convert to 1-dimentional array
    x_train = x_train.reshape(60000, 784)
    # change type
    x_train = x_train.astype('float64')
    # Normalization(0~pi/2)
    x_train /= 255.0
    x_train = np.arcsin(x_train)

    backends = Aer.backends()
    # print("Aer backends:",backends)

    qubit = 12
    qc = QuantumCircuit(qubit, qubit)

    # apply hadamard gates
    qc.h(range(2, qubit))

    # apply c10Ry gates (representing color data)
    for i in range(len(x_train[img_num])):
        if x_train[img_num][i] != 0:
            c10mary(qc, 2 * x_train[img_num][i], format(i, '010b'), 0, 1, [i for i in range(2, 12)])

    qc.measure(range(qubit), range(qubit))

    backend_sim = Aer.get_backend('qasm_simulator')
    # print(qc.depth())
    numOfShots = 1024000
    result = execute(qc, backend_sim, shots=numOfShots).result()
    # circuit_drawer(qc).show()
    # plot_histogram(result.get_counts(qc))

    print(result.get_counts(qc))

    # generated image
    genimg = np.array([])

    #### decode
    for i in range(len(x_train[img_num])):
        try:
            genimg = np.append(genimg, [np.sqrt(result.get_counts(qc)[format(i, '010b') + '01'] / numOfShots)])
        except KeyError:
            genimg = np.append(genimg, [0.0])

    # inverse nomalization
    genimg *= 32.0 * 255.0
    x_train = np.sin(x_train)
    x_train *= 255.0

    # convert type
    genimg = genimg.astype('int')

    # back to 2-dimentional data
    genimg = genimg.reshape((28, 28))

    plt.imshow(genimg, cmap='gray', vmin=0, vmax=255)
    plt.savefig('gen_' + str(img_num) + '.png')
    plt.show()
