%matplotlib inline
from qiskit import QuantumCircuit, transpile
from qiskit.visualization import plot_histogram, array_to_latex
from qiskit.circuit.library import QFT, GroverOperator
from qiskit.primitives import Sampler
from qiskit_aer import AerSimulator
from math import pi, sqrt, sin
import matplotlib.pyplot as plt
from itertools import product
import pandas as pd
import functools
from utils import get_image_path
from utils import fibonacci as recursive_fibonacci
from tqdm.notebook import tqdm_notebook

def simulate(qc, shots=1000, backend=AerSimulator()):
    transpiled_circuit = transpile(qc, backend) if not isinstance(backend, Sampler) else qc
    return backend.run(transpiled_circuit, shots=shots)
    
def get_counts(job):
    return job.result().get_counts()


def fibonacci(n):
    qc = QuantumCircuit(n)
    qc.ry(pi/2, range(n))

    if(n > 1):
        for i in range(0, n-1):
            qc.cry(-pi/2, i, i+1)
    return qc

qc = fibonacci(4)
qc.measure_all()
qc.draw('mpl', filename=get_image_path("fibonacci-circuit.png"))

job = simulate(qc)
counts = get_counts(job)

def complete_results(counts):
    results = {}
    for bit_string in product('01', repeat=len(qc.qubits)):
        bit_string = ''.join(bit_string)
        results[bit_string] = 0 if counts.get(bit_string) is None else counts[bit_string]
    return results

fig= plot_histogram(complete_results(counts))
fig.savefig(get_image_path('fibonacci-4.png'))
display(fig)

def get_fibonacci_number(counts):
    return len(counts)

# F(4)
get_fibonacci_number(counts)

def evaluate_fibonacci_quantum(n, shots=1000):
    qc = fibonacci(n)
    qc.measure_all()
    job = simulate(qc, shots=shots)
    counts = get_counts(job)
    return get_fibonacci_number(counts)

x_values = list(range(1,15))
y_values = {"classical":[recursive_fibonacci(n+1) for n in x_values], "quantum":[]}

def comparation(shots=1000):
    y_values["quantum"] = []
    for n in x_values:
        y_values["quantum"].append(evaluate_fibonacci_quantum(n, shots))
    return y_values

y_values_1000 = comparation(shots=1000)
y_values_4000 = comparation(shots=4000)
y_values_6000 = comparation(shots=6000)

x = pd.Series(x_values)

y_1000 = pd.DataFrame(y_values_1000)
y_4000 = pd.DataFrame(y_values_4000)
y_6000 = pd.DataFrame(y_values_6000)

f, (ax1, ax2, ax3) = plt.subplots(1, 3, sharey=True, figsize=(30, 10))

#plt.title("Comparation between quantum and classical Fibonacci algorithm")
ax1.plot(x, y_1000)
ax1.legend(labels=["classical", "quantum"])
ax1.set_title("shots=1000")

ax2.plot(x, y_4000)
ax2.legend(labels=["classical", "quantum"])
ax2.set_title("shots=4000")

ax3.plot(x, y_6000)
ax3.legend(labels=["classical", "quantum"])
ax3.set_title("shots=6000")

plt.show()
f.savefig(get_image_path('comparation.png'))

def fibonacci_oracle(n):
    qc = QuantumCircuit(n, name="oracle")
    
    #from https://arxiv.org/pdf/2006.10656.pdf
    predicate = lambda k: k & k >> 1 == 0

    max_value = int('1'*n, 2)
    for k in range(max_value):
        if(predicate(k)):
            bit_string = format(k, 'b').zfill(n)[::-1]
            zeros = [index for index, bit in enumerate(bit_string) if bit == '0' ] 
            qc.x(zeros)
            qc.mcp(pi, list(range(1, n)), 0)
            qc.x(zeros)
            
    return qc

def show_unitary(qc):
    sim = AerSimulator()
    qc_unitary = qc.copy()
    qc_unitary.save_unitary()
    unitary = sim.run(transpile(qc_unitary, sim), shots=1000).result().get_unitary()
    display(array_to_latex(unitary, max_size=100000))

n = 4
oracle = fibonacci_oracle(n)
show_unitary(oracle)
oracle.draw('mpl')

def create_grover_operator(n_iterations, oracle):
    grover = GroverOperator(oracle, name="grover").repeat(n_iterations).to_gate().control()
    return grover

def create_qft_dagger(n):
    return QFT(n, inverse=True).to_gate(label="QFT†")

def create_quantum_couting_algorithm(n, oracle):
    total_qubits = n*2

    qc = QuantumCircuit(total_qubits, n)
    qc.h(range(total_qubits))
    iterations = 1
    for qubit in range(n):
        qc.append(create_grover_operator(iterations, oracle), [qubit] + list(range(n, total_qubits)))
        iterations *= 2
    
    qc.append(create_qft_dagger(n), list(range(n)))
    qc.measure(range(n), range(n))
    
    return qc

qc = create_quantum_couting_algorithm(n, oracle)
qc.draw('mpl', filename=get_image_path("quantum_counting.png"))

def get_dist(job):
    dists = job.result().quasi_dists
    print(dists)
    return dists[0]

def get_total_results(dist):
    measured_str = max(dist, key=dist.get)
    theta =  (measured_str/(2**n))*pi*2
    N = 2 ** n
    M = N * (sin(theta/2)**2)
    return M

def evaluate_fibonacci_quantum_2(n, shots):
    progress = tqdm_notebook(total=3)
    
    oracle = fibonacci_oracle(n)
    progress.update(1)
    
    qc = create_quantum_couting_algorithm(n, oracle)
    progress.update(1)
    
    dist = get_dist(simulate(qc, backend=Sampler()))
    progress.update(1)
    
    return get_total_results(dist)

x_values = list(range(2,7))
y_values = {"classical":[recursive_fibonacci(n+1) for n in x_values], "quantum":[]}

def comparation_2(shots=1000):
    y_values["quantum"] = []
    
    for n in x_values:
        y_values["quantum"].append(evaluate_fibonacci_quantum_2(n, shots))
    return y_values

y_values_1000 = comparation_2(shots=1000)
y_values_4000 = comparation_2(shots=4000)
y_values_6000 = comparation_2(shots=6000)
y_values_10000 = comparation_2(shots=10000)

x = pd.Series(x_values)

y_1000 = pd.DataFrame(y_values_1000)
y_4000 = pd.DataFrame(y_values_4000)
y_6000 = pd.DataFrame(y_values_6000)
y_10000 = pd.DataFrame(y_values_10000)

f, ax = plt.subplots(2, 2, sharey=True, figsize=(30, 10))

ax1 = ax[0][0]
ax2 = ax[0][1]
ax3 = ax[1][0]
ax4 = ax[1][1]

#plt.title("Comparation between quantum and classical Fibonacci algorithm")
ax1.plot(x, y_1000)
ax1.legend(labels=["classical", "quantum"])
ax1.set_title("shots=1000")

ax2.plot(x, y_4000)
ax2.legend(labels=["classical", "quantum"])
ax2.set_title("shots=4000")

ax3.plot(x, y_6000)
ax3.legend(labels=["classical", "quantum"])
ax3.set_title("shots=6000")

ax4.plot(x, y_10000)
ax4.legend(labels=["classical", "quantum"])
ax4.set_title("shots=10000")

plt.show()
f.savefig(get_image_path('comparation2.png'))