import numpy as np
import networkx as nx
from networkx import hypercube_graph
from networkx.drawing.nx_agraph import graphviz_layout
import matplotlib.pyplot as plt
from typing import Callable
from qiskit import QuantumCircuit, Aer, QuantumRegister, ClassicalRegister, transpile
from qiskit.quantum_info import Operator
from qiskit.tools.visualization import plot_histogram, plot_bloch_vector
from qiskit_aer import AerSimulator
from numpy import pi


coords = [1, np.pi / 2, 0]  # [Radius, Theta, Phi]
plot_bloch_vector(coords, coord_type="spherical")


qc = QuantumCircuit(2)
qc.h(0)
qc.x(1)
qc.x(0)
qc.y(1)
qc.h(0)
qc.draw("mpl", initial_state=True)


qc = QuantumCircuit(2)
qc.x(0)
qc.cx(0, 1)
qc.x(0)
qc.draw("mpl", initial_state=True)


pos = {1: [0.75, 1.0], 2: [0.75, 0.15], 3: [0.5, -0.5], 4: [1.0, -0.5]}
U = [
    [0.1, 0.2, 0.6, 0.1],
    [0.4, 0.3, 0.2, 0.1],
    [0.4, 0.2, 0.1, 0.3],
    [0.1, 0.0, 0.1, 0.8],
]

G = nx.DiGraph(directed=True)
for i in range(1, 5):
    G.add_node(i)
    for j in range(1, 5):
        G.add_edge(i, j, weight=U[i - 1][j - 1])

nx.draw(G, pos=pos, with_labels=True)
nx.draw_networkx_edge_labels(
    G, pos=pos, edge_labels=nx.get_edge_attributes(G, "weight")
)
plt.show()


n_steps = 50  # Number of steps the walker will take
pr = 0.5  # Probability of the walker stepping to the right
initial_position = 0  # Initial position of our walker


def random_walk(pr: float, n_steps: int, initial_pos: int) -> int:
    """Run a random walk with N steps and a probability of stepping to the right of pr.

    Args:
        pr: probability of stepping to the right
        N: number of steps
        i: initial position

    Returns:
        The final position of the walker
    """
    position = initial_pos

    for _ in range(n_steps):
        coin_flip = list(
            np.random.choice(2, 1, p=[1 - pr, pr])
        )  # Flips our weighted coin
        position += 2 * coin_flip[0] - 1  # Moves the walker according to the coin flip

    return position


print(f"The walker is located at: x = {random_walk(pr, n_steps, initial_position)}")


def dist(n_runs: int, n_steps: int):
    """Plot the distribution of the random walk.

    Args:
        runs: number of runs
        N: number of steps
    """
    positions = range(-1 * n_steps, n_steps + 1)
    instances = [0 for _ in range(-1 * n_steps, n_steps + 1)]

    for _ in range(n_runs):
        result = random_walk(pr, n_steps, initial_position)
        instances[positions.index(result)] += 1

    plt.bar(positions, [n / n_runs for n in instances])
    plt.show()


dist(10000, n_steps)


# Number of qubits
n = 3
# Initial state
initial_state = np.array([1] + [0] * (2**n - 1))
assert initial_state.size == 2**n

# Hadamard gate
hadamard_gate_base = 1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])
hadamard_gate = hadamard_gate_base
for i in range(n - 1):
    hadamard_gate = np.kron(hadamard_gate, hadamard_gate_base)

# Grover's gate
grover_gate = np.full((2**n, 2**n), 1 / 2 ** (n - 1)) - np.eye(2**n)

# Apply gates
hadamard_res = np.dot(initial_state, hadamard_gate)
hadamard_res = np.square(hadamard_res)
grover_res = np.dot(initial_state, grover_gate)
grover_res = np.square(grover_res)

# Plot results
fig, axes = plt.subplots(2)
plt.suptitle("Hadamard vs Grover gates (|000>)")
for ax, res in zip(axes, [hadamard_res, grover_res]):
    ax.set_ylabel("Probability")
    ax.set_ylim(0, 1)
    ax.bar(tuple(bin(i)[2:].zfill(n) for i in range(res.size)), res)
plt.show()


n_walker_qubits = 4  # The number of qubits used to represent the position of the walker


