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

# importing Qiskit
from qiskit import Aer, transpile, assemble
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
from qiskit.circuit.library import QFT

# import basic plot tools
from qiskit.visualization import plot_histogram, plot_bloch_multivector

def qft_dagger(qc, n):
    """n-qubit QFTdagger the first n qubits in circ"""
    # Don't forget the Swaps!
    for qubit in range(n//2):
        qc.swap(qubit, n-qubit-1)
    for j in range(n):
        for m in range(j):
            qc.cp(-math.pi/float(2**(j-m)), m, j)
        qc.h(j)

#Modified to remove prints

def generalised_qpe_timing(amt_estimation_qubits, angle, shots=4096):
    # Create and set up circuit
    qpe3 = QuantumCircuit(amt_estimation_qubits+1, amt_estimation_qubits)

    # Apply H-Gates to counting qubits:
    for qubit in range(amt_estimation_qubits):
        qpe3.h(qubit)

    # Prepare our eigenstate |psi>:
    qpe3.x(amt_estimation_qubits)

    repetitions = 1
    for counting_qubit in range(amt_estimation_qubits):
        for i in range(repetitions):
            qpe3.cp(angle, counting_qubit, amt_estimation_qubits);
        repetitions *= 2

    # Do the inverse QFT:
    qft_dagger(qpe3, amt_estimation_qubits)

    # Measure of course!
    qpe3.barrier()
    for n in range(amt_estimation_qubits):
        qpe3.measure(n,n)

    #print(qpe3)
    # Let's see the results!
    aer_sim = Aer.get_backend('aer_simulator')
    t_qpe3 = transpile(qpe3, aer_sim)
    qobj = assemble(t_qpe3, shots=shots)
    results = aer_sim.run(qobj).result()
    answer = results.get_counts()
    ##return(plot_histogram(answer))
    ##comment out the return if you want to see the histogram
    ## do not output value to not skew timing
    return ((int(answer.most_frequent(), 2)/2**amt_estimation_qubits))
    
    
generalised_qpe_timing(8,(2*math.pi/3),10000)
#generalised_qpe(2,(2*math.pi/4))

%timeit -p 8 -r 1 -n 1 generalised_qpe_timing(8,(2*math.pi/3),10000)
%timeit -p 8 -r 1 -n 1 generalised_qpe_timing(8,(2*math.pi/3),10000)
%timeit -p 8 -r 1 -n 1 generalised_qpe_timing(8,(2*math.pi/3),10000)
%timeit -p 8 -r 1 -n 1 generalised_qpe_timing(8,(2*math.pi/3),10000)
%timeit -p 8 -r 1 -n 1 generalised_qpe_timing(8,(2*math.pi/3),10000)
%timeit -p 8 -r 1 -n 1 generalised_qpe_timing(8,(2*math.pi/3),10000)
%timeit -p 8 -r 1 -n 1 generalised_qpe_timing(8,(2*math.pi/3),10000)
%timeit -p 8 -r 1 -n 1 generalised_qpe_timing(8,(2*math.pi/3),10000)
%timeit -p 8 -r 1 -n 1 generalised_qpe_timing(8,(2*math.pi/3),10000)
%timeit -p 8 -r 1 -n 1 generalised_qpe_timing(8,(2*math.pi/3),10000)