# imports
import math

# importing Qiskit
from qiskit import QuantumCircuit, Aer, assemble, transpile
from qiskit_aer import AerSimulator
from qiskit.visualization import plot_histogram

f = open("cnf_test2.txt")
info = f.readline().split(' ')
n_var = int(info[0])
n_clause = int(info[1])

lines = [line.rstrip('\n') for line in f]

formula = []
for i in lines:
    l = i.split(' ')
    formula.append([int(l[0]),int(l[1]),int(l[2])])

n_qubits = n_var + 3*n_var*(2 + n_clause) + n_clause + 1
n_measures = 7*n_var
iters = 3*n_var

qc = QuantumCircuit(n_qubits)

for i in range(n_var):
    qc.h(i)

ini = n_var+n_clause
theta = 2*math.acos(math.sqrt(2/3))
for i in range(iters):
    qc.u(theta,0,0,ini)
    qc.x(ini)
    qc.ch(ini, ini+1)
    qc.x(ini)
    ini += 2+n_clause

qc.x(n_qubits-1)
qc.h(n_qubits-1)

initial = qc.to_gate()
initial.name = "Initialization"
qc.draw()

qc = QuantumCircuit(n_qubits)

for it in range(iters):
    first_bit = n_var+(2+n_clause)*it+n_clause
    first_clause = n_var+(2+n_clause)*it
    nc = 0
    for clause in formula:
        control = list()

        for lit in clause:
            lit_aux = abs(lit)-1
            if lit < 0:
                qc.x(lit_aux)
            control.append(lit_aux)

        qc.mct(control,first_clause+nc)

        for lit in clause:
            if lit < 0:
                lit_aux = abs(lit)-1
                qc.x(lit_aux)

        nc += 1

    nc = 0
    for clause in formula:
        lit1 = abs(clause[0])-1
        lit2 = abs(clause[1])-1
        lit3 = abs(clause[2])-1
        control = list()
        control.append(first_bit)
        control.append(first_bit+1)
        control.append(first_clause+nc)


        for n in range(nc):
            control.append(first_clause+n)
            qc.x(first_clause+n)

        qc.x(first_bit)
        qc.x(first_bit+1)
        qc.mct(control,lit1)
        qc.x(first_bit+1)
        qc.x(first_bit)

        qc.x(first_bit+1)
        qc.mct(control,lit2)
        qc.x(first_bit+1)

        qc.x(first_bit)
        qc.mct(control,lit3)
        qc.x(first_bit)

        for n in range(nc):
            qc.x(first_clause+n)

        nc += 1


nc = 0
for clause in formula:
    control = list()

    for lit in clause:
        lit_aux = abs(lit)-1
        if lit < 0:
            qc.x(lit_aux)
        control.append(lit_aux)

    qc.mct(control,n_qubits-1-n_clause+nc)

    for lit in clause:
        if lit < 0:
            lit_aux = abs(lit)-1
            qc.x(lit_aux)

    nc += 1

    
schoning_step = qc.to_gate()
schoning_step.name = "Schoning"
qc.draw(fold=10000)

qc = QuantumCircuit(n_qubits)

qc.append(schoning_step, list(range(n_qubits)))

control = list()
for nc in range(n_clause):
    qc.x(n_qubits-1-n_clause+nc)
    control.append(n_qubits-1-n_clause+nc)

qc.mct(control,n_qubits-1)

for nc in range(n_clause):
    qc.x(n_qubits-1-n_clause+nc)

qc.append(schoning_step.inverse(), list(range(n_qubits)))

oracle = qc.to_gate()
oracle.name = "Oracle"
qc.draw()

qc = QuantumCircuit(n_qubits)

l = list()
aux = list()

ini = n_var
for i in range(iters):
    for j in range(n_clause):
        aux.append(ini+j)
    ini += n_clause+2

for j in range(n_clause):
        aux.append(ini+j)


for i in range(n_var):
    qc.h(i)
    qc.x(i)
    l.append(i)

ini = n_var+n_clause
theta = 2*math.acos(math.sqrt(2/3))
for i in range(iters):
    qc.x(ini)
    qc.ch(ini, ini+1)
    qc.x(ini)
    qc.u(-theta,0,0,ini)

    qc.x(ini)
    qc.x(ini+1)
    l.append(ini)
    l.append(ini+1)
    ini += 2+n_clause

