import pennylane as qml
from pennylane import numpy as np
import matplotlib.pyplot as plt

def coefficients_to_values(coefficients):
    """Returns the value representation of a polynomial
    
    Args:
        coefficients (array[complex]): a 1-D array of complex 
            coefficients of a polynomial with 
            index i representing the i-th degree coefficient

    Returns: 
        array[complex]: the value representation of the 
            polynomial 
    """
    ##################
    # YOUR CODE HERE #
    ################## 
    
    # pass
    return np.fft.fft(coefficients)

A = [4, 3, 2, 1]
print(coefficients_to_values(A))


def values_to_coefficients(values):
    """Returns the coefficient representation of a polynomial
    
    Args:
        values (array[complex]): a 1-D complex array with 
            the value representation of a polynomial 

    Returns: 
        array[complex]: a 1-D complex array of coefficients
    """
    
    ##################
    # YOUR CODE HERE #
    ################## 
    return np.fft.ifft(values)


A = [10.+0.j,  2.-2.j,  2.+0.j,  2.+2.j]
print(values_to_coefficients(A))


def nearest_power_of_2(x):
    """Given an integer, return the nearest power of 2. 
    
    Args:
        x (int): a positive integer

    Returns: 
        int: the nearest power of 2 of x
    """
    ##################
    # YOUR CODE HERE #
    ################## 
    # k = np.log2(x)
    
    k = np.ceil(np.log2(x))
    return (2**(int(k)))
nearest_power_of_2(7)

def fft_multiplication(poly_a, poly_b):
    """Returns the result of multiplying two polynomials
    
    Args:
        poly_a (array[complex]): 1-D array of coefficients 
        poly_b (array[complex]): 1-D array of coefficients 

    Returns: 
        array[complex]: complex coefficients of the product
            of the polynomials
    """
    ##################
    # YOUR CODE HERE #
    ################## 
    # len(poly_b) = b
    # len(poly_a) = a
    b = len(poly_b)
#     print(b)
    a = len(poly_a) 
#     print(a)
    b = b-1
    a= a-1
    # Calculate the number of values required
    nn = a + b + 1
    # Figure out the nearest power of 2
    n = nearest_power_of_2(nn)
    # print(n)
    # Pad zeros to the polynomial
    # if a<n:
    poly_a = np.pad(poly_a, (0,n-len(poly_a)))
    # if b<n:
    poly_b = np.pad(poly_b, (0,n-len(poly_b)))
    # Convert the polynomials to value representation 
    aa = coefficients_to_values(poly_a)
    bb = coefficients_to_values(poly_b)
#     print(a)
#     print(b)
    # Multiply
    m = aa*bb
    # Convert back to coefficient representation
    res = values_to_coefficients(m)
#     final = np.zeros((a+b))
#     final = []
#     for i in range(a+b):
#         final[i] = res[i]
    return res
A = np.array([1,1])
B = np.array([1,-1])
fft_multiplication(A, B)

dev = qml.device("default.qubit", wires=1)

@qml.qnode(dev)
def one_qubit_QFT(basis_id):
    """A circuit that computes the QFT on a single qubit. 
    
    Args:
        basis_id (int): An integer value identifying 
            the basis state to construct.
    
    Returns:
        array[complex]: The state of the qubit after applying QFT.
    """
    # Prepare the basis state |basis_id>
    bits = [int(x) for x in np.binary_repr(basis_id, width=dev.num_wires)]
    qml.BasisStatePreparation(bits, wires=[0])

    ##################
    # YOUR CODE HERE #
    ##################
    qml.Hadamard(wires=0)
    return qml.state()


n_bits = 2
dev = qml.device("default.qubit", wires=n_bits)

@qml.qnode(dev)
def two_qubit_QFT(basis_id):
    """A circuit that computes the QFT on two qubits using qml.QubitUnitary. 
    
    Args:
        basis_id (int): An integer value identifying the basis state to construct.
    
    Returns:
        array[complex]: The state of the qubits after the QFT operation.
    """
    
    # Prepare the basis state |basis_id>
    bits = [int(x) for x in np.binary_repr(basis_id, width=dev.num_wires)]
    qml.BasisStatePreparation(bits, wires=[0, 1])
    
    ##################
    # YOUR CODE HERE #
    ##################
    U = 0.5*(np.array([[1,1,1,1], [1, complex(0,1), -1, complex(0,-1)], [1, -1, 1, -1], [1, complex(0,-1), -1, complex(0,1)]]))
    qml.QubitUnitary(U, wires=[0,1])
    return qml.state()


dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev)
def decompose_two_qubit_QFT(basis_id):
    """A circuit that computes the QFT on two qubits using elementary gates.
    
    Args:
        basis_id (int): An integer value identifying the basis state to construct.
    
    Returns:
        array[complex]: The state of the qubits after the QFT operation.
    """
    # Prepare the basis state |basis_id>
    bits = [int(x) for x in np.binary_repr(basis_id, width=dev.num_wires)]
    qml.BasisStatePreparation(bits, wires=[0, 1])
    
    ##################
    # YOUR CODE HERE #
    ##################
    qml.Hadamard(wires=0)
    qml.ctrl(qml.S, control=1)(wires=0)
    qml.Hadamard(wires=1)
    qml.SWAP(wires=[0,1])
    # pass
    return qml.state()


