#initialization
import matplotlib.pyplot as plt
import numpy as np

# importing Qiskit
from qiskit import IBMQ, Aer, assemble, transpile
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
from qiskit.providers.ibmq import least_busy

# import basic plot tools
from qiskit.visualization import plot_histogram

# general diffuser copied from https://qiskit.org/textbook/ch-algorithms/grover.html#2qubits-simulation
def diffuser(nqubits):
    qc = QuantumCircuit(nqubits)
    # Apply transformation |s> -> |00..0> (H-gates)
    for qubit in range(nqubits):
        qc.h(qubit)
    # Apply transformation |00..0> -> |11..1> (X-gates)
    for qubit in range(nqubits):
        qc.x(qubit)
    # Do multi-controlled-Z gate
    qc.h(nqubits-1)
    qc.mct(list(range(nqubits-1)), nqubits-1)  # multi-controlled-toffoli
    qc.h(nqubits-1)
    # Apply transformation |11..1> -> |00..0>
    for qubit in range(nqubits):
        qc.x(qubit)
    # Apply transformation |00..0> -> |s>
    for qubit in range(nqubits):
        qc.h(qubit)
    return qc

# THIS IS THE MAIN IMPORTANT PART OF THE IMPLEMENTATION
# we made lots of stubs for unimportant but time-consuming functions
# After putting around 10-15 hours in to the project, we realized it would take way too much time to hammer in every
# single function.

def oracle(nqubits, expressions):
#     iterate over expressions and define circuit
    qc = QuantumCircuit(nqubits + len(expressions))
    x_reset_target_indices = [] # list of target qubits which have x gates that need to be reversed
    for i, exp in enumerate(expressions):
        # if OR: apply X gates
        if (exp[0] == "OR"):
            for variable in exp[1:]:
                qubit_index = int(variable[1]) # each variable is in format xi where i is the qubit index
                qc.x(qubit_index)
                
        # apply multi qubit toffoli
        control_indices = []
        x_reset_input_indices = [] # indices to be reset if we apply x gates
        
        for variable in exp[1:]:
            if (variable[0] == "x"):
                qubit_index = int(variable[1])
                control_indices.append(qubit_index)
            else:
                # variable[0] = "e" -> the qubit index is the target qubit from the prior expression ei
                qubit_index = nqubits + int(variable[1]) # this references the target qubit after evaluating expression ei
                control_indices.append(qubit_index)
            if (variable[len(variable-1)] == "~"):
                qc.x(qubit_index)
                x_reset_indices.append(qubit_index)
                
        target_index = nqubits + i
        qc = multi_qubit_toffoli(qc, control_indices, target_index)
        
        for index in x_reset_input_indices:
            # reset the inputs if we applied x gates 
            qc.x(index)
        
        # if OR: apply X gate on target qubit for logic and X gates on control qubits to reset input
        if (exp[0] == "OR"):
            qc.x(target_index)
            x_reset_target_indices.append(i) # append to this list so that we can reset the circuit at end of oracle
            for variable in exp[1:]:
                qubit_index = int(variable[1])
                qc.x(qubit_index)
    for index in x_reset_target_indices:
        # reset the rest of the circuit
        qc.x(index)
    return qc

def multi_qubit_toffoli(qc, control_indices, target_index):
#     perform the multi qubit toffoli gate on the quantum circuit and return the circuit
    return qc

def init(nqubits):
    qc = QuantumCircuit(nqubits)
    for i in range(nqubits):
        qc.h(i)
    return qc

def parse_exp(exp):
# returns something like [e1=["OR", "x0", "x1~"], e2=["AND", "e1", "x2"], e3=[...], ...]
# where each expression ei is connected by a series of ANDs
# make sure that any negations get distributed
# We will write our expressions as ei=["AND", "x0", "x1~"] so that the last character is "~" if the variable has negation
    li = []
    return li

def check_solution(boolean_list, expression):
    # boolean list would represent the variables in the expression by index
    return false

expression = "((x0 & x1) & (x1 & ~x2) & x3)"
nqubits = 2
circuit = init(nqubits)
oracle_qc = oracle(nqubits, expression)
diffuser_qc = diffuser(nqubits)

# now join the circuits together
solved = false
boolean_list
while (solved):
#     add oracle and diffuser on our circuit
    circuit = circuit + oracle_qc + diffuser_qc
#     measure result
    results = circuit.measure_all().counts()
    boolean_list = results
#     check solution with expression
    solved = checks_solution(boolean_list, expression)
print(boolean_list)