# Importing libraries
from qiskit import QuantumCircuit, execute, Aer, IBMQ, QuantumRegister, ClassicalRegister
from qiskit.compiler import transpile, assemble
from qiskit.tools.jupyter import *
from qiskit.visualization import *
import numpy as np
import matplotlib.pyplot as plt
from qiskit.converters import circuit_to_gate
from qiskit.circuit.library.standard_gates import ZGate


def oracle(n, targets):
    '''
    Given the number of qubits "n" and a list of target states "targets", the function returns the quantum circuit 
    for the oracle.
    '''
    qbits = QuantumRegister(n,'q')
    circ = QuantumCircuit(qbits, name="Oracle")
    
    # Creating a sequential circuit for each target state.
    for target in targets:
        
        # Each string representing a target state is reversed so that the indexing matches with the indices of
        # the corresponding qubits.
        target = list(target)
        target.reverse()
        target = ''.join(target)

        # Flip zero bits
        for i in range(len(target)):
            if target[i]=='0':
                circ.x(i)

        # Apply a n-1 CZ
        CZ = ZGate().control(n-1)
        circ.append(CZ, qbits)

        # Flip back
        for i in range(len(target)):
            if target[i]=='0':
                circ.x(i)
                
# Uncomment the following line if you the circuit to be converted to a gate. This step is not essential to the
# function of the quantum circuit.
#     circ = circuit_to_gate(circ)
            
    return circ


def diffusion(n):
    '''
    Given the number of qubits "n", the function returns the circuit for the n-qubit diffusion operation.
    '''
    qbits = QuantumRegister(n,'q')
    circ = QuantumCircuit(qbits, name="Diffusion")
    
    circ.h(qbits)
    # Call the oracle function with target state set to the n-qubit zero state.
    circ.append(oracle(n,['0'*n]), qbits)
    circ.h(qbits)
    
    return circ


def grover(n, oracle, m,threshold=0):
    '''
    The "grover" function returns the quantum circuit to implement Grover's search algorithm for an "n"-qubit 
    system to search for "m" targets. The argument "oracle" is the quantum circuit for the Oracle for the
    target state/states. The argument "threshold" is the user-given threshold on the success probability. The 
    default value of "threshold is zero."
    '''
    qbits = QuantumRegister(n,'q')
    circ = QuantumCircuit(n, name="Grover")
    
    # Initial state: Prepare equal superposition
    circ.h(qbits)
    
    # Number of iterations required
    N = 2**n    # Size of the database
    
    theta = np.arcsin(np.sqrt(m/N))%(2*np.pi)    # The value of theta
    
    ''' The following loop searches for the optimum value of "p" and the corresponding number of iterations
    such that the success probability of the search is equal to or more than the threshold.'''
    p=0    # Initialized to zero
    while True:
        
        k = (2*p+1)*np.pi/(4*theta)-0.5     # The optimal k is close to this value.
        tk_down = theta*(1+2*np.floor(k))    # Final angle theta_k with k rounded down.
        tk_up = theta*(1+2*np.ceil(k))    # Final angle theta_k with k rounded up.
        
        # Condition to check if the success probability crosses the threshold.
        if max(np.sin(tk_down)**2, np.sin(tk_up)**2) >= threshold:
            
            # If tk_down satisfies the condition choose the floor of k.
            if np.sin(tk_down)**2 >= threshold:
                k = int(np.floor(k))
                print(np.sin(tk_down)**2)    # Print the value of success probability.
                
            # If tk_up satisfies the condition choose the ceil of k.
            else:
                k = int(np.ceil(k))
                print(np.sin(tk_up)**2)    # Print the value of success probability.
                
            print('p =',p)    # Print the value of optimal p.
            break    # Break the loop if the optimal p is found. 
            
        p+=1
        
    print('k =',k)    # Print the value of optimal k.
    
    # Append the oracle followed by diffusion operation k number of times.
    for i in range(k):
        circ.append(oracle, qbits)
        circ.append(diffusion(n), qbits)
    
    return circ

'''
Each n-qubit computational basis state is represented by its corresponding n-digit binary string. For example,
the state |001101> is represented by the string '0001101'.
'''
n=    # Set the number of qubits.
form = f'{{0:0{str(n)}b}}'    # Format to convert decimal number to n-digit binary string.
all_states = [form.format(i) for i in range(2**n)]    # A list of all possible orthogonal states

tars = []    # "tars" is a list of binary strings for all the target states.
success_rate =      # Set the threshold for success probability.
orc = oracle(n, tars)    # The circuit for Oracle given the target states in the list tars.

# Construction of the circuit for Grover's search.
circuit = grover(n, orc, len(tars), success_rate)

# Applying projective measurement to each qubit at the end of the circuit.
circuit.measure_all()

# Running the "circuit" using the QASM simulator.
sim = Aer.get_backend('qasm_simulator')    # Calling the qasm_simulator.
job = execute(circuit, sim, shots=1000000)    # Executing the circuit with 1,000,000 shots.
result = job.result()
count = result.get_counts(circuit)    # "count" is a dict of counts after the projective measurements.

''' If the count of a state is zero in the output, it is not included in the dict.
The following loop includes the absent states in the dict and sets their counts to zero. '''
for x in all_states:
    if x not in count.keys():
        count[x]=0

final_counts = count

''' Uncomment the following lines, if you want the collective counts of the target and non-target states.'''
# collective_count = {'Others':0, 'Target':0}
# for x in all_states:
#     if x in tars:
#         collective_count['Target'] += count[x]
#     else:
#         collective_count['Others'] += count[x]
        
# final_counts = collective_count

plot_histogram(final_counts)    # Plot the probability histogram of the final_counts.