from IPython.display import Image, display
Image('ryoko_pillars.png')

Image('lights_out_rule.png')

Image('lights_out_prob.png')

import numpy as np

# import qiskit libraries

from qiskit import IBMQ, Aer,QuantumRegister, ClassicalRegister, QuantumCircuit, execute
from qiskit.providers.ibmq import least_busy
from qiskit.quantum_info import Operator
from qiskit.tools.jupyter import *

#provider = IBMQ.load_account()

#import functions to plot
from qiskit.visualization import plot_histogram

# The starting pattern is represented by this list of numbers.
# Please use it as an input for your solution.
lights = [0, 1, 1, 1, 0, 0, 1, 1, 1]

def week2a_ans_func(lights):
    ##### build your quantum circuit here
    def lights_oracle(qc, sol, board, orc_out, anc, c_out):

        flip_matrix = [[1,1,0,1,0,0,0,0,0],
                       [1,1,1,0,1,0,0,0,0],
                       [0,1,1,0,0,1,0,0,0],
                       [1,0,0,1,1,0,1,0,0],
                       [0,1,0,1,1,1,0,1,0],
                       [0,0,1,0,1,1,0,0,1],
                       [0,0,0,1,0,0,1,1,0],
                       [0,0,0,0,1,0,1,1,1],
                       [0,0,0,0,0,1,0,1,1]]

        for sol_bit in range(len(flip_matrix)):
            for board_bit in range(len(flip_matrix[sol_bit])):
                if flip_matrix[sol_bit][board_bit] == 1:

                    qc.cx(sol[sol_bit], board[board_bit])

        qc.barrier()
        for i in range(len(flip_matrix)):
            qc.x(board[i])

        qc.barrier()
        qc.mct([board[0], board[1],board[2],board[3],board[4],board[5],board[6],board[7],board[8]],
               orc_out, 
               [anc[0], anc[1],anc[2],anc[3],anc[4],anc[5],anc[6]], 
               mode='v-chain') 

        qc.barrier()

        for i in range(len(flip_matrix)):
            qc.x(board[i])

        for sol_bit in range(len(flip_matrix)):
            for board_bit in range(len(flip_matrix[sol_bit])):
                if flip_matrix[sol_bit][board_bit] == 1:

                    qc.cx(sol[sol_bit], board[board_bit])

    def diffuser(nqubits):
        qc = QuantumCircuit(nqubits)

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

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

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

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

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

        U_f0 = qc.to_gate()
        U_f0.name = "V"
        return U_f0                

#def groverLights(lights):

    n = len(lights)
    
    r = int(np.round(np.pi/(4*np.arcsin(np.sqrt(1/3**n)))-1/2))
    
    print(f'{3*n-1} qubits, {r} rounds')
    
    board = QuantumRegister(n, 'board')
    sol = QuantumRegister(n, 'solution')
    anc = QuantumRegister(n-2, 'anc')
    orc_out = QuantumRegister(1, 'oracle_out') 
    c_out = ClassicalRegister(n, 'classical_bits')

    qc = QuantumCircuit(sol, board, orc_out, anc, c_out)

    for l in range(n):
        if lights[l] == 1:
            qc.x(board[l])

    qc.x(orc_out)     
    qc.h(orc_out)
    qc.barrier() 
    
    qc.h(sol)
    qc.barrier()

    r = 1
    
    for _ in range(r):
        lights_oracle(qc, sol, board, orc_out, anc, c_out)
        qc.barrier()
        qc.append(diffuser(n), sol)
        qc.barrier()   
        
    qc.measure(sol, c_out)
    qc = qc.reverse_bits()
    #####  In addition, please make it a function that can solve the problem even with different inputs (lights). We do validation with different inputs.
    
    return qc

# Submission code
from qc_grader import prepare_ex2a, grade_ex2a, submit_ex2a

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

result = job.result()
count = result.get_counts()
original_problem_set_counts = count[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 quantum cost is related to rank only in the third week.
grade_ex2a(job)

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

