import numpy as np
import numpy.fft
import matplotlib.pyplot as plt
from qiskit.aqua.circuits import StateVectorCircuit
from qiskit.circuit.library import QFT, Barrier
from qiskit import QuantumCircuit, QuantumRegister, execute, Aer

def normalize(ψ):
    n = np.linalg.norm(ψ)
    return ψ if n == 0 else ψ / n

def wavefunction(circuit):
    backend = Aer.get_backend('statevector_simulator')
    return execute(circuit.copy(), backend).result().get_statevector()

def exercise1():
    
    l = 4 # Number of qubits
    N = 2**l
    L = 8
    Δx = L/N
    x = -L/2 + Δx * np.arange(N)
    #
    # Discretize our function, and normalize it to turn it into a
    # quantum state
    f = normalize(np.exp(-x**2/2.0))
    #
    # Use Qiskit to build a quantum circuit that constructs the
    # quantum state associated to the discretize function 'f'
    f_qc = StateVectorCircuit(f).construct_circuit()
    
    fig, ax = plt.subplots(nrows=1,ncols=1,figsize=(8,4))
    ax.plot(x, f, '-')
    ax.plot(x, np.abs(wavefunction(f_qc)), '.')
    ax.set_title(f'$f$ with $N={N}$')
    
exercise1()

from qiskit.circuit.library import QFT

QFT(3, insert_barriers=True).draw()

def fQFT(nqubits, **kwdargs):
    circuit = QuantumCircuit(nqubits)
    return circuit.compose(QFT(nqubits, **kwdargs), qubits=list(reversed(range(0,nqubits))))

fQFT(3).draw()

fQFT(3, insert_barriers=True, do_swaps=True, inverse=True).draw()

def exercise3():
    l = 4 # Number of qubits
    N = 2**l
    L = 8
    Δx = L/N
    x = -L/2 + Δx * np.arange(N)
    #
    # The original function
    f = normalize(np.exp(-x**2/2.0))
    #
    # Construct the circuit for the function 'f'
    f_qc = StateVectorCircuit(f).construct_circuit()
    #
    # Construct the circuit for the Quantum Fourier Transform
    qft_circuit = QFT(l, do_swaps=True)
    #
    # The QFT acts on the qubits of the circuit in the opposite
    # order than we would like. We explicit this order when
    # we compose the circuits
    qft_order = list(reversed(range(0, l)))
    #
    Ff = f_qc.compose(qft_circuit, qubits=qft_order)
    #
    # We can also compute the classical fourier transform using
    # Python's numpy library
    DFTf = np.fft.fft(f)
    
    fig, ax = plt.subplots(nrows=1,ncols=1,figsize=(8,4))
    ax.plot(x, f, '-', label='f')
    ax.plot(x, normalize(np.abs(DFTf)), label='$\\tilde{f}$')
    ax.plot(x, np.abs(wavefunction(Ff)), '.', label='$QFT$')
    ax.set_title(f'$f$ with $N={N}$')
    ax.legend()
    
exercise3()

aux=QuantumCircuit(4)
aux.cx(1,3)
aux.cx(1,2)
aux.draw()

steps = QuantumCircuit(4)
steps.compose(fQFT(3, insert_barriers=True), qubits=[0,1,2], inplace=True)
steps.cx(2,3)
steps.compose(fQFT(4, insert_barriers=True), inplace=True)
steps.draw()

def exercise6():
    l = 4 # Number of input qubits
    N = 2**l
    L = 8
    Δx = L/N
    x = -L/2 + Δx * np.arange(N)
    
    lfinal = l+2 # Number of qubits for interpolation
    xfinal = -L/2 + (Δx/2**(lfinal-l)) * np.arange(2**lfinal)
    #
    # Our circuit (with the final number of qubits)
    steps = QuantumCircuit(lfinal)
    #
    # The original function
    f = normalize(np.exp(-x**2/2.0))
    #
    # The interpolated one
    ffinal = normalize(np.exp(-xfinal**2/2.0))
    #
    # Construct the circuit for the function 'f'
    f_qc = StateVectorCircuit(f).construct_circuit()
    steps.compose(f_qc, qubits=range(0,l), inplace=True)
    #
    # Construct the circuit for the Quantum Fourier Transform
    steps.compose(fQFT(l, insert_barriers=True),
                  qubits=list(range(0, l)), inplace=True)
    #
    # Implement padding
    for i in range(l, lfinal):
        steps.cx(l-1,i)
    #
    # Implement the inverse transform
    steps.compose(fQFT(lfinal, inverse=True, insert_barriers=True),
                  inplace=True)
    #print(steps.draw())
    
    fig, ax = plt.subplots(nrows=1,ncols=1,figsize=(8,4))
    ax.plot(xfinal, ffinal, '-', label='exact')
    ax.plot(xfinal, np.abs(wavefunction(steps)), '.', label='interpolated')
    ax.set_title(f'Interpolate $f$ from {l} to {lfinal} qubits')
    ax.legend()
    
exercise6()