import numpy as np
from numpy import pi
# importing Qiskit
from qiskit import QuantumCircuit, transpile, assemble, Aer, IBMQ
from qiskit.providers.ibmq import least_busy
from qiskit.tools.monitor import job_monitor
from qiskit.visualization import plot_histogram, plot_bloch_multivector

# Qiskit's way of implementing QFT:
def qft_rotations(circuit, n):
    """Performs qft on the first n qubits in circuit (without swaps)"""
    if n == 0:
        return circuit
    n -= 1
    circuit.h(n)
    for qubit in range(n):
        circuit.cp(pi/2**(n-qubit), qubit, n)
    # At the end of our function, we call the same function again on
    # the next qubits (we reduced n by one earlier in the function)
    qft_rotations(circuit, n)

def swap_registers(circuit, n):
    for qubit in range(n//2):
        circuit.swap(qubit, n-qubit-1)
    return circuit

def qft(circuit, n):
    """QFT on the first n qubits in circuit"""
    qft_rotations(circuit, n)
    swap_registers(circuit, n)
    return circuit

# See network archetechture for 4 qubits. There's a clear pattern with the Hadamard gates and CROT gates.
qc = QuantumCircuit(4)
qc = qft(qc,4)
qc.draw()

# My way, which doesn't use recursion:
def qft_new(circuit, n):
    for qubit in range(n):
        idx = n-qubit-1
        circuit.h(idx)
        for i in range(idx):
            circuit.cp(pi/2**(idx-i), i, idx)
    for j in range(n//2):
        circuit.swap(j, n-j-1)
    return circuit

# Outputs the same circuit for 4 qubits as the original way.
qc = QuantumCircuit(4)
qc = qft_new(qc,4)
qc.draw()

# Inverse QFT can be be found by simply using the .inverse() method:
qft_circ = qft_new(QuantumCircuit(4), 4)
inv = qft_circ.inverse()
inv.draw()

# The following example was given in the textbook:

qc = QuantumCircuit(3)
# Encode the state |1,0,1>
qc.x(0)
qc.x(2)
sim = Aer.get_backend("aer_simulator")
qc_init = qc.copy()
qc_init.save_statevector()
statevector = sim.run(qc_init).result().get_statevector()
plot_bloch_multivector(statevector)

# Perform QFT, then see the transformed states
qft_new(qc,3)
qc.save_statevector()
statevector = sim.run(qc).result().get_statevector()
plot_bloch_multivector(statevector)

# New example
qc_new = QuantumCircuit(2)
# Encode a custom entangled state
qc_new.initialize([0,1/(2**0.5),0,-1/(2**0.5)])
qc_init = qc_new.copy()
qc_init.save_statevector()
statevector = sim.run(qc_init).result().get_statevector()
plot_bloch_multivector(statevector)
# Note: Bloch sphere images might not provide any insight into encoded state since the two qubits are entangled

# Perform QFT
qc_transformed = qft_new(qc_new,2)
qc_final = qc_transformed.copy()
qc_final.save_statevector()
statevector = sim.run(qc_final).result().get_statevector()
plot_bloch_multivector(statevector)
# Note: once again, Bloch sphere images might not provide any insight