import numpy as np

from random import random
from qiskit import *
from qiskit.aqua.utils.controlled_circuit import get_controlled_circuit

num_bits_estimate = 3
# For 2x2 matrix one qubit is enough
q = QuantumRegister(1, name="q")
# In QPE we use n ancillas to estimate n bits from the phase
a = QuantumRegister(num_bits_estimate, name="a") 
# For n ancillary qubit measurment we need n cllasical bits
c = ClassicalRegister(num_bits_estimate, name="c") 

# Create a quantum circuit
circuit = QuantumCircuit(q, a, c)

# |1> eigenstate initialization
circuit.x(q[0])
circuit.draw('mpl')

E_1, E_2 = (2 * np.pi * random(), 2 * np.pi * random())
print("We are going to estimate E_2 via QPE algorithm \nE_2 = {}".format(E_2))

# circuit for unitary operator exp(iHt)
t = 1
unitary = QuantumCircuit(q)

unitary.u1(E_2 * t, q[0]) # q[0] is the only qubit in q register
unitary.x(q[0])
unitary.u1(E_1 * t, q[0])
unitary.x(q[0])
unitary.draw('mpl')

for ancillary in a:
    circuit.h(ancillary)

for n in range(a.size):
    for m in range(2**n):
        get_controlled_circuit(unitary, a[n], circuit)

# inverse QFT without SWAP gates
for n in reversed(range(a.size)):
    circuit.h(a[n])
    if n != 0:
        for m in reversed(range(n)):
            angle = -2*np.pi / (2**(n - m + 1))
            circuit.cu1(angle, a[n], a[m])

# measurements on the ancillary qubits stored in c classical register
for n in reversed(range(a.size)):
    circuit.measure(a[n],c[n])
circuit.draw('mpl')

backend = BasicAer.get_backend('qasm_simulator')
shots = 1024  # how many time execute the algorithm
job = execute(circuit, backend, shots=shots)
result = job.result()
counts = result.get_counts()

phase_bits = max(counts, key=counts.get) # take the most often obtaned result

phase = 0
for index, bit in enumerate(reversed(phase_bits)):
    phase += int(bit) / 2**(index + 1)
    
estimated_E_2 = 2 * np.pi * phase / t

print("Eigenvalue of the Hamiltonian: {}".format(E_2))
print("Estimated eigenvalue of the Hamiltonian: {}".format(estimated_E_2))