dev = qml.device("default.qubit", wires=3)

@qml.qnode(dev)
def three_qubit_QFT(basis_id):
    """A circuit that computes the QFT on three qubits.
    
    Args:
        basis_id (int): An integer value identifying the basis state to construct.
        
    Returns:
        array[complex]: The state of the qubits after the QFT operation.
    """
    # Prepare the basis state |basis_id>
    bits = [int(x) for x in np.binary_repr(basis_id, width=dev.num_wires)]
    qml.BasisStatePreparation(bits, wires=[0, 1, 2])
    ##################
    # YOUR CODE HERE #
    ################## 
    for i in range(3):
        # j=i+1
        # print("test1")
        qml.Hadamard(wires=i)
        for jj in range(i+1,3,1):
            # print("test2")
#             k = (1/(2**(jj)))*(np.pi)*(complex(0,1))
#             U = np.array([[1,0], [0, np.exp(k)]])
#             qml.ctrl(qml.QubitUnitary(U, wires=i), control=jj)
            # Below lines of code give the same result as the above one
            k = (1/(2**(jj)))*(np.pi)                    
            qml.ctrl(qml.PhaseShift(k, wires=i), control = jj)
    qml.SWAP(wires=[0,2])
    return qml.state()
print(qml.draw(three_qubit_QFT)(3))

dev = qml.device('default.qubit', wires=4)

            
def swap_bits(n_qubits):
    """A circuit that reverses the order of qubits, i.e.,
    performs a SWAP such that [q1, q2, ..., qn] -> [qn, ... q2, q1].
    
    Args:
        n_qubits (int): An integer value identifying the number of qubits.
    """
    ##################
    # YOUR CODE HERE #
    ##################
    # if n%2==0:
    # j = 
    for i in range(int(n_qubits/2)):
        qml.SWAP(wires=[i,n_qubits-1-i])
    # pass

@qml.qnode(dev) 
def qft_node(basis_id, n_qubits):
    # Prepare the basis state |basis_id>
    bits = [int(x) for x in np.binary_repr(basis_id, width=n_qubits)]
    qml.BasisStatePreparation(bits, wires=range(n_qubits))
    # qft_rotations(n_qubits)
    swap_bits(n_qubits)
    return qml.state()


dev = qml.device('default.qubit', wires=4)
# @qml.qnode(dev) 
def qft_rotations(n_qubits):
    """A circuit performs the QFT rotations on the specified qubits.
    
    Args:
        n_qubits (int): An integer value identifying the number of qubits.
    """
    
    ##################
    # YOUR CODE HERE #
    ################## 
    n = n_qubits
    for i in range(n):
        # j=i+1
        # print("test1")
        qml.Hadamard(wires=i)
        for jj in range(i+1,n,1):
            # print("test2")
            qml.ControlledPhaseShift(np.pi/(2**(jj-i)), wires=[jj,i])
    pass

@qml.qnode(dev) 
def qft_node(basis_id, n_qubits):
    # Prepare the basis state |basis_id>
    bits = [int(x) for x in np.binary_repr(basis_id, width=n_qubits)]
    qml.BasisStatePreparation(bits, wires=range(n_qubits))
    qft_rotations(n_qubits)
    swap_bits(n_qubits)
    return qml.state()

# print(qml.draw(qft_rotations)(3))
print(qml.draw(qft_node)(3,3))


dev = qml.device('default.qubit', wires=4)

def qft_recursive_rotations(n_qubits, wire=0):
    """A circuit that performs the QFT rotations on the specified qubits
        recursively.
        
    Args:
        n_qubits (int): An integer value identifying the number of qubits.
        wire (int): An integer identifying the wire 
                    (or the qubit) to apply rotations on.
    """

    ##################
    # YOUR CODE HERE #
    ################## 
    i = wire
    if i == n_qubits-1:
        qml.Hadamard(wires=i)
#     for i in range(n_qubits):
#         qml.Hadamard(wires=i)
#         if n_qubits!=i:
    else:
        qml.Hadamard(wires=i)
        for j in range(i+1, n_qubits, 1):
            qml.ControlledPhaseShift(np.pi/(2**(j-i)), wires=[j,i])
        qft_recursive_rotations(n_qubits, i+1)
    pass

@qml.qnode(dev) 
def qft_node(basis_id, n_qubits):
    # Prepare the basis state |basis_id>
    bits = [int(x) for x in np.binary_repr(basis_id, width=n_qubits)]
    qml.BasisStatePreparation(bits, wires=range(n_qubits))
    qft_recursive_rotations(n_qubits)
    swap_bits(n_qubits)
    return qml.state()
print(qml.draw(qft_node)(3,3))