#qc.mct(l,n_qubits-1)
original = l
total = len(l)
l2 = list()
i = 0
while total != 1:
    for j in range(math.ceil(total/4)):
        qc.mct(l[4*j:4*j+4],aux[i])
        l2.append(aux[i])
        i+=1

    l = l2
    total = len(l)
    l2 = list()

diff_part1 = qc.to_gate()
diff_part1.name = "diff_part1"
qc.draw()

qc = QuantumCircuit(n_qubits)

qc.append(diff_part1, list(range(n_qubits)))

qc.mct(l,n_qubits-1)

qc.append(diff_part1.inverse(), list(range(n_qubits)))

diffuser = qc.to_gate()
diffuser.name = "Diffuser"
qc.draw()

rounds = math.pi/4.0*math.sqrt((4/3)**n_var)
print("Result:",rounds)
print("Iterations R ≤",math.ceil(rounds))

if math.floor(rounds) == 0:
    rounds = round(rounds)
else:
    rounds = math.floor(rounds)

print("Iterations chosen:", rounds)

qc = QuantumCircuit(n_qubits,n_measures)

qc.append(initial, list(range(n_qubits)))
for i in range(rounds):
    qc.append(oracle, list(range(n_qubits)))
    qc.append(diffuser, list(range(n_qubits)))

actual = 0

for i in range(n_var):
    qc.measure([i],[i])
    actual += 1

ini = n_var+n_clause
for i in range(iters):
    qc.measure([ini],[actual])
    qc.measure([ini+1],[actual+1])
    ini += 2+n_clause
    actual += 2

qc.draw()

# Select the AerSimulator from the Aer provider
simulator = AerSimulator(method='matrix_product_state')
#simulator = AerSimulator(method='statevector')
#simulator = AerSimulator(method='extended_stabilizer')

# Run and get counts, using the matrix_product_state method
tcirc = transpile(qc, simulator)
result = simulator.run(tcirc, shots=1).result()
print('This succeeded?: {}'.format(result.success))
counts = result.get_counts(0)
print(counts)

assig = list(next(iter(counts))[::-1][0:n_var])
steps = next(iter(counts))[::-1][n_var:]

print(assig)
print(steps)

def check_formula(formula, assig):
    c = 0
    for clause in formula:
        lit1_pos = abs(clause[0])-1
        lit2_pos = abs(clause[1])-1
        lit3_pos = abs(clause[2])-1
        total = 0
        
        if (clause[0] < 0 and assig[lit1_pos] == '0') or (clause[0] > 0 and assig[lit1_pos] == '1'):
            total += 1
            
        if (clause[1] < 0 and assig[lit2_pos] == '0') or (clause[1] > 0 and assig[lit2_pos] == '1'):
            total += 1
        
        if (clause[2] < 0 and assig[lit3_pos] == '0') or (clause[2] > 0 and assig[lit3_pos] == '1'):
            total += 1
            
        if total == 0:
            return (False,c)
        
        c += 1
    
    return (True,-1)

def schoning(formula,assig,steps,verbose=True):
    
    n = len(assig)
    
    if verbose:
            print("initial guess","->",assig)
            print("")

    for i in range(3*n):
        checking = check_formula(formula, assig)
        if checking[0]:
            return True
        
        lit = int(steps[2*i])+2*int(steps[2*i+1])
        var = abs(formula[checking[1]][lit])-1
        
        if assig[var] == '0':
            assig[var] = '1'
        else:
            assig[var] = '0'
        
        if verbose:
            print("iter",i+1,"->",assig)
            
    return check_formula(formula, assig)[0]

schoning(formula,assig,steps)

simulator = AerSimulator(method='matrix_product_state')
tcirc = transpile(qc, simulator)
result = simulator.run(tcirc, shots=1024).result()
print('This succeeded?: {}'.format(result.success))
counts = result.get_counts(0)

total = 0
true = 0
for key in counts:
    assig = list(key[::-1][0:n_var])
    steps = key[::-1][n_var:]
    r = schoning(formula,assig,steps,verbose=False)
    if r:
        true += counts[key]
    total += counts[key]
print(true/total*100,"% of success")