def coined_walk_circuit(
    n_walker_qubits: int,
) -> tuple[QuantumCircuit, QuantumRegister, QuantumRegister, ClassicalRegister]:
    """Create a quantum circuit for the quantum walk.

    Args:
        n_walker_qubits: number of qubits used to represent the position of the walker

    Returns:
        a tuple containing the circuit, the walker qubits, the coin qubits and the classical register to store the output
    """
    walker_r = QuantumRegister(n_walker_qubits, name="w")
    coin_r = QuantumRegister(1, name="c")  # The coin requires only one qubit
    classic_r = ClassicalRegister(n_walker_qubits, name="output")
    qc = QuantumCircuit(walker_r, coin_r, classic_r)
    return qc, walker_r, coin_r, classic_r


qc, _, _, _ = coined_walk_circuit(n_walker_qubits)
qc.draw(output="mpl", initial_state=True)


initial_position = 2 ** (
    n_walker_qubits - 1
)  # The initial position of the walker is the middle of all possible positions
initial_coin_value = 1  # The initial value of the coin is 0


def initialize_coined_walk_circuit(
    qc: QuantumCircuit,
    walker_r: QuantumRegister,
    coin_r: QuantumRegister,
    initial_position: int,
    initial_coin_value: int,
) -> QuantumCircuit:
    """Initialize the circuit with the initial position of the walker and the initial value of the coin.

    Args:
        qc: the quantum circuit
        walker_r: the quantum register containing the walker's position qubits
        coin_r: the quantum register containing the coin qubit
        initial_position: the initial position of the walker
        initial_coin_value: the initial value of the coin

    Returns:
        the initialized quantum circuit
    """
    if initial_coin_value == 1:
        qc.x(coin_r)
    for i in range(n_walker_qubits):
        if initial_position & (1 << i):
            qc.x(walker_r[n_walker_qubits - i - 1])
    qc.barrier()
    return qc


qc, walker_r, coin_r, _ = coined_walk_circuit(n_walker_qubits)
qc = initialize_coined_walk_circuit(
    qc, walker_r, coin_r, initial_position, initial_coin_value
)
qc.draw(output="mpl", initial_state=True)


def coined_walk_step(
    qc: QuantumCircuit, walker_r: QuantumRegister, coin_r: QuantumRegister
) -> QuantumCircuit:
    """Single step of the quantum walk.

    Args:
        qc: quantum circuit
        walker_r: the quantum register containing the walker's position qubits
        coin_r: the quantum register containing the coin qubit

    Returns:
        quantum circuit with an added walk step
    """

    # "Flip" the coin vector
    qc.h(coin_r)

    # Implement the Addition Operator
    for i in reversed(range(len(walker_r))):
        # Qubits with less significant bits than the current one
        controls = [walker_r[v] for v in range(len(walker_r) - 1, i, -1)]
        controls.append(coin_r)  # The coin qubit is also used as a control
        qc.mcx(controls, walker_r[i])  # Multi-controlled X gate
        if i != 0:
            qc.x(walker_r[i])  # Flip the qubit

    # Switch the coin vector
    qc.x(coin_r)

    # Implement the Subtraction Operator
    for i in range(len(walker_r)):
        if i != 0:
            qc.x(walker_r[i])  # Reverse the flip
        # Qubits with less significant bits than the current one
        controls = [walker_r[v] for v in range(len(walker_r) - 1, i, -1)]
        controls.append(coin_r)  # The coin qubit is also used as a control
        qc.mcx(controls, walker_r[i])  # Multi-controlled X gate

    # Revert the switch of the coin vector
    qc.x(coin_r)

    return qc


qc, walker_r, coin_r, _ = coined_walk_circuit(n_walker_qubits)
qc = coined_walk_step(qc, walker_r, coin_r)
qc.draw(output="mpl", initial_state=True)


n_steps = 5  # The number of steps of the quantum walk


