#!pip install qiskit
import numpy as np
import matplotlib.pyplot as plt
from math import floor, pi, sqrt
from time import process_time
from qiskit import *
from qiskit.quantum_info.operators import Operator
from qiskit.visualization import plot_histogram
from qiskit.circuit import Gate


def get_bitstring_permutations(index, lst, n, args):
    """
    This function populates a list with all the combinations of bit strings of length n
    """
    if (index == n):
        # append combination to list
        lst.append(list.copy(args))
    else:
        # handle the case where the preceding bit is 0
        args[index] = 0
        get_bitstring_permutations(index + 1, lst, n, args)

        # handle the case where the preceding bit is 1
        args[index] = 1
        get_bitstring_permutations(index + 1, lst, n, args)

def generate_uf(f, n):
    """
    Parameters: f is an anonymous function and n is the number of bits in input: f:{0,1}^n -> {0,1}^n
    This function returns an oracle gate representing the function f for all x in {0,1}^n and y in {0,1}^n,
    the desired result is of the oracle is mapping the input x,y> to |x, y + f(x)> where + is addition modulo 2.
    The function first finds the list of bitstring permutations of n bits, it then establishes a mapping which is
    representative of the decimal number of the bitstring represents. For each |x,y>, it calculates
    |x, y + f(x)>. Finally it constructs a permutation gate which treats each permutation as a different basis vector
    in the 2^(n+1) dimensional complex Hilbert space that represents a system of 2*n qubits.
    Returns: the permutation gate
    """
    # generate list of all bitstrings of size n
    bitstrings = []
    get_bitstring_permutations(0, bitstrings, n + 1, [0] * (n + 1))

    # initialize mapping and permutation list
    perm_dict = dict()
    perm_list = []

    # populate mapping
    for permutation, bitstring in enumerate(bitstrings):
        values = [0] * (len(bitstrings))
        values[permutation] = 1
        perm_dict["".join(str(bit) for bit in bitstring)] = values

    # Send each |xy> to |x, f(x) + y>
    for bitstring in bitstrings:
        params = bitstring[:n]
        params.append((f(params) + bitstring[-1]) % 2)
        perm_list.append(perm_dict["".join(str(bit) for bit in params)])
    return Operator(np.array(perm_list))

def apply_H(circuit, apply_to_list):
    """
    Apply Hadamards to all specified qubits (if apply_to_list[index] == 1).
    Designed for a large amount of Hadamards being applied at once.
    """
    for index, qubit in enumerate(apply_to_list):
        if qubit == 1:
            circuit.h(index)
    return circuit

def initialize_dj(n):
    """
    This function sets initial states and applies Hadamards to each qubit
    Note: apply_H isn't called because it is actually more efficient to initialize in one loop as opposed to 2.
    """
    # apply H to first n qubits and X H to the last qubit (ancilla qubit)
    quantum_register = QuantumRegister(n + 1)
    classical_register = ClassicalRegister(n)
    quantum_circuit = QuantumCircuit(quantum_register, classical_register)
    # In qiskit, all quantum registers start in the low energy |0> state so we must apply an x gate to our helper bit
    # in order to put it in the state |1>
    quantum_circuit.x(quantum_register[-1])
    for index in range(n + 1):
        quantum_circuit.h(quantum_register[index])
    quantum_circuit.barrier()
    return quantum_circuit, quantum_register, classical_register

def initialize_bv(states):
    """
    This function sets initial states and applies Hadamards to each qubit
    Note: apply_H isn't called because it is actually more efficient to initialize in one loop as opposed to 2.
    """
    n = len(states)
    circuit = QuantumCircuit(n, n-1)
    for index, state in enumerate(states):
        if state == 1:
            circuit.x(index)
        circuit.h(index)
    return circuit




