from qiskit import BasicAer, execute
from qiskit.circuit import QuantumRegister, ClassicalRegister, QuantumCircuit

from qiskit.circuit.library.standard_gates import PhaseGate
from qiskit.circuit.library.basis_change import QFT
import math
from src.util.util import run_qc

q = QuantumRegister(3)
b = QuantumRegister(1)
c = ClassicalRegister(3)
qc = QuantumCircuit(q, b, c)

qc.x(q[0])

qc.draw(output='mpl')

def increment(circuit, register, apply_QFT=True):
    q = register
    num = len(q)
    qc = circuit
    
    if apply_QFT == True:
        qc.barrier()
        qc = qc.compose(QFT(num_qubits=num, approximation_degree=0, do_swaps=True, \
                           inverse=False, insert_barriers=True, name='qft'))
        qc.barrier()
    
    for i, qubit in enumerate(q):
        qc.rz(math.pi/2**(num-1-i), qubit)
    
    if apply_QFT == True:
        qc.barrier()
        qc = qc.compose(QFT(num_qubits=num, approximation_degree=0, do_swaps=True, \
                          inverse=True, insert_barriers=True, name='iqft'))
        qc.barrier()
    
    return qc
    

test = increment(qc, q)
test.draw(output='mpl')

def control_increment(circuit, qregister, cregister, apply_QFT=True):
    q = qregister
    c = cregister
    numq = len(q)
    numc = len(c)
    qc = circuit
    if apply_QFT == True:
        qc.barrier()
        qc = qc.compose(QFT(num_qubits=numq, approximation_degree=0, do_swaps=True, \
                            inverse=False, insert_barriers=True, name='qft'))
        qc.barrier()
    
    for i, qubit in enumerate(q):
        ncp = PhaseGate(math.pi/2**(numq-i-1)).control(numc)
        qc.append(ncp, [*c, qubit])
        
    if apply_QFT == True:
        qc.barrier() 
        qc = qc.compose(QFT(num_qubits=numq, approximation_degree=0, do_swaps=True, \
                            inverse=True, insert_barriers=True, name='iqft'))
        qc.barrier()
    return qc

test2 = control_increment(qc, q, b)
test2.draw(output='mpl')

## Test increment without control
test = increment(qc, q)

test.measure(q[:], c[:])

run_qc(qc)

test.draw(output='mpl')

# Test control increment

q = QuantumRegister(3)
b = QuantumRegister(1)
c = ClassicalRegister(3)
qc = QuantumCircuit(q, b, c)

qc.x(q[0])

qc = control_increment(qc, q, b)

# Flipping control qubit to 1
qc.x(b[0])

qc = control_increment(qc, q, b)

# Flipping control qubit to 1
qc.x(b[0])

qc = control_increment(qc, q, b)

# Should equal 010
qc.measure(q[:], c[:])

run_qc(qc)
qc.draw(output="mpl")

