import numpy as np
from math import pi
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit import Aer, execute
from qiskit.circuit import Gate
from qiskit.visualization import plot_histogram
from qiskit.circuit.library.standard_gates import ZGate, XGate

# Define circuit parameters
a = '10x01xx'
n = len(a)
k = 7

circ = QuantumCircuit(k+n, k)
circ.reset(range(k+n))
circ.h(range(k+n))
circ.barrier()
for i in range(k):
    name = '$G_' + str(i) + '$'
    circ.append(QuantumCircuit(n, name=name).to_gate().control(1), [n+i]+list(range(n)))
circ.barrier()
circ.append(Gate(name="$QFT^{-1}$", num_qubits=k, params=[]), range(n, k+n))
circ.barrier()
circ.measure(range(n, k+n), range(k))
circ.draw('mpl', reverse_bits=True, scale=0.5)

def Gop(j):
    p = 2 ** j
    ctrl_bits = []
    ctrl_state = ''
    for i in range(n):
        if a[n-i-1] != 'x':
            ctrl_bits.append(i+1)
            ctrl_state += a[n-i-1]
    G = QuantumCircuit(n+1, name='G'+str(j))
    for i in range(p):
        G.append(XGate().control(len(ctrl_bits), ctrl_state=ctrl_state[::-1]), ctrl_bits + [0])
        G.h(range(1, n+1))
        G.x(range(1, n+1))
        G.append(ZGate().control(n-1), reversed(range(1, n+1)))
        G.x(range(1, n+1))
        G.h(range(1, n+1))
    return G

Gop(3).draw('mpl', reverse_bits=True, scale=0.5)

QFT_inv = QuantumCircuit(k, name='QFT^{-1}')
for i in reversed(range(k)):
    if i != k-1:
        QFT_inv.barrier()
    for j in reversed(range(i+1,k)):
        QFT_inv.cu1(-pi/(2 ** (j-i)), i, j)
    QFT_inv.h(i)
QFT_inv.draw('mpl', reverse_bits=True)

circ = QuantumCircuit(k+n+1, k)
circ.reset(range(k+n+1))
circ.x(0)
circ.h(range(k+n+1))
circ.z(n+1)
circ.barrier()
for i in range(k):
    circ.append(Gop(i).to_gate().control(1), [n+i+1]+list(range(n+1)))
circ.barrier()
circ.append(QFT_inv, range(n+1, k+n+1))
circ.barrier()
circ.measure(range(n+1, k+n+1), range(k))
circ.draw(reverse_bits=True, scale=0.5)

delta = 64
simulator = Aer.get_backend('qasm_simulator')
result = execute(circ, simulator, shots=delta).result()
counts = result.get_counts(circ)
plot_histogram(counts)

x = list(counts.keys())
x = [pi*int(i[::-1], 2)/(2 ** k) for i in x]
p = list(counts.values())
p = p.index(max(p))
theta = min(x[p], pi-x[p])
m_estimate = (2 ** n) * (theta ** 2)
m = 2 ** a.count('x')
print('Estimated Count:')
print(m_estimate)
print('Actual Count:')
print(m)