def coined_walk(
    n_walker_qubits: int,
    n_steps: int,
    initial_position: int,
    initial_coin_value: int,
    init_func: Callable,
) -> QuantumCircuit:
    """Create a quantum circuit for the quantum walk.

    Args:
        n_walker_qubits: number of qubits used to represent the position of the walker
        n_steps: number of steps of the quantum walk
        initial_position: initial position of the walker
        initial_coin_value: initial value of the coin
        init_func: function used to initialize the circuit

    Returns:
        quantum circuit for the quantum walk
    """
    qc, walker_r, coin_r, classic_r = coined_walk_circuit(n_walker_qubits)
    qc = init_func(qc, walker_r, coin_r, initial_position, initial_coin_value)
    for i in range(n_steps):
        step_gate = coined_walk_step(QuantumCircuit(walker_r, coin_r), walker_r, coin_r)
        qc.append(step_gate.to_gate(label=f"Step #{i + 1}"), walker_r[:] + coin_r[:])
    qc.barrier()
    qc.measure(walker_r, reversed(classic_r))
    return qc


qc = coined_walk(
    n_walker_qubits,
    n_steps,
    initial_position,
    initial_coin_value,
    initialize_coined_walk_circuit,
)
qc.draw(output="mpl", initial_state=True)


n_walker_qubits = 6  # The number of qubits used to represent the position of the walker
n_steps = 30  # The number of steps of the quantum walk
initial_position = 2 ** (
    n_walker_qubits - 1
)  # The initial position of the walker is the middle of all possible positions
initial_coin_value = 0  # The initial value of the coin is 0
n_runs = 5000  # The number of runs of the quantum circuit

qc = coined_walk(
    n_walker_qubits,
    n_steps,
    initial_position,
    initial_coin_value,
    initialize_coined_walk_circuit,
)


def plot_results(qc: QuantumCircuit, n_runs: int, title: str):
    """Plot the results of the quantum walk.

    Args:
        qc: quantum circuit for the quantum walk
        n_runs: number of runs of the quantum circuit

    Returns:
        histogram of the results of the quantum walk
    """
    simulator: AerSimulator = Aer.get_backend("aer_simulator")
    transpiled_qc = transpile(qc, simulator)
    results = simulator.run(transpiled_qc, shots=n_runs).result()
    counts = results.get_counts(transpiled_qc)
    counts = {int(k, 2): v for k, v in counts.items()}
    return plot_histogram(counts, title=title)


plot_results(qc, n_runs, "Quantum Coined Walk")


def initialize_balance_coined_walk_circuit(
    qc: QuantumCircuit,
    walker_r: QuantumRegister,
    coin_r: QuantumRegister,
    initial_position: int,
    initial_coin_value: int,
) -> QuantumCircuit:
    """Initialize the circuit with the initial position of the walker and the initial value of the coin.
    Set the coin qubit to |+> and apply the S gate to it to make sure that the coin operator keeps being balanced.

    Args:
        qc: the quantum circuit
        walker_r: the quantum register containing the walker's position qubits
        coin_r: the quantum register containing the coin qubit
        initial_position: the initial position of the walker
        initial_coin_value: the initial value of the coin

    Returns:
        the initialized quantum circuit
    """
    if initial_coin_value == 1:
        qc.x(coin_r)
    qc.h(coin_r)
    qc.s(coin_r)
    for i in range(n_walker_qubits):
        if initial_position & (1 << i):
            qc.x(walker_r[n_walker_qubits - i - 1])
    qc.barrier()
    return qc


qc, walker_r, coin_r, _ = coined_walk_circuit(n_walker_qubits)
qc = initialize_balance_coined_walk_circuit(
    qc, walker_r, coin_r, initial_position, initial_coin_value
)
qc.draw(output="mpl", initial_state=True)


n_walker_qubits = 6  # The number of qubits used to represent the position of the walker
n_steps = 30  # The number of steps of the quantum walk
initial_position = 2 ** (
    n_walker_qubits - 1
)  # The initial position of the walker is the middle of all possible positions
initial_coin_value = 0  # The initial value of the coin is 0
n_runs = 5000  # The number of runs of the quantum circuit

qc = coined_walk(
    n_walker_qubits,
    n_steps,
    initial_position,
    initial_coin_value,
    initialize_balance_coined_walk_circuit,
)

plot_results(qc, n_runs, "Balanced Quantum Coined Walk")


n = 3  # Number of qubits

G = hypercube_graph(n)  # Create hypercube graph
# Draw graph
nx.draw(
    G,
    with_labels=True,
    pos=graphviz_layout(G, prog="dot"),
    node_color="#dddddd",
    node_size=2500,
)
plt.show()


n = 2  # Dimension of the hypercube


