from qiskit.quantum_info.operators import Operator
from qiskit.circuit.library import CU1Gate


#helper functions

def make_permutation_matrix(n, permutation):
    r = np.zeros((n,n), dtype=int)
    for i in range(n):
        r[permutation(i), i] = 1
    return r

def mult_mat(x,k,N):
    n = math.ceil(math.log(N, 2))
    return make_permutation_matrix(
    2**n,
    permutation=lambda y:  y*pow(x,k) % N if y<N else y)

def c_mult_mat(x,k,N):
    n = math.ceil(math.log(N, 2))
    permutation = lambda y: y if y%2==0 or (y-1)/2 >= N else 2*(int((y-1)/2)*pow(x,k) % N) + 1
    return make_permutation_matrix(2*(2**n), permutation )

def mult_op(x,k,N):
    return Operator(mult_mat(x,k,N))



#controlled-U oracle 

def c_mult_op(x,k,N):
    return Operator(c_mult_mat(x,k,N))


# QFT and IQFT
def qft(circ, q, n):
    """n-qubit QFT on q in circ."""
    for j in range(n):
        circ.h(q[j])
        for k in range(j+1,n):
            circ.cu1(math.pi/float(2**(k-j)), q[k], q[j])
    for i in range(n//2):
        circ.swap(q[i], q[n-i-1])
        
def iqft(circ,q,n):
    for j in range(n):
        circ.h(q[j])
        for k in range(j+1,n):
            #circ.cu1(-math.pi/float(2**(k-j)), q[k], q[j])
             gate = CU1Gate(-np.pi/float(2**(k-j)) )
             circ.append(gate, [q[k],q[j]])
    for i in range(n//2):
        circ.swap(q[i], q[n-i-1])

