#initialization
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
import math

# importing Qiskit
from qiskit import IBMQ, Aer
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister, execute

# import basic plot tools
from qiskit.visualization import plot_histogram

from qiskit.circuit.library import QFT
def qft_dagger(circ, qubits, n):
    circ.append(QFT(n).inverse(), qubits)
    
# Add UROTx based on CU1 gates
def apply_crot(circ, control_qubit, target_qubit, theta, exponent):
    circ.cu1(2*math.pi*theta*exponent, control_qubit, target_qubit)

def prepare_circuit(n, theta):
    qpe = QuantumCircuit(n+1,n)
    qpe.x(n)
    qpe.h(range(n))

    for x in range(n):
        exponent = 2**(n-x-1)
        apply_crot(qpe, x, n, theta, exponent)

    qpe.barrier()
    qft_dagger(qpe, range(n), n)
    qpe.barrier()
    for i in range(n):
        qpe.measure(i,i)

    return qpe

# The idea is that we are able to estimate theta (the angle) thanks to the QFT Dagger.
# The more qubits we apply on top, the more accuracy we get. Take a look

import operator

backend = Aer.get_backend('qasm_simulator')
shots = 2048
theta = 0.24125
for qubits in range(2, 15):
    circuit = prepare_circuit(qubits, theta)
    results = execute(circuit, backend=backend, shots=shots).result()
    counts = results.get_counts(circuit)
    
    highest_probability_outcome = max(counts.items(), key=operator.itemgetter(1))[0][::-1]
    measured_theta = int(highest_probability_outcome, 2)/2**qubits
    print("Using %d qubits with theta = %.6f, measured_theta = %.6f." % (qubits, theta, measured_theta))


## Look at the lovely circuit
circuit.draw('text')

