from mpqp.gates import *
from mpqp.execution.result import Result
from mpqp import QCircuit, Barrier
from mpqp.execution import run, IBMDevice
from math import floor
import numpy as np

class QFT(QCircuit):
      
    def __init__(self,n_qubits,inverse=False):
        
        super().__init__(n_qubits, nb_cbits=n_qubits)
        self.inverse = inverse
        self._build()
    
    def _build(self):
        for j in range(self.nb_qubits):
            self.add(H(j))
            self.add([CRk(i+1, i, j) for i in range(j+1, self.nb_qubits)])
            self.add(Barrier())
        self.add([SWAP(i, self.nb_qubits-1-i) for i in range(int(floor(self.nb_qubits / 2)))])
        
        if self.inverse == True:
           self.inverse()

qft_5 = QFT(n_qubits=5)
print(qft_5)

result = run(qft_5, IBMDevice.AER_SIMULATOR_STATEVECTOR)
if isinstance(result, Result):
    print(result.amplitudes)

qc1 = QCircuit([H(1)]) + QFT(2)

print(qc1)

result1 = run(qc1, IBMDevice.AER_SIMULATOR_STATEVECTOR) 
if isinstance(result1, Result):
    print(result1.amplitudes)

W_circuit = QCircuit([Ry(2*np.arccos(1/np.sqrt(3)),0),Ry(-np.pi/4,1),CZ(0,1),Ry(np.pi/4,1),CNOT(1,2),CNOT(0,1),X(0)])

qc2 = W_circuit + QFT(3)
print(qc2)

result2 = run(qc2, IBMDevice.AER_SIMULATOR_STATEVECTOR)
if isinstance(result2, Result):
    print(result2.amplitudes)