def deutsch_jozsa_algorithm(f, n, shots=1024, threshold=0.9):
    """
    This function is intended to determine if f is constant or balanced for a given function f s.t.
        f:{0,1}^n -> {0,1}. The algorithm initializes the qubits with H for the first n qubits and X and H for the last
        qubit. The algorithm then constructs a Uf oracle gate based on the function input f. It then applies Uf to all
        the qubits and applies H to the first n qubits. Finally, the simulator is run on the circuit and measures the
        results. If upon measurement, the first n qubits are all 0, 1 is returned and the function is constant,
        otherwise 0 is returned and the function is balanced.
    This function has an anonymous function and integer n as parameters.
    This function uses 9q-squared-qvm, so it assumes that n <= 9.
    """
    # apply H to first n qubits and X H to the last qubit (ancilla qubit)
    quantum_circuit, quantum_register, classical_register = initialize_dj(n)

    # Generate Uf oracle from f (anonymous function)
    uf_gate = generate_uf(f, n)
    # Applying the uf_gate must be done with the qubits in the reverse order due to the implementation of qiskit
    rv_qr = list()
    for index in range(n + 1):
        rv_qr.append(index)
    rv_qr.reverse()
    quantum_circuit.unitary(uf_gate, rv_qr)
    quantum_circuit.barrier()

    # Apply Hadamards to first n qubits
    for index in range(n):
        quantum_circuit.h(quantum_register[index])

    # Run simulator
    quantum_simulator = Aer.get_backend('qasm_simulator')
    quantum_circuit.measure(quantum_register[0:n], classical_register)
    # quantum_circuit.draw('mpl')
    # plt.show()

    # Evaluate the job results
#     start = process_time()
    job = execute(quantum_circuit, quantum_simulator, shots=shots)
#     end = process_time()
#     total_time = end - start
    results = job.result()
    counts = results.get_counts(quantum_circuit)

    # NOTE: 1 = constant, 0 = balanced
    # To compensate for the error rates of actual quantum machines, the threshold for being consider balanced is
    # set to be threshold * shots, or a percentage of shots given at runtime.

    # Function is constant
    key = '0' * n
    if key in counts:
        if counts[key] >= threshold * shots:
            return 1
    key = '1' * n
    if key in counts:
        if counts[key] >= threshold * shots:
            return 1

    # Function is balanced
    return 0

def f_balanced(args):
    return args[0]

def f_constant(args):
    return 1

input_size = 10

times = []
for i in range(1,input_size+1):
    start = process_time()
    res = deutsch_jozsa_algorithm(f_balanced, i, shots=1)
    times.append(process_time()-start)
print(times)
plt.plot(range(1,input_size+1), times)
for x,y in zip(range(1,input_size+1),times):

    label = "{:.3f}".format(y)
    plt.annotate(label, # this is the text
                 (x,y), # this is the point to label
                 textcoords="offset points", # how to position the text
                 xytext=(0,10), # distance from text to points (x,y)
                 ha='right') # horizontal alignment can be left, right or center
plt.xticks(np.arange(0, input_size+1, step=1))
plt.yticks(np.arange(0, 15, step=1))
plt.ylabel("Execution Time (s)")
plt.xlabel("Input Size (n)")
plt.title("DJ: Execution Time for Balanced Function vs. Input Size")

times = []
for i in range(1,input_size+1):
    start = process_time()
    res = deutsch_jozsa_algorithm(f_constant, i, shots=1)
    times.append(process_time()-start)
print(times)
plt.plot(range(1,input_size+1), times)
for x,y in zip(range(1,input_size+1),times):
    label = "{:.3f}".format(y)
    plt.annotate(label, # this is the text
                 (x,y), # this is the point to label
                 textcoords="offset points", # how to position the text
                 xytext=(0,10), # distance from text to points (x,y)
                 ha='right') # horizontal alignment can be left, right or center
plt.xticks(np.arange(0, input_size+1, step=1))
plt.yticks(np.arange(0, 14, step=1))
plt.ylabel("Execution Time (s)")
plt.xlabel("Input Size (n)")
plt.title("DJ: Execution Time for Constant Function vs. Input Size")

