from IPython.display import Image, display
Image("ryoko.png", width="70")

Image('asteroids_example.png')

Image('asteroids_beam_example.png')

Image('false_asteroids_example.png')

Image('asteroids_example.png')

problem_set = \
    [[['0', '2'], ['1', '0'], ['1', '2'], ['1', '3'], ['2', '0'], ['3', '3']],
    [['0', '0'], ['0', '1'], ['1', '2'], ['2', '2'], ['3', '0'], ['3', '3']],
    [['0', '0'], ['1', '1'], ['1', '3'], ['2', '0'], ['3', '2'], ['3', '3']],
    [['0', '0'], ['0', '1'], ['1', '1'], ['1', '3'], ['3', '2'], ['3', '3']],
    [['0', '2'], ['1', '0'], ['1', '3'], ['2', '0'], ['3', '2'], ['3', '3']],
    [['1', '1'], ['1', '2'], ['2', '0'], ['2', '1'], ['3', '1'], ['3', '3']],
    [['0', '2'], ['0', '3'], ['1', '2'], ['2', '0'], ['2', '1'], ['3', '3']],
    [['0', '0'], ['0', '3'], ['1', '2'], ['2', '2'], ['2', '3'], ['3', '0']],
    [['0', '3'], ['1', '1'], ['1', '2'], ['2', '0'], ['2', '1'], ['3', '3']],
    [['0', '0'], ['0', '1'], ['1', '3'], ['2', '1'], ['2', '3'], ['3', '0']],
    [['0', '1'], ['0', '3'], ['1', '2'], ['1', '3'], ['2', '0'], ['3', '2']],
    [['0', '0'], ['1', '3'], ['2', '0'], ['2', '1'], ['2', '3'], ['3', '1']],
    [['0', '1'], ['0', '2'], ['1', '0'], ['1', '2'], ['2', '2'], ['2', '3']],
    [['0', '3'], ['1', '0'], ['1', '3'], ['2', '1'], ['2', '2'], ['3', '0']],
    [['0', '2'], ['0', '3'], ['1', '2'], ['2', '3'], ['3', '0'], ['3', '1']],
    [['0', '1'], ['1', '0'], ['1', '2'], ['2', '2'], ['3', '0'], ['3', '1']]]

import numpy as np

# import qiskit libraries

from qiskit import IBMQ, Aer,QuantumRegister, ClassicalRegister, QuantumCircuit, execute
from qiskit.tools.jupyter import *
from qiskit.circuit.library import OR as or_gate

provider = IBMQ.load_account()

#import functions to plot
from qiskit.visualization import plot_histogram



