import numpy as np
import pylab as plt

from qiskit import QuantumCircuit, execute, Aer, IBMQ
from qiskit.visualization import plot_histogram, plot_bloch_multivector


qc = QuantumCircuit(3)

qc.h(2)
qc.draw('mpl')

# UROT_2 gate to x1 depending on x2

from numpy import pi
qc.cu1(pi/2, 1, 2) # CROT from qubit 1 to qubit 2 hence the angle: pi/2^{2-1}
qc.draw('mpl')

qc.cu1(pi/4, 0, 2) # CROT from qubit 2 to qubit 0 hence the angle: pi/2^{2-0}
qc.draw('mpl')

# Repeat the process for 2 and 3 
qc.h(1)  # Hadamard on 1
qc.cu1(pi/2, 0, 1) # CROT from qubit 0 to qubit 1 hence the angle: pi/2^{1-0}
qc.h(0)  # Hadamard on 0
qc.draw('mpl')

# Now swap the qubit 0 and 2 to complete QFT. [NOT CLEAR TO ME]
qc.swap(0,2)
qc.draw('mpl')

def qft_rotations(ckt, n):
    if n == 0:
        return ckt
    n -= 1
    ckt.h(n)
    for qubit in range(n):
        ckt.cu1( pi/2**(n-qubit), qubit, n )

qc = QuantumCircuit(4)
qft_rotations(qc,4)
qc.draw('mpl')

def qft_rotations(ckt, n):
    if n == 0:
        return ckt
    n -= 1
    ckt.h(n)
    for qubit in range(n):
        ckt.cu1( pi/2**(n-qubit), qubit, n )
        
    qft_rotations(ckt, n)
        
qc = QuantumCircuit(4)
qft_rotations(qc,4)
qc.draw('mpl')

from qiskit_textbook.widgets import scalable_circuit
scalable_circuit(qft_rotations)

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

def qft(ckt, n):
    qft_rotations(ckt, n)
    swap_registers(ckt, n)
    return ckt

qc = QuantumCircuit(4)
qft(qc, 4)
qc.draw('mpl')

print (f"5 and it's binary: {bin(5)}")
#bin(9)

qc = QuantumCircuit(3)
qc.x(0)
qc.x(2)
qc.draw('mpl')

backend = Aer.get_backend('statevector_simulator')
sv = execute(qc, backend=backend).result().get_statevector()
plot_bloch_multivector(sv)

print (sv.real)

qft(qc, 3)
qc.draw('mpl')

sv = execute(qc, backend=backend).result().get_statevector()
plot_bloch_multivector(sv)


print ('The statevectors are:\n', sv.round(1))
print ('Angle subtended by each states (as a factor of pi):\n',(np.angle(sv)/np.pi).round(2))


# 10
qc = QuantumCircuit(2)
qc.x(0)
qc.draw('mpl')

backend = Aer.get_backend('statevector_simulator')
sv = execute(qc, backend=backend).result().get_statevector()
plot_bloch_multivector(sv)

qft(qc, 2)
qc.draw('mpl')

backend = Aer.get_backend('statevector_simulator')
sv = execute(qc, backend=backend).result().get_statevector()
plot_bloch_multivector(sv)





def inverse_qft(circuit, n):
    """Does the inverse QFT on the first n qubits in circuit"""
    # First we create a QFT circuit of the correct size:
    qft_circ = qft(QuantumCircuit(n), n)
    # Then we take the inverse of this circuit
    invqft_circ = qft_circ.inverse()
    # And add it to the first n qubits in our existing circuit
    circuit.append(invqft_circ, circuit.qubits[:n])
    # use .decompose() which allows us to see the individual gates
    return circuit.decompose() 


nqubits = 3
number = 6
qc = QuantumCircuit(nqubits)
for qubit in range(nqubits):
    qc.h(qubit)
qc.u1(number*pi/4,0)
qc.u1(number*pi/2,1)
qc.u1(number*pi,2)

qc.draw('mpl')

backend = Aer.get_backend("statevector_simulator")
sv = execute(qc, backend=backend).result().get_statevector()
plot_bloch_multivector(sv)


print ('The statevectors are:\n', sv.round(1))
print ('Angle subtended by each states (as a factor of pi):\n',(np.angle(sv)/np.pi).round(2))


qc = inverse_qft(qc,nqubits)
qc.measure_all()
qc.draw('mpl')

backend = Aer.get_backend("qasm_simulator")

shots = 2048
job = execute(qc, backend=backend, shots=shots, optimization_level=3)

counts = job.result().get_counts()
plot_histogram(counts)

from qiskit import circuit, quantum_info

num_qubits = 6
qft = circuit.library.QFT(num_qubits, approximation_degree=0, insert_barriers=False)
qft.draw('mpl', fold=500, reverse_bits = True)


# Plot classical vs. quantum Fourier transform, for all basis states..
# By linearity of Fourier transform, all states map correctly.
# The difference in height is due to normalization in quantum.

plt.figure(figsize=(16,16))
for i, basis_state in enumerate([0, 1, 2, 3]):
    state = quantum_info.Statevector.from_int(basis_state, dims=2**num_qubits)
    #print (state)
    classical_fourier = np.fft.ifft(state.data)      # classical library function
    quantum_fourier = state.evolve(qft).data         # quantum library (simulation via quantum_info)

    width = 0.5
    
    #fig.suptitle('Input State: ' + str(basis_state), fontsize=16)
    
    plt.subplot(4,2, 2*i+1)
    plt.title('state: %1i real'%i)
    plt.bar(np.arange(2**num_qubits)-width/2., classical_fourier.real, width,  label='classical')
    plt.bar(np.arange(2**num_qubits)+width/2., quantum_fourier.real, width, label='quantum')
    plt.legend(frameon=False)
    
    plt.subplot(4,2, 2*i+2)
    plt.title('state: %1i Imaginary part'%i)
    plt.bar(np.arange(2**num_qubits)-width/2., classical_fourier.imag, width, label='classical')
    plt.bar(np.arange(2**num_qubits)+width/2., quantum_fourier.imag, width, label='quantum')
    plt.legend(frameon=False)



import matplotlib.pyplot as plt

Nt = 256
t = np.arange(Nt)
fft = np.fft.fft(np.sin(t))
ifft = np.fft.ifft(np.sin(t))
freq = np.fft.fftfreq(Nt)

plt.figure(figsize=(16,4))
plt.subplot(121)
plt.title('real')
plt.plot(freq, fft.real, '-o', label='fft (real)')
plt.plot(freq, Nt*ifft.real, '-o', label='Nt * ifft (real)')
plt.legend(fontsize='large')
plt.xlabel('frequency', fontsize='large')

plt.subplot(122)
plt.title('imag')
plt.plot(freq, fft.imag, '-o', label='fft (imag)')
plt.plot(freq, Nt*ifft.imag, '-o', label='Nt * ifft (imag)')
plt.legend(fontsize='large')
plt.xlabel('frequency', fontsize='large')

plt.show()

from qiskit import circuit, quantum_info

num_qubits = 6
qft = circuit.library.QFT(num_qubits, insert_barriers=False)
qft.draw('mpl', fold=500, reverse_bits = True)

basis_state = 2
state = quantum_info.Statevector.from_int(basis_state, dims=2**num_qubits)
quantum_fourier = state.evolve(qft).data

width = 0.5

plt.subplot(121)
plt.bar(np.arange(2**num_qubits)+width/2., quantum_fourier.real, width, label='real')

plt.subplot(122)
plt.bar(np.arange(2**num_qubits)+width/2., quantum_fourier.imag, width, label='imag')
plt.legend(frameon=False)