def bernstein_vazirani_algorithm(f, n):
    initialize_list = [0] * n

    # calculate b by f(0^n) = b
    b = f(initialize_list)

    # Initialize circuit by applying H to first n qubits and X H to last qubit (ancilla qubit)
    initialize_list.append(1)
    qubits = list(range(len(initialize_list)))
    circuit = initialize_bv(initialize_list)

    # Generate Uf oracle from f (anonymous function)
    uf_gate = generate_uf(f, n)
    # Applying the uf_gate must be done with the qubits in the reverse order due to the implementation of qiskit
    rv_qubits = qubits[::-1]
    circuit.unitary(uf_gate, rv_qubits)

    # Apply H to all qubits except for the last qubit
    apply_to_list = [1] * n
    apply_to_list.append(0)
    circuit = apply_H(circuit, apply_to_list)
    circuit.measure(range(n),range(n))

    # run simulator and measure qubits
    simulator = Aer.get_backend("qasm_simulator")
    job = execute(circuit, simulator, shots=1)
    result = job.result()
    counts = result.get_counts(circuit)
    plot_histogram(counts)
    for count in counts:
        a = count
    return a,b

def f_bv(x):
    n = len(x)
    y = 0
    for i,x_i in enumerate(x):
        y = (y + (x_i*(i%2)))%2
    y = (y+(n%2))%2
    return(y)

times = []
for i in range(1,input_size+1):
    start = process_time()
    res = bernstein_vazirani_algorithm(f_bv, i)
    times.append(process_time()-start)
print(times)
plt.plot(range(1,input_size+1), times)
for x,y in zip(range(1,input_size+1),times):

    label = "{:.3f}".format(y)
    plt.annotate(label, # this is the text
                 (x,y), # this is the point to label
                 textcoords="offset points", # how to position the text
                 xytext=(0,10), # distance from text to points (x,y)
                 ha='right') # horizontal alignment can be left, right or center
plt.xticks(np.arange(0, input_size+1, step=1))
plt.yticks(np.arange(0, 15, step=1))
plt.ylabel("Execution Time (s)")
plt.xlabel("Input Size (n)")
plt.title("BV: Execution Time vs. Input Size")

def get_Z0(n):
    """
    This function generates the Z0 gate satisfying the conditions for x in {0,1}^n Z0|x> -> -|x> iff x = 0^n
        otherwise Z0|x> -> |x>
    The parameter to this function is only the size n, a 2^n x 2^n dimensional matrix is created satisfying the
        conditions above.
    This function has one dependency, the DefGate function defined in pyquil.quil
    This function is designed to absorb the negative in G, so the returned gate is actually -Z0
    Returns -Z0
    """
    # Create a 2^n x 2^n matrix with all 0's
    gate = np.zeros((2 ** n, 2 ** n), dtype=int)
    # since this is -Z0, set first element to 1 not -1
    gate[0][0] = 1
    # set all other elements on the diagonal to -1, again not 1 because this is -Z0
    for i in range(1, 2 ** n):
        gate[i][i] = -1
    # Return gate
    return Operator(gate)


def get_Zf(f, n):
    """
    This function generates the Zf gate satisfying the condition for x in {0,1}^n where Zf|x> -> (-1)^f(X)|x>
    This function requires that f(x) be calculated for all x, so f is passed as an anonymous function, the other
        parameter is n.
    The function has one dependency, the DefGate function defined in pyquil.quil
    This function finds all permutations of bitstrings of length n, then initializes a 2^n x 2^n matrix of all 0's,
        and sets all elements along the diagonal to either 1 or -1 depending on f(x)
    Finally a gate representation of this matrix is returned.
    """
    # generate bitstring permutations
    bitstrings = list()
    get_bitstring_permutations(0, bitstrings, n, [0] * n)
    # initialize a 2^n x 2^n matrix of all 0's
    gate = np.zeros((2 ** n, 2 ** n), dtype=int)
    # set diagonals of matrix based on f(x)
    for i in range(2 ** n):
        gate[i][i] = -1 if f(bitstrings[i]) == 1 else 1
    # create and return gate
    return Operator(gate)

