import matplotlib.pyplot as plt
import numpy as np
import math
from numpy import pi

from qiskit import Aer, transpile, assemble
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
from qiskit.visualization import plot_histogram

# function definitions for QFT and inverse QFT 

def qft_rotations(circuit, n):
    """Performs qft on the first n qubits in circuit (without swaps)"""
    if n == 0:
        return circuit
    n -= 1
    circuit.h(n)
    for qubit in range(n):
        circuit.cp(pi/2**(n-qubit), qubit, n)
    # At the end of our function, we call the same function again on
    # the next qubits (we reduced n by one earlier in the function)
    qft_rotations(circuit, n)
    
def swap_registers(circuit, n):
    for qubit in range(n//2):
        circuit.swap(qubit, n-qubit-1)
    return circuit

def qft(circuit, n):
    """QFT on the first n qubits in circuit"""
    qft_rotations(circuit, n)
    swap_registers(circuit, n)
    return circuit

def inverse_qft(circuit, n):
    """Does the inverse QFT on the first n qubits in circuit"""
    # First we create a QFT circuit of the correct size:
    qft_circ = qft(QuantumCircuit(n), n)
    # Then we take the inverse of this circuit
    invqft_circ = qft_circ.inverse()
    # And add it to the first n qubits in our existing circuit
    circuit.append(invqft_circ, circuit.qubits[:n])
    return circuit.decompose() # .decompose() allows us to see the individual gates

nq_c = 4    # counting (precision) qubits
nq_a = 2    # qubits to store the eigenvector
nq = nq_c + nq_a
nc = nq_c   # classical bits

rq = QuantumRegister(nq)
rc = ClassicalRegister(nc)

qc = QuantumCircuit(rq, rc)

# prepare the last eigenvector uf the U operator
for i_a in range(nq_c, (nq_c+nq_a)):
     qc.x(rq[i_a])

# apply Hadamard gates on the counting qubits
for i_c in range(nq_c):
     qc.h(rq[i_c])

from qiskit.quantum_info.operators import Operator
from qiskit.extensions import PhaseGate

# prepare the U operator starting from a T-gate (PhaseGate)
# single qubit operator
theta=math.pi *(0.25 + 0.01)
u = Operator(PhaseGate(theta,label="U"))
#u.data

# two qubits operator
u = u.tensor(u)
#u.data

# build the controlled version of the U operator
rq_a = QuantumRegister(nq_a)
qc_a = QuantumCircuit(rq_a)
qc_a.append(u, rq_a)
u_gate = qc_a.to_gate()
cu_gate = qc_a.to_gate().control(1)

# apply the controlled U operator
repetitions = 1
for counting_qubit in range(nq_c):
    for i in range(repetitions):
        qc.append(cu_gate, [counting_qubit, nq_c, (nq_c+1)]); # This is C-U
    repetitions *= 2

# apply the inverse QFT
qc.barrier()
inverse_qft(qc, nq_c)

<qiskit.circuit.quantumcircuit.QuantumCircuit at 0x7f2e21dd8fd0>

# measure the counting qubits
qc.barrier()
for n in range(nq_c):
    qc.measure(n, n)

qasm_sim = Aer.get_backend('qasm_simulator')
shots = 2048
t_qc = transpile(qc, qasm_sim)
qobj = assemble(t_qc, shots=shots)
results = qasm_sim.run(qobj).result()
answer = results.get_counts()

plot_histogram(answer, figsize=(14,5))