%matplotlib inline
from qiskit import QuantumCircuit, transpile
from qiskit.circuit.library import QFT
from qiskit.visualization import plot_histogram, plot_bloch_multivector
from qiskit.quantum_info import Statevector
from qiskit_aer import AerSimulator
import numpy as np
from fractions import Fraction
from utils import get_image_path

def create_qpe(phase, total_counting_qubits=3):
    total_qubits = total_counting_qubits+1
    
    qc = QuantumCircuit(total_qubits, total_counting_qubits)
    qc.x(-1)
    qc.h(range(total_counting_qubits))
    
    repetitions = 1
    for qubit in range(total_counting_qubits):
        for _ in range(repetitions):
            qc.cp(phase, qubit, -1)
        repetitions *= 2
    
    qc.append(QFT(total_counting_qubits, inverse=True), range(total_counting_qubits))
    qc.measure(range(total_counting_qubits), range(total_counting_qubits))
    return qc

sim = AerSimulator()

counting_qubits = 3
lambda_ = np.pi
qc = create_qpe(lambda_)
display(qc.draw('mpl', filename=get_image_path('qpe.png')))
job = sim.run(transpile(qc, sim), shots=1000).result().get_counts()
display(plot_histogram(job))

def get_theta(counting_qubits, counts):
    N = 2**(counting_qubits)
    estimated_value = int(max(counts, key=counts.get), 2)
    return Fraction(estimated_value, N)

theta = get_theta(counting_qubits, job)
theta

def get_phi(counting_qubits, theta):
    fraction = Fraction(2*float(theta)).limit_denominator(2**counting_qubits)
    return float(fraction)*np.pi 

phi = get_phi(counting_qubits, theta)
phi

def check(phi):
    qc = QuantumCircuit(1)
    qc.x(0)
    qc.p(phi, 0)
    display(Statevector(qc).draw('latex'))

check(phi)

lambda_ = np.pi/2
counting_qubits = 4
qc = create_qpe(lambda_, counting_qubits)
job = sim.run(transpile(qc, sim), shots=1000).result().get_counts()
plot_histogram(job)

theta = get_theta(counting_qubits, job)
phi = get_phi(counting_qubits, theta)
print(f"theta: {theta}\nphi:{phi}")

check(phi)