def grovers_algorithm(f, n, shots=1024, threshold=0.9):
    """
    This function is intended to determine if there exists an x in {0,1}^n s.t. f(x) = 1 for a given function f s.t.
        f:{0,1}^n -> {0,1}. The algorithm first constructs Zf, -Z0 gates, initializes with Hanamard matrices, and
        applies G = -H^n o Z0 o H^n o Zf. This algorithm is not deterministic, so G is applied multiple times. More
        specifically, G is run (pi / 4 * sqrt(n)) times. Furthermore, there are 10 trials to minimize the chance of a
        false negative.
    This function has an anonymous function and integer n as parameters.
    This function runs the algorithm as described for each 10 trials, and then checks if for any of the outputted states
        x, if f(x) = 1. If this is true, then 1 is returned, otherwise 0 is returned. The function returns 0 if there
        is an issue with the simulator.
    This function uses 9q-squared-qvm, so it assumes that n <= 9
    """
    # Initialize the circuit and apply Hadamards to all qubits
    quantum_circuit, quantum_register, classical_register = initialize_dj(n)

    # Needed for application of custom gates since Operator works in reverse
    rv_qr = list()
    for index in range(n):
        rv_qr.append(index)
    rv_qr.reverse()

    # Define and generate Z0 gate (really -Z0)
    z0_gate = get_Z0(n)

    # Define and generate Zf gate
    zf_gate = get_Zf(f, n)

    # Determine the number of times to apply G
    iteration_count = floor(pi / 4 * sqrt(2 ** n))
    # Apply G iteration_count times
    for i in range(iteration_count):
        # Apply Zf
        quantum_circuit.unitary(zf_gate, rv_qr)
        # Apply H to all qubits
        for index in range(n):
            quantum_circuit.h(quantum_register[index])
        # Apply -Z0
        quantum_circuit.unitary(z0_gate, rv_qr)
        # Apply H to all qubits
        for index in range(n):
            quantum_circuit.h(quantum_register[index])
        quantum_circuit.barrier()

    # Run simulator
    quantum_simulator = Aer.get_backend('qasm_simulator')
    quantum_circuit.measure(quantum_register[0:n], classical_register)

    # Display circuit diagram
    quantum_circuit.draw('mpl')
    plt.show()

    # Execute and evaluate the job results
    job = execute(quantum_circuit, quantum_simulator, shots=shots)
    results = job.result()
    counts = results.get_counts(quantum_circuit)
    return counts

    # Parse results and return 1 or 0 accordingly
#    dict = {}
#     for key in counts:
#         if counts[key] >= (shots/(2**n)):
#            dict[key] = counts[key]
#     for key in dict:
#         poop = list(key)
#         poop = [int(i) for i in poop]
#         poop.reverse()
#         if f(poop) == 1:
#             return 1
#     return 0



def f_0(args):
    return 0

def f_1(args):
    return 1

def f_100(args):
    n = len(args)
    marked = [0]*(n-1)
    marked.append(1)
    if(args == marked):
        return 1
    return 0

input_size = 8
times = []
for i in range(1,input_size+1):
    start = process_time()
    res = grovers_algorithm(f_0, i, shots=1024, threshold=0.9)
    times.append(process_time()-start)
print(times)
plt.plot(range(1,input_size+1), times)
for x,y in zip(range(1,input_size+1),times):

    label = "{:.4f}".format(y)
    plt.annotate(label, # this is the text
                 (x,y), # this is the point to label
                 textcoords="offset points", # how to position the text
                 xytext=(0,10), # distance from text to points (x,y)
                 ha='right') # horizontal alignment can be left, right or center
plt.xticks(np.arange(0, input_size+1, step=1))
plt.yticks(np.arange(0, 3.5, step=0.5))
plt.ylabel("Execution Time (s)")
plt.xlabel("Input Size (n)")
plt.title("Grover's: Execution Time of Zero-Function vs. Input Size")

