from qiskit import QuantumCircuit as qc
from qiskit import QuantumRegister as qr
from qiskit import transpile
from qiskit.providers import Backend
from qiskit.quantum_info import DensityMatrix as dm
from qiskit.result import Counts
from qiskit.visualization import plot_histogram
from qiskit_ibm_runtime import QiskitRuntimeService
from math import pi, sqrt
from heapq import nlargest

"""Number of qubits."""
N: int = 5

"""Set of m nonnegative integers to search for using Grover's algorithm (i.e. TARGETS in base 10)."""
SEARCH_VALUES: set[int] = { 11, 9, 0, 3 }

"""Amount of times to simulate the algorithm."""
SHOTS: int = 10000

"""Set of m N-qubit binary strings representing target state(s) (i.e. SEARCH_VALUES in base 2)."""
TARGETS: set[str] = { f"{s:0{N}b}" for s in SEARCH_VALUES }

"""N-qubit quantum register."""
QUBITS: qr = qr(N, "qubit")

"""Name of backend to run the algorithm on."""
BACKEND_NAME: str = "ibmq_qasm_simulator"

"""Backend to run the algorithm on."""
BACKEND: Backend = QiskitRuntimeService().backend(BACKEND_NAME)

def oracle(targets: set[str] = TARGETS, name: str = "Oracle") -> qc:
    """Mark target state(s) with negative phase.

    Args:
        targets (set[str]): N-qubit binary string(s) representing target state(s). Defaults to TARGETS.
        name (str, optional): Quantum circuit's name. Defaults to "Oracle".

    Returns:
        qc: Quantum circuit representation of oracle.
    """
    # Create N-qubit quantum circuit for oracle
    oracle = qc(QUBITS, name = name)

    for target in targets:
        # Reverse target state since Qiskit uses little-endian for qubit ordering
        target = target[::-1]
        
        # Flip zero qubits in target
        for i in range(N):
            if target[i] == "0":
                oracle.x(i)                   # Pauli-X gate

        # Simulate (N - 1)-control Z gate
        oracle.h(N - 1)                       # Hadamard gate
        oracle.mcx(list(range(N - 1)), N - 1) # (N - 1)-control Toffoli gate
        oracle.h(N - 1)                       # Hadamard gate

        # Flip back to original state
        for i in range(N):
            if target[i] == "0":
                oracle.x(i)                   # Pauli-X gate

    return oracle

# Generate and display quantum circuit for oracle
grover_oracle = oracle()
grover_oracle.draw("mpl", style = "iqp")

def diffuser(name: str = "Diffuser") -> qc:
    """Amplify target state(s) amplitude, which decreases the amplitudes of other states
    and increases the probability of getting the correct solution (i.e. target state(s)).

    Args:
        name (str, optional): Quantum circuit's name. Defaults to "Diffuser".

    Returns:
        qc: Quantum circuit representation of diffuser (i.e. Grover's diffusion operator).
    """
    # Create N-qubit quantum circuit for diffuser
    diffuser = qc(QUBITS, name = name)
    
    diffuser.h(QUBITS)                                    # Hadamard gate
    diffuser.append(oracle({ "0" * N }), list(range(N)))  # Oracle with all zero target state
    diffuser.h(QUBITS)                                    # Hadamard gate
    
    return diffuser

# Generate and display quantum circuit for diffuser
grover_diffuser = diffuser()
grover_diffuser.draw("mpl", style = "iqp")


def grover(oracle: qc = oracle(), diffuser: qc = diffuser(), name: str = "Grover Circuit") -> tuple[qc, dm]:
    """Create quantum circuit representation of Grover's algorithm,
    which consists of 4 parts: (1) state preparation/initialization,
    (2) oracle, (3) diffuser, and (4) measurement of resulting state.
    
    Steps 2-3 are repeated an optimal number of times (i.e. Grover's
    iterate) in order to maximize probability of success of Grover's algorithm.

    Args:
        oracle (qc, optional): Quantum circuit representation of oracle. Defaults to oracle().
        diffuser (qc, optional): Quantum circuit representation of diffuser. Defaults to diffuser().
        name (str, optional): Quantum circuit's name. Defaults to "Grover Circuit".

    Returns:
        tuple[qc, dm]: Quantum circuit representation of Grover's algorithm and its density matrix.
    """
    # Create N-qubit quantum circuit for Grover's algorithm
    grover = qc(QUBITS, name = name)

    # Intialize qubits with Hadamard gate (i.e. uniform superposition)
    grover.h(QUBITS)
    
    # Apply barrier to separate steps
    grover.barrier()
    
    # Apply oracle and diffuser (i.e. Grover operator) optimal number of times
    for _ in range(int((pi / 4) * sqrt((2 ** N) / len(TARGETS)))):
        grover.append(oracle, list(range(N)))
        grover.append(diffuser, list(range(N)))

    # Generate density matrix representation of Grover's algorithm
    density_matrix = dm(grover)

    # Measure all qubits once finished
    grover.measure_all()
    
    return grover, density_matrix

# Save density matrix and generate and display quantum circuit for Grover's algorithm
grover_circuit, density_matrix = grover(grover_oracle, grover_diffuser)
grover_circuit.draw("mpl", style = "iqp")

def outcome(winners: list[str], counts: Counts) -> None:
    """Print top measurement(s) (state(s) with highest frequency)
    and target state(s) in binary and decimal form, determine
    if top measurement(s) equals target state(s), then print result.

    Args:
        winners (list[str]): State(s) (N-qubit binary string(s))
        with highest probability of being measured.
        counts (Counts): Each state and its respective frequency.
    """
    print("WINNER(S):")
    print(f"Binary = {winners}\nDecimal = {[ int(key, 2) for key in winners ]}\n")
        
    print("TARGET(S):")
    print(f"Binary = {TARGETS}\nDecimal = {SEARCH_VALUES}\n")

    if not all(key in TARGETS for key in winners): print ("Target(s) not found...")

    else:
        winners_frequency, total = 0, 0

        for value, frequency in counts.items():
            if value in winners:
                winners_frequency += frequency
            total += frequency
        
        print(f"Target(s) found with {winners_frequency / total:.2%} accuracy!")

# Simulate Grover's algorithm with Qiskit backend and get results
with BACKEND.open_session() as session:
    results = BACKEND.run(transpile(grover_circuit, BACKEND, optimization_level = 2), shots = SHOTS).result()

# Get each state's frequency
counts = results.get_counts()

# Find winner(s) (i.e. state(s) with highest count)
winners = nlargest(len(TARGETS), counts, key = counts.get)

# Print outcome
outcome(winners, counts)

# Display simulation results as histogram
plot_histogram(data = counts, number_to_keep = len(TARGETS))

# City plot of density_matrix
density_matrix.draw("city")

# Bloch sphere representation of density_matrix
# reverse_bits = True since Qiskit uses little-endian for qubit ordering
density_matrix.draw("bloch", reverse_bits = True)

# Hinton plot of density_matrix
density_matrix.draw("hinton")

# Qsphere representation of density_matrix
density_matrix.draw("qsphere")