import numpy as np
from qiskit import QuantumCircuit, Aer, execute
from qiskit.visualization import circuit_drawer
from qiskit.tools.visualization import plot_histogram

# Total number of inputs
n = 8

# Random n numbers
a = np.random.randint(0, 2**n, size=n)

# Check if n is a power of 2
if not (n and (not (n & (n - 1)))) or len(a) != n:
    raise ValueError("n must be a power of 2 and len(a) must be n")

# Initialize the circuit and the registers
circuit = QuantumCircuit(2*n, name="bitonic_sort")
a_reg = [i for i in range(n)]
c_reg = [i + n for i in range(n)]

# Apply Hadamard gates to all input qubits
circuit.h(a_reg)

# Apply the bitonic sort operations recursively
for k in range(1, n):
    for j in range(k, 0, -1):
        circuit.cswap(j + k - 1, j - 1, c_reg[k])
        for i in range(j - 1):
            circuit.cswap(j + k - 1, i, c_reg[k])

# Measure the qubits and run the circuit
circuit.measure_all()
circuit_drawer(circuit, output='mpl')

backend = Aer.get_backend('qasm_simulator')
shots = 1024
job = execute(circuit, backend=backend, shots=shots)
result = job.result()
counts = result.get_counts()
plot_histogram(counts, figsize=(100,30))

# Extract the sorted output from the measurement results
output = [int(key[::-1], 2) for key in sorted(counts.keys())]

print(f"Input array: {a}")
print(f"Sorted array: {output}")