times = []
for i in range(1,input_size+1):
    start = process_time()
    res = grovers_algorithm(f_1, i, shots=1024, threshold=0.9)
    times.append(process_time()-start)
print(times)
plt.plot(range(1,input_size+1), times)
for x,y in zip(range(1,input_size+1),times):

    label = "{:.4f}".format(y)
    plt.annotate(label, # this is the text
                 (x,y), # this is the point to label
                 textcoords="offset points", # how to position the text
                 xytext=(0,10), # distance from text to points (x,y)
                 ha='right') # horizontal alignment can be left, right or center
plt.xticks(np.arange(0, input_size+1, step=1))
plt.yticks(np.arange(0, 3.5, step=0.5))
plt.ylabel("Execution Time (s)")
plt.xlabel("Input Size (n)")
plt.title("Grover's: Execution Time of One-Function vs. Input Size")

times = []
for i in range(1,input_size+1):
    start = process_time()
    res = grovers_algorithm(f_100, i, shots=1024, threshold=0.9)
    times.append(process_time()-start)
print(times)
plt.plot(range(1,input_size+1), times)
for x,y in zip(range(1,input_size+1),times):

    label = "{:.4f}".format(y)
    plt.annotate(label, # this is the text
                 (x,y), # this is the point to label
                 textcoords="offset points", # how to position the text
                 xytext=(0,10), # distance from text to points (x,y)
                 ha='right') # horizontal alignment can be left, right or center
plt.xticks(np.arange(0, input_size+1, step=1))
plt.yticks(np.arange(0, 3.5, step=0.5))
plt.ylabel("Execution Time (s)")
plt.xlabel("Input Size (n)")
plt.title("Grover's: Execution Time of First-Bit-Function vs. Input Size")

def generate_uf_simons(f, n):
    """
    Parameters: f is an anonymous function and n is the number of bits in input: f:{0,1}^n -> {0,1}^n
    This function returns an oracle gate representing the function f for all x in {0,1}^n and y in {0,1}^n,
    the desired result is of the oracle is mapping the input x,y> to |x, y + f(x)> where + is addition modulo 2.
    The function first finds the list of bitstring permutations of n bits, it then establishes a mapping which is
    representative of the decimal number of the bitstring represents. For each |x,y>, it calculates
    |x, y + f(x)>. Finally it constructs a permutation gate which treats each permutation as a different basis vector
    in the 2^(n+1) dimensional complex Hilbert space that represents a system of 2*n qubits.
    Returns: the permutation gate
    """

    # generate list of all bitstrings of size n
    bitstrings = []
    get_bitstring_permutations(0, bitstrings, 2 * n, [0] * (2 * n))
    
    # initialize mapping and permutation list
    perm_dict = dict()
    perm_list = []
    
    # populate mapping
    for permutation, bitstring in enumerate(bitstrings):
        values = [0]*(len(bitstrings))
        values[permutation] = 1
        perm_dict["".join(str(bit) for bit in bitstring)] = values

    # Send each |xy> to |x, f(x) + y>
    for bitstring in bitstrings:
        params = bitstring[:n]
        params2 = bitstring[n:2 * n]
        f_values = f(bitstring[:n])
        for i in range(n):
            params.append((params2[i] + f_values[i]) % 2)
        perm_list.append(perm_dict["".join(str(bit) for bit in params)])
    return Operator(np.array(perm_list))

def simons_solver(Y, n):
    """
    Inputs: Y is a linear system of n-1 equations in matrix form. n is the dimension of the input into f.
    This function acts as a binary linear matrix solver.
    Returns: the key string s, if found, or the zero bitstring
    """
    # Create all possible bit strings to test for s
    bitstrings = []
    get_bitstring_permutations(0, bitstrings, n, [0] * n)

    # For each possible s, test to see if it's a candidate
    for s in bitstrings:
        if s == [0] * n:
            continue
        candidate = True

        # For each equation in Y, bit by bit test that y*s = [0]*n
        for y in Y:
            value = 0
            for i in np.arange(n):
                value = value + s[i] * y[i]

            # If a bit doesn't evaluate to 0...
            if (value % 2 == 1):
                candidate = False
        if (candidate):
            return s

    return [0] * n

