import pennylane as qml
from pennylane import numpy as np

# Basis State preparation 
def state_preparation(basis_id, n_qubits):
    bits = [int(x) for x in np.binary_repr(basis_id, width=n_qubits)]
    return qml.BasisStatePreparation(bits, wires=range(n_qubits))

#Testing
state_preparation(3,5)

# Swappping of quantum circuits 
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.
    """
    for i in range(int(n_qubits/2)):
        qml.SWAP(wires=[i,n_qubits-1-i])

# Conditional Rotation matrix
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.
    """
    n = n_qubits
    for i in range(n):
        qml.Hadamard(wires=i)
        for jj in range(i+1,n,1):
            qml.ControlledPhaseShift(np.pi/(2**(jj-i)), wires=[jj,i])
    pass

def qft_node(basis_id, basis_id2, n_qubits):
    ''' A circuit performs addition of two numbers 
    
    Args:
        basis_id (int): An integer value specifying the first number to be added.
        basis_id2 (int): An integer value specifying the second number to be added.
        n_qubits (int): An integer value identifying the number of 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))
    
    bits2 = [int(x) for x in np.binary_repr(basis_id2, width=n_qubits)]
    qml.BasisStatePreparation(bits2, wires=range(n_qubits, 2*(n_qubits),1))
    qft_rotations(n_qubits)
    for i in range(0,n_qubits,1):
        k=0
        for j in range(i+n_qubits,(2*n_qubits),1):
            qml.ControlledPhaseShift((np.pi/(2**(k))), wires=[j,i])
            k+=1
    (qml.adjoint)(qft_rotations)(n_qubits)
    return qml.sample(wires=[x for x in range(n_qubits)])

# basis_id stores a, basis_id2 stores b
basis_id = int(input("Enter number a: "))
basis_id2 = int(input("Enter number b: "))
if len(np.binary_repr(basis_id))>len(np.binary_repr(basis_id2)):
    n_qubits = len(np.binary_repr(basis_id)) + 1
else:
    n_qubits = len(np.binary_repr(basis_id2)) + 1
dev = qml.device('default.qubit', wires=(2*n_qubits), shots=10)
qnode = qml.QNode(qft_node, dev)

k = qnode(basis_id, basis_id2, n_qubits)
poww, summ = 0, 0
for l in reversed(k[0]):
    summ = summ + (2**poww)*l
    poww+=1
print(summ)

def multiplier(basis_id, basis_id2, n_qubits):
    ''' A circuit performs multiplication of two numbers 
    
    Args:
        basis_id (int): An integer value specifying the first number to be multiplied.
        basis_id2 (int): An integer value specifying the second number to be multiplied.
        n_qubits (int): An integer value identifying the number of qubits.
    '''
    
    # Basis state preparation for 0, a, b respectively
    
    bits1 = [int(x) for x in np.binary_repr(0, width=2*n_qubits)]
    qml.BasisStatePreparation(bits1, wires=range(2*(n_qubits)))
    
    bits2 = [int(x) for x in np.binary_repr(basis_id, width=n_qubits)]
    qml.BasisStatePreparation(bits2, wires=range(2*n_qubits, 3*n_qubits))
    
    bits3 = [int(x) for x in np.binary_repr(basis_id2, width=n_qubits)]
    qml.BasisStatePreparation(bits3, wires=range(3*n_qubits, 4*(n_qubits),1))
       
    # Applying QFT on first 2*n_qubits as they store the basis state for |0>
    qft_rotations(2*n_qubits)
    
    # Applying Conditional Rotation operator with controls a_j and b_m and target as |0>
    for i in range(3*n_qubits, 4*n_qubits,1):
        for j in range(2*n_qubits, 3*n_qubits,1):
            for m in range(2*n_qubits):
                qml.ctrl(qml.PhaseShift(np.pi/(2**(-(5*n_qubits)-m+j+i+1)), wires=m),control=[j,i])
    
    # Applying inverse QFT 2*n_qubits
    (qml.adjoint)(qft_rotations)(2*n_qubits)
    
    return qml.sample(wires=[x for x in range(2*n_qubits)])

# basis_id stores a, basis_id2 stores b
basis_id = int(input("Enter number a: "))
basis_id2 = int(input("Enter number b: "))

# n_qubits takes the value of the number whose binary representation is the greatest among the two
if len(np.binary_repr(basis_id))>len(np.binary_repr(basis_id2)):
    n_qubits = len(np.binary_repr(basis_id)) 
else:
    n_qubits = len(np.binary_repr(basis_id2))

# Creating QNode
dev = qml.device('default.qubit', wires=(4*n_qubits), shots=10)
qnode = qml.QNode(multiplier, dev)

# Below code returns the sample containing measurements on the wires
# print(qnode(basis_id,basis_id2,n_qubits))  
# Below prints quantum circuit
# print(qml.draw(qnode)(basis_id,basis_id2,n_qubits))

# k stores the sample of measurement values on each wire
k = qnode(basis_id, basis_id2, n_qubits)

# Converting k[0] to integer value
# Considering only k[0] as all elements of the sample are identical
poww, summ = 0, 0
for l in reversed(k[0]):
    summ = summ + (2**poww)*l
    poww+=1
print(summ)