def week3_ans_func(problem_set):
    ##### build your quantum circuit here
    ##### In addition, please make it a function that can solve the problem even with different inputs (problem_set). We do validation with different inputs. 

    def qRAM(qc,add,data, auxiliary, problem_set):

        problem_set_mapping_dict = {'0' : {'0': 0, '1': 1, '2': 2, '3':3},
                                    '1' : {'0': 4, '1': 5, '2': 6, '3':7},
                                    '2' : {'0': 8, '1': 9, '2': 10, '3':11},
                                    '3' : {'0': 12, '1': 13, '2': 14, '3':15}}
        for ind in range(len(data)):
            binary = f'{ind:04b}'
            if(binary[0] == '0'):
                qc.x(add[0])
            if(binary[1] == '0'):
                qc.x(add[1])
            if(binary[2] == '0'):
                qc.x(add[2])
            if(binary[3] == '0'):
                qc.x(add[3])

            for i, edg in enumerate(problem_set[ind]):

                qc.ch(add[0],auxiliary[0]) 
                qc.cz(add[1],auxiliary[0]) 
                qc.ch(add[0], auxiliary[0]) 
                #qc.ccx(add[0], add[1], auxiliary[0])

                qc.ch(add[2],auxiliary[1])
                qc.cz(add[3],auxiliary[1])
                qc.ch(add[2], auxiliary[1])
                #qc.ccx(add[2], add[3], auxiliary[1])

                qc.ch(auxiliary[0],auxiliary[2])
                qc.cz(auxiliary[1],auxiliary[2])
                qc.ch(auxiliary[0],auxiliary[2])
                #qc.ccx(auxiliary[0], auxiliary[1], auxiliary[2])

                qc.cx(auxiliary[2], data[problem_set_mapping_dict[edg[0]][edg[1]]])

                qc.ch(auxiliary[0],auxiliary[2])
                qc.cz(auxiliary[1],auxiliary[2])
                qc.ch(auxiliary[0],auxiliary[2])

                qc.ch(add[2],auxiliary[1])
                qc.cz(add[3],auxiliary[1])
                qc.ch(add[2], auxiliary[1])

                qc.ch(add[0],auxiliary[0])
                qc.cz(add[1],auxiliary[0])
                qc.ch(add[0], auxiliary[0])


            if(binary[0] == '0'):
                qc.x(add[0])
            if(binary[1] == '0'):
                qc.x(add[1])
            if(binary[2] == '0'):
                qc.x(add[2])
            if(binary[3] == '0'):
                qc.x(add[3])
            qc.barrier()    
            
    def diffuser(nqubits):

        qc_diff = QuantumCircuit(nqubits)

        for qubit in range(nqubits):
            qc_diff.h(qubit)

        for qubit in range(nqubits):
            qc_diff.x(qubit)

        qc_diff.h(nqubits-1)
        qc_diff.mct(list(range(nqubits-1)), nqubits-1)  # multi-controlled-toffoli
        qc_diff.h(nqubits-1)

        for qubit in range(nqubits):
            qc_diff.x(qubit)

        for qubit in range(nqubits):
            qc_diff.h(qubit)

        U_f0 = qc_diff.to_gate()
        U_f0.name = "V"

        return U_f0   
    
    def add_row_counters(qc, data, auxiliary):

        row_mat = [[0, 1, 2, 3],
                   [4, 5, 6, 7],
                   [8, 9, 10, 11],
                   [12, 13, 14, 15]]

        for i, curr_row in enumerate(row_mat):

            flag_list = [0] * len(data)
            for curr_flag_index in range(len(data)):
                if curr_flag_index in curr_row:
                    flag_list[curr_flag_index] = 1

            qc.append(or_gate(len(data), flag_list, mcx_mode='noancilla'), [*data, auxiliary[i]])


    def add_column_counters(qc, data, auxiliary):

        col_mat = [[0, 4, 8, 12],
                   [1, 5, 9, 13],
                   [2, 6, 10, 14],
                   [3, 7, 11, 15]]


        for i, curr_col in enumerate(col_mat):

            flag_list = [0] * len(data)
            for curr_flag_index in range(len(data)):
                if curr_flag_index in curr_col:
                    flag_list[curr_flag_index] = 1

            qc.append(or_gate(len(data), flag_list, mcx_mode='noancilla'), [*data, auxiliary[i]])        

    def add_lone_counters(qc, data, auxiliary, count_out):

        lone_mat = [[0, 1, 2, 3, 4, 8, 12],
                    [1, 0, 2, 3, 5, 9, 13],
                    [2, 0, 1, 3, 6, 10, 14],
                    [3, 0, 1, 2, 7, 11, 15],
                    [4, 5, 6, 7, 0, 8, 12],
                    [5, 4, 6, 7, 1, 9, 13],
                    [6, 4, 5, 7, 2, 10, 14],
                    [7, 4, 5, 6, 3, 11, 15],
                    [8, 9, 10, 11, 0, 4, 12],
                    [9, 8, 10, 11, 1, 5, 13],
                    [10, 8, 9, 11, 2, 6, 14],
                    [11, 8, 9, 10, 3, 7, 15],
                    [12, 13, 14, 15, 0, 4, 8],
                    [13, 12, 14, 15, 1, 5, 9],
                    [14, 12, 13, 15, 2, 6, 10],
                    [15, 12, 13, 14, 3, 7, 11]]

        qc.barrier()

        #curr_lone_row = [0, 1, 2, 3, 4, 8, 12]
        for curr_lone_row in lone_mat:

            n = len(curr_lone_row)

            for j in range(1,n):
                qc.x(data[curr_lone_row[j]])

            qc.ch(data[curr_lone_row[0]], auxiliary[0])
            qc.cz(data[curr_lone_row[1]], auxiliary[0])
            qc.ch(data[curr_lone_row[0]],auxiliary[0]) 
            #qc.ccx(data[curr_lone_row[0]], data[curr_lone_row[1]], auxiliary[0])

            qc.ch(data[curr_lone_row[2]],auxiliary[1])
            qc.cz(data[curr_lone_row[3]],auxiliary[1])
            qc.ch(data[curr_lone_row[2]],auxiliary[1]) 
            #qc.ccx(data[curr_lone_row[2]], data[curr_lone_row[3]], auxiliary[1])

            qc.ch(auxiliary[0],auxiliary[2])
            qc.cz(auxiliary[1],auxiliary[2])
            qc.ch(auxiliary[0],auxiliary[2])
            #qc.ccx(auxiliary[0], auxiliary[1], auxiliary[2])

            qc.ch(data[curr_lone_row[0]],auxiliary[0])
            qc.cz(data[curr_lone_row[1]],auxiliary[0])
            qc.ch(data[curr_lone_row[0]],auxiliary[0]) 
            #qc.ccx(data[curr_lone_row[0]], data[curr_lone_row[1]], auxiliary[0])

            qc.ch(data[curr_lone_row[4]],auxiliary[0])
            qc.cz(data[curr_lone_row[5]],auxiliary[0])
            qc.ch(data[curr_lone_row[4]],auxiliary[0]) 
            #qc.ccx(data[curr_lone_row[4]], data[curr_lone_row[5]], auxiliary[0])

            qc.ch(auxiliary[0],auxiliary[3])
            qc.cz(auxiliary[2],auxiliary[3]) 
            qc.ch(auxiliary[0],auxiliary[3])
            #qc.ccx(auxiliary[0], auxiliary[2], auxiliary[3])

            #------------------------------------------------------------------------------------------------------------------------
            #qc.ch(data[curr_lone_row[6]], count_out[2])
            #qc.cz(auxiliary[3],count_out[2])
            #qc.ch(data[curr_lone_row[2]],count_out[2]) 
            qc.ccx(data[curr_lone_row[6]], auxiliary[3], count_out[2]) #this CCX is important to avoid
            #unnecessary addition of relative phase
            #------------------------------------------------------------------------------------------------------------------------

            qc.ch(auxiliary[0],auxiliary[3])
            qc.cz(auxiliary[2],auxiliary[3]) 
            qc.ch(auxiliary[0],auxiliary[3])

            qc.ch(data[curr_lone_row[4]],auxiliary[0])
            qc.cz(data[curr_lone_row[5]],auxiliary[0])
            qc.ch(data[curr_lone_row[4]],auxiliary[0]) 

            qc.ch(data[curr_lone_row[0]],auxiliary[0])
            qc.cz(data[curr_lone_row[1]],auxiliary[0])
            qc.ch(data[curr_lone_row[0]],auxiliary[0]) 

            qc.ch(auxiliary[0],auxiliary[2])
            qc.cz(auxiliary[1],auxiliary[2])
            qc.ch(auxiliary[0],auxiliary[2])

            qc.ch(data[curr_lone_row[2]],auxiliary[1])
            qc.cz(data[curr_lone_row[3]],auxiliary[1])
            qc.ch(data[curr_lone_row[2]],auxiliary[1]) 

            qc.ch(data[curr_lone_row[0]], auxiliary[0])
            qc.cz(data[curr_lone_row[1]], auxiliary[0])
            qc.ch(data[curr_lone_row[0]],auxiliary[0]) 
            for j in range(1,n):
                qc.x(data[curr_lone_row[j]])

            qc.barrier
        
    def u3a(data, auxiliary,oracle, count_out):

        qc_1 = QuantumCircuit(data, auxiliary,oracle, count_out)

        #Oracle -----------------------------------------------
        #Adding Row counters to check if each row has at least 1 asteroid and then flipping a count_out bit if all rows has at least one asteroid
        add_row_counters(qc_1, data, auxiliary) #compute
        qc_1.mct([auxiliary[0], auxiliary[1], auxiliary[2], auxiliary[3]], count_out[0], mode = 'noancilla')
        add_row_counters(qc_1, data, auxiliary) #uncompute

        #Adding column counters to check if each column has at least 1 asteroid and then flipping a count_out bit if all columns has at least one asteroid
        add_column_counters(qc_1, data, auxiliary) #compute
        qc_1.mct([auxiliary[0], auxiliary[1], auxiliary[2], auxiliary[3]], count_out[1], mode = 'noancilla')
        add_column_counters(qc_1, data, auxiliary) #uncompute

        #qc_1.barrier()
        #lone asteroid counter circuit to add
        add_lone_counters(qc_1, data, auxiliary, count_out)

        return qc_1
            
    def u3a_dagger(data, auxiliary, oracle, count_out):

        qc_1 = QuantumCircuit(data, auxiliary,oracle, count_out)

        #Oracle -----------------------------------------------
        #Adding Row counters to check if each row has at least 1 asteroid and then flipping a count_out bit if all rows has at least one asteroid
        add_row_counters(qc_1, data, auxiliary) #compute
        qc_1.mct([auxiliary[0], auxiliary[1], auxiliary[2], auxiliary[3]], count_out[0], mode = 'noancilla')
        add_row_counters(qc_1, data, auxiliary) #uncompute

        #Adding column counters to check if each column has at least 1 asteroid and then flipping a count_out bit if all columns has at least one asteroid
        add_column_counters(qc_1, data, auxiliary) #compute
        qc_1.mct([auxiliary[0], auxiliary[1], auxiliary[2], auxiliary[3]], count_out[1], mode = 'noancilla')
        add_column_counters(qc_1, data, auxiliary) #uncompute

        #qc_1.barrier()
        #lone asteroid counter circuit to add
        add_lone_counters(qc_1, data, auxiliary, count_out)

        qc_2 = qc_1.inverse()

        return qc_2
            
    add = QuantumRegister(4, name = 'address')
    data = QuantumRegister(16, name = 'data')
    auxiliary = QuantumRegister(4, name = 'auxiliary')
    count_out = QuantumRegister(3, name = 'counter_outputs')
    oracle = QuantumRegister(1, name = 'oracle')
    cbits = ClassicalRegister(4, name = 'solution')


    qc = QuantumCircuit(add, data, auxiliary,oracle, count_out, cbits)

    qc.x(oracle)
    qc.h(oracle)

    qc.h(add)
    qc.barrier()

    #Add qRAM
    qRAM(qc,add,data, auxiliary, problem_set) #compute
    qc.barrier()
    
    # counters
    qc.extend(u3a(data, auxiliary,oracle, count_out))
    qc.barrier()

    # Phase flipping oracle
    #qc.cx(auxiliary[6], oracle)
    qc.mct([count_out[0], count_out[1], count_out[2]], oracle, mode = 'noancilla')
    # -------------------------------------------------------
    
    # counters - dagger
    qc.extend(u3a_dagger(data, auxiliary,oracle, count_out))
    qc.barrier()

    #qRAM
    qc.barrier()
    qRAM(qc,add,data, auxiliary, problem_set)  #uncompute
    qc.barrier()

    #diffuser
    qc.append(diffuser(4), add)
    qc.barrier()

    #measure
    qc.measure(add[0], cbits[0])
    qc.measure(add[1], cbits[1]) 
    qc.measure(add[2], cbits[2]) 
    qc.measure(add[3], cbits[3]) 

    
    return qc

# Submission code
from qc_grader import grade_ex3, prepare_ex3, submit_ex3

# Execute your circuit with following prepare_ex3() function.
# The prepare_ex3() function works like the execute() function with only QuantumCircuit as an argument.
job = prepare_ex3(week3_ans_func)

result = job.result()
counts = result.get_counts()
original_problem_set_counts = counts[0]

original_problem_set_counts
# The bit string with the highest number of observations is treated as the solution.

# Check your answer by executing following code.
# The quantum cost of the QuantumCircuit is obtained as the score. The lower the cost, the better.
grade_ex3(job)

# Submit your results by executing following code. You can submit as many times as you like during the period. 
submit_ex3(job)