def graph_walk_circuit(
    n: int,
) -> tuple[QuantumCircuit, QuantumRegister, QuantumRegister, ClassicalRegister]:
    """Create a quantum circuit for the graph walk.

    Args:
        n: number of qubits used to represent the position of the walker

    Returns:
        quantum circuit for the graph walk, walker register, coin register, and output register
    """
    walker_r = QuantumRegister(2**n, name="w")
    coin_r = QuantumRegister(n, name="c")
    classic_r = ClassicalRegister(2**n, name="output")
    qc = QuantumCircuit(walker_r, coin_r, classic_r)
    return qc, walker_r, coin_r, classic_r


qc, _, _, _ = graph_walk_circuit(n)
qc.draw(output="mpl", initial_state=True)


n = 2  # Dimension of the hypercube

def grover_coin(coin_r: QuantumRegister) -> Operator:
    matrix_size = 2 ** len(coin_r)
    grover_matrix = np.full((matrix_size, matrix_size), 2 / matrix_size) - np.eye(
        matrix_size
    )
    return Operator(grover_matrix)


coin_r = QuantumRegister(n, name="c")
qc = QuantumCircuit(coin_r)
gc = grover_coin(coin_r)
qc.unitary(gc, coin_r[:], label="Grover Coin")
qc.draw(output="mpl", initial_state=True)


n = 2  # Dimension of the hypercube


def shift_operator(
    walker_r: QuantumRegister, coin_r: QuantumRegister
) -> QuantumCircuit:
    """Create a quantum circuit for the shift operator.

    Args:
        walker_r: the quantum register containing the walker's position qubits
        coin_r: the quantum register containing the coin qubit

    Returns:
        quantum circuit for the shift operator
    """
    qc = QuantumCircuit(walker_r, coin_r)
    for i in reversed(range(len(walker_r))):
        qc.mcx(coin_r, walker_r[i])
        qc.x(coin_r[-1])
        for j in range(1, len(coin_r)):
            if i & ((1 << j) - 1) == 0:
                qc.x(coin_r[-(j + 1)])
    return qc


qc, walker_r, coin_r, _ = graph_walk_circuit(n)
qc = shift_operator(walker_r, coin_r)
qc.draw(output="mpl", initial_state=True)


n = 2  # Dimension of the hypercube


def graph_walk_step(
    walker_r: QuantumRegister, coin_r: QuantumRegister
) -> QuantumCircuit:
    """Create a quantum circuit for one step of the graph walk.

    Args:
        walker_r: the quantum register containing the walker's position qubits
        coin_r: the quantum register containing the coin qubit

    Returns:
        quantum circuit for one step of the graph walk
    """
    shift = shift_operator(walker_r, coin_r)
    gc = grover_coin(coin_r)
    walk_step = QuantumCircuit(walker_r, coin_r)
    walk_step.unitary(gc, coin_r, label="Grover Coin")
    walk_step.compose(shift, inplace=True)
    return walk_step


walk_step = graph_walk_step(walker_r, coin_r)
walk_step.draw(output="mpl", initial_state=True)


n = 2  # Dimension of the hypercube
n_steps = 3  # Number of steps of the quantum walk


def graph_walk(
    n: int,
    n_steps: int,
) -> QuantumCircuit:
    """Create a quantum circuit for the quantum walk over a graph.

    Args:
        n: dimension of the hypercube
        n_steps: number of steps of the quantum walk
        initial_position: initial position of the walker
        initial_coin_value: initial value of the coin

    Returns:
        quantum circuit for the quantum walk over a graph
    """
    qc, walker_r, coin_r, classic_r = graph_walk_circuit(n)
    for i in range(n_steps):
        step_gate = graph_walk_step(walker_r, coin_r).to_gate(label=f"Step #{i + 1}")
        qc.append(step_gate, walker_r[:] + coin_r[:])
    qc.barrier()
    qc.measure(walker_r, reversed(classic_r))
    return qc


qc = graph_walk(n, n_steps)
qc.draw(output="mpl", initial_state=True)


n = 2  # Dimension of the hypercube
n_steps = 30  # Number of steps of the quantum walk
n_runs = 5000  # Number of runs of the quantum circuit

qc = graph_walk(n, n_steps)
plot_results(qc, n_runs, "Quantum Walk on a Hypercube")