def simons_algorithm(f, n):
    """
    Inputs: f is a blackbox function (f:{0,1}^n -> {0,1}^n) that is either one-to-one or two-to-one. n is the
    dimension of the input into f. This function finds and returns the key s, if one exists, for a two-to-one
    function by first creating a matrix U_f that represents f, then applying the appropriate quantum gates to
    generate a linear equation. By running the circuit until we generate n-1 unique equations, the set of equations can solve for s. The Classical solver returns s.
    Returns: the key string s, if found, or the zero bitstring
    """
    #Generate the oracle gate
    oracle = generate_uf_simons(f, n)
    #Initialize the circuit
    circuit = QuantumCircuit(2*n, 2*n)
    #initialize the simulator, use qasm_simulator
    simulator = Aer.get_backend("qasm_simulator")
    indices = list(range(2*n))
    indices.reverse()
    #apply Hadamards to first n qubits
    for i in range(n):
        circuit.h(i)
    #apply oracle gate
    circuit.unitary(oracle, indices, label="oracle")
    #apply Hadamards again to first n qubits
    for i in range(n):
        circuit.h(i)
    indices = list(range(n))
    #measure first n qubits
    circuit.measure(indices, indices)
    #Run the entire process 20 times
    for i in range(20):
        s = set()
        s_trials = []
        #Run quantum circuit until at least n-1 unique eqautions are obtained
        while(len(s) < n-1):
            job = execute(circuit, simulator, shots=1)
            result = job.result()
            counts = result.get_counts()
            for count in counts:
                s.add(count[2*n:n-1:-1])
        for bitstring in s:
            s_trials.append([int(bit) for bit in bitstring])
        s_trials = np.array(s_trials)
        #Solve system of equations
        val = simons_solver(s_trials, n)
        if val == [0] * n:
            continue
        #if the correct function value is found, no need to keep searching
        f_val = f(val)
        if f_val == f([0]*n):
            return val
    #s not found, return 0 bit string
    return [0] * n

def f_oto(x):
    return x

def f_tto(x):
    res = [0]
    if (len(x)>1):
        res = res + x[1:]
    return res


input_size = 5
times = []
for i in range(1,input_size+1):
    start = process_time()
    res = simons_algorithm(f_oto, i)
    times.append(process_time()-start)
print(times)

plt.plot(range(1,input_size+1), times)
for x,y in zip(range(1,input_size+1),times):

    label = "{:.4f}".format(y)
    plt.annotate(label, # this is the text
                 (x,y), # this is the point to label
                 textcoords="offset points", # how to position the text
                 xytext=(0,10), # distance from text to points (x,y)
                 ha='right') # horizontal alignment can be left, right or center
plt.xticks(np.arange(0, input_size+1, step=1))
plt.yticks(np.arange(0, 450, step=50))
plt.ylabel("Execution Time (s)")
plt.xlabel("Input Size (n)")
plt.title("Simon's: Execution Time of One-to-One Function vs. Input Size")

input_size = 5
times = []
for i in range(1,input_size+1):
    start = process_time()
    res = simons_algorithm(f_tto, i)
    times.append(process_time()-start)
print(times)

plt.plot(range(1,input_size+1), times)
for x,y in zip(range(1,input_size+1),times):

    label = "{:.4f}".format(y)
    plt.annotate(label, # this is the text
                 (x,y), # this is the point to label
                 textcoords="offset points", # how to position the text
                 xytext=(0,10), # distance from text to points (x,y)
                 ha='right') # horizontal alignment can be left, right or center
plt.xticks(np.arange(0, input_size+1, step=1))
plt.yticks(np.arange(0, 10, step=1))
plt.ylabel("Execution Time (s)")
plt.xlabel("Input Size (n)")
plt.title("Simon's: Execution Time of Two-to-One Function vs. Input Size")

