from qiskit import QuantumCircuit
from qiskit.tools.visualization import plot_histogram
from qiskit.extensions import UnitaryGate
import numpy as np

def logicalAnd(circuit,input_qubits,output_qubit):
    circuit.mcx(input_qubits,output_qubit)
    
def logicalOr(circuit,input_qubits,output_qubit):
    for i in input_qubits:
        circuit.x(i)
    circuit.x(output_qubit)
    circuit.mcx(input_qubits,output_qubit)
    for i in input_qubits:
        circuit.x(i)

def CNF(maxterms):
    n = len(maxterms[0])
    m = len(maxterms)
    circuit = QuantumCircuit(n+m+1)
    o = n
    for t in maxterms:
        inp = []
        for i in range(0,n):
            if t[i] == -1:
                circuit.x(i)
                inp.append(i)
            elif t[i] == 1:
                inp.append(i)
        logicalOr(circuit,inp,o)
        for i in range(0,n):
            if t[i] == -1:
                circuit.x(i)
        o = o+1
    logicalAnd(circuit,list(range(n,n+m)),n+m)
    return circuit.to_gate()

def diffusion(n):
    m = np.identity(2**n)
    for i in range(1,2**n):
        m[i][i] = -1
    return UnitaryGate(m)

def grover(maxterms,rounds = 0):
    n = len(maxterms[0]) #number of variables
    N = 2**n #solution space size
    m = len(maxterms) #number of maxterms
    if rounds == 0:
        r = int(np.ceil(np.pi/(4*np.arcsin(1/np.sqrt(N)))-0.5))
    else:
        r = rounds
    circuit = QuantumCircuit(n+m+1,n)
    circuit.x(n+m)
    circuit.h(n+m)
    for i in range(n):
        circuit.h(i)
    cnf = CNF(maxterms)
    cnf.label = 'Uf'
    diff = diffusion(n)
    diff.label = 'Diff'
    for i in range(r):
        circuit.append(cnf,range(n+m+1))
        for i in range(n,n+m):
            circuit.reset(i)
        for i in range(n):
            circuit.h(i)
        circuit.append(diff,range(n))
        for i in range(n):
            circuit.h(i)
    for i in range(n):
        circuit.measure(i,i)
    return circuit

maxterms = [[1,1,1],[-1,-1,-1]] # This seems to behave like f(x,y,z) = xyz + x'y'z' (DNF) even though the function is CNF
circuit = grover(maxterms)
#circuit = grover(maxterms)
circuit.draw('mpl')

from qiskit import Aer, execute
simulator = Aer.get_backend('qasm_simulator')
result = execute(circuit, simulator, shots=1000).result()
counts = result.get_counts(circuit)
print(counts)
plot_histogram(counts)

