import math
number = 3
int(math.ceil(math.log(number, 2)))

import math
import sys
import time

import numpy as np
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister

sys.path[1:1] = ["_common", "_common/qsim", "shors/_common", "quantum-fourier-transform/qsim"]
sys.path[1:1] = ["../../_common", "../../_common/qsim", "../../shors/_common", "../../quantum-fourier-transform/qsim"]

from qiskit import *
from shors_utils import getAngles, getAngle, modinv, generate_base

n=2
num_qubits = 2*n + 3

qr_counting = QuantumRegister(1)

qr_mult = QuantumRegister(n)      # Register for multiplications
qr_aux = QuantumRegister(n+2)     # Register for addition and multiplication
cr_data = ClassicalRegister(2*n)  # Register for measured values of QFT
cr_aux = ClassicalRegister(1)

print(qr_counting)

print(qr_mult,qr_aux,cr_data,cr_aux)

qc = QuantumCircuit(qr_counting, qr_mult, qr_aux, cr_data, cr_aux, name="main")
print(qc)

qc.x(qr_mult[0])
print(qc)

num_gates = 0
depth = 0
############### QFT Circuit

def qft_gate(input_size):
    global QFT_, num_gates, depth
    qr = QuantumRegister(input_size); qc = QuantumCircuit(qr, name="qft")
    
    # Generate multiple groups of diminishing angle CRZs and H gate
    for i_qubit in range(0, input_size):
    
        # start laying out gates from highest order qubit (the hidx)
        hidx = input_size - i_qubit - 1
        
        # if not the highest order qubit, add multiple controlled RZs of decreasing angle
        if hidx < input_size - 1:   
            num_crzs = i_qubit
            for j in range(0, num_crzs):
                divisor = 2 ** (num_crzs - j)
                qc.crz( math.pi / divisor , qr[hidx], qr[input_size - j - 1])
                num_gates += 1
                depth += 1
            
        # followed by an H gate (applied to all qubits)
        qc.h(qr[hidx])
        num_gates += 1
        depth += 1
        
        qc.barrier()
    
    if QFT_ == None or input_size <= 5:
        if input_size < 9: QFT_ = qc
        
    return qc

############### Inverse QFT Circuit

def inv_qft_gate(input_size):
    global QFTI_, num_gates, depth
    qr = QuantumRegister(input_size); qc = QuantumCircuit(qr, name="inv_qft")
    
    # Generate multiple groups of diminishing angle CRZs and H gate
    for i_qubit in reversed(range(0, input_size)):
    
        # start laying out gates from highest order qubit (the hidx)
        hidx = input_size - i_qubit - 1
        
        # precede with an H gate (applied to all qubits)
        qc.h(qr[hidx])
        num_gates += 1
        depth += 1
        
        # if not the highest order qubit, add multiple controlled RZs of decreasing angle
        if hidx < input_size - 1:   
            num_crzs = i_qubit
            for j in reversed(range(0, num_crzs)):
                divisor = 2 ** (num_crzs - j)
                qc.crz( -math.pi / divisor , qr[hidx], qr[input_size - j - 1])
                num_gates += 1
                depth += 1
            
        qc.barrier()  
    
    if QFTI_ == None or input_size <= 5:
        if input_size < 9: QFTI_= qc
        
    return qc

base=2
QC_ = None
PHIADD_ = None
CCPHIADDMODN_ = None
CMULTAMODN_ = None
CUA_ = None
QFT_ = None
QFTI_ = None
def phiADD(num_qubits, a):

    qc = QuantumCircuit(num_qubits, name = "\u03C6ADD")

    angle = getAngles(a, num_qubits)
    for i in range(0, num_qubits):
        # addition
        qc.p(angle[i], i)

    global PHIADD_
    if PHIADD_ == None or num_qubits <= 3:
        if num_qubits < 4: PHIADD_ = qc

    return qc

#Single controlled version of the phiADD circuit
def cphiADD(num_qubits, a):
    phiadd_gate = phiADD(num_qubits,a).to_gate()
    cphiadd_gate = phiadd_gate.control(1)
    return cphiadd_gate

#Doubly controlled version of the phiADD circuit
def ccphiADD(num_qubits, a):
    phiadd_gate = phiADD(num_qubits,a).to_gate()
    ccphiadd_gate = phiadd_gate.control(2)
    return ccphiadd_gate

# Circuit that implements doubly controlled modular addition by a (num qubits should be bit count for number N)
def ccphiADDmodN(num_qubits, a, N):
    qr_ctl = QuantumRegister(2)
    qr_main = QuantumRegister(num_qubits+1)
    qr_ancilla = QuantumRegister(1)
    qc = QuantumCircuit(qr_ctl, qr_main,qr_ancilla, name = "cc\u03C6ADDmodN")

    # Generate relevant gates for circuit
    ccphiadda_gate = ccphiADD(num_qubits+1, a)
    ccphiadda_inv_gate = ccphiADD(num_qubits+1, a).inverse()
    phiaddN_inv_gate = phiADD(num_qubits+1, N).inverse(); phiaddN_inv_gate.name = "inv_\u03C6ADD"
    cphiaddN_gate = cphiADD(num_qubits+1, N)

    # Create relevant temporary qubit lists
    ctl_main_qubits = [i for i in qr_ctl]; ctl_main_qubits.extend([i for i in qr_main])
    anc_main_qubits = [qr_ancilla[0]]; anc_main_qubits.extend([i for i in qr_main])

    #Create circuit
    qc.append(ccphiadda_gate, ctl_main_qubits)
    qc.append(phiaddN_inv_gate, qr_main)

    qc.append(inv_qft_gate(num_qubits+1), qr_main)
    qc.cx(qr_main[-1], qr_ancilla[0])
    qc.append(qft_gate(num_qubits+1), qr_main)

    qc.append(cphiaddN_gate, anc_main_qubits)
    qc.append(ccphiadda_inv_gate, ctl_main_qubits)

    qc.append(inv_qft_gate(num_qubits+1), qr_main)

    qc.x(qr_main[-1])
    qc.cx(qr_main[-1], qr_ancilla[0])
    qc.x(qr_main[-1])

    qc.append(qft_gate(num_qubits+1), qr_main)

    qc.append(ccphiadda_gate, ctl_main_qubits)

    global CCPHIADDMODN_
    if CCPHIADDMODN_ == None or num_qubits <= 2:
        if num_qubits < 3: CCPHIADDMODN_ = qc

    return qc
def ccphiADDmodN_inv(num_qubits, a, N):
    cchpiAddmodN_circ = ccphiADDmodN(num_qubits, a, N)
    cchpiAddmodN_inv_circ = cchpiAddmodN_circ.inverse()
    cchpiAddmodN_inv_circ.name = "inv_cchpiAddmodN"
    return  cchpiAddmodN_inv_circ

# Creates circuit that implements single controlled modular multiplication by a. n represents the number of bits
# needed to represent the integer number N
def cMULTamodN(n, a, N):
    qr_ctl = QuantumRegister(1)
    qr_x = QuantumRegister(n)
    qr_main = QuantumRegister(n+1)
    qr_ancilla = QuantumRegister(1)
    qc = QuantumCircuit(qr_ctl, qr_x, qr_main,qr_ancilla, name = "cMULTamodN")

    # quantum Fourier transform only on auxillary qubits
    qc.append(qft_gate(n+1), qr_main)

    for i in range(n):
        ccphiADDmodN_gate = ccphiADDmodN(n, (2**i)*a % N, N)

        # Create relevant temporary qubit list
        qubits = [qr_ctl[0]]; qubits.extend([qr_x[i]])
        qubits.extend([i for i in qr_main]); qubits.extend([qr_ancilla[0]])

        qc.append(ccphiADDmodN_gate, qubits)

    # inverse quantum Fourier transform only on auxillary qubits
    qc.append(inv_qft_gate(n+1), qr_main)

    global CMULTAMODN_
    if CMULTAMODN_ == None or n <= 2:
        if n < 3: CMULTAMODN_ = qc

    return qc
def controlled_Ua(n,a,exponent,N):
    qr_ctl = QuantumRegister(1)
    qr_x = QuantumRegister(n)
    qr_main = QuantumRegister(n)
    qr_ancilla = QuantumRegister(2)
    qc = QuantumCircuit(qr_ctl, qr_x, qr_main,qr_ancilla, name = f"C-U^{a**exponent}")

    # Generate Gates
    a_inv = modinv(a**exponent,N)
    cMULTamodN_gate = cMULTamodN(n, a**exponent, N)
    cMULTamodN_inv_gate = cMULTamodN(n, a_inv, N).inverse(); cMULTamodN_inv_gate.name = "inv_cMULTamodN"

    # Create relevant temporary qubit list
    qubits = [i for i in qr_ctl]; qubits.extend([i for i in qr_x]); qubits.extend([i for i in qr_main])
    qubits.extend([i for i in qr_ancilla])

    qc.append(cMULTamodN_gate, qubits)

    for i in range(n):
        qc.cswap(qr_ctl, qr_x[i], qr_main[i])

    qc.append(cMULTamodN_inv_gate, qubits)

    global CUA_
    if CUA_ == None or n <= 2:
        if n < 3: CUA_ = qc

    return qc

# perform modular exponentiation 2*n times
for k in range(2*n):
    qc.barrier()

    # Reset the counting qubit to 0 if the previous measurement was 1
    qc.x(qr_counting).c_if(cr_aux,1)
    qc.h(qr_counting)

    cUa_gate = controlled_Ua(n, base,2**(2*n-1-k), number)

    # Create relevant temporary qubit list
    qubits = [qr_counting[0]]; qubits.extend([i for i in qr_mult]);qubits.extend([i for i in qr_aux])

    qc.append(cUa_gate, qubits)

    # perform inverse QFT --> Rotations conditioned on previous outcomes
    for i in range(2**k):
        qc.p(getAngle(i, k), qr_counting[0]).c_if(cr_data, i)

    qc.h(qr_counting)
    qc.measure(qr_counting[0], cr_data[k])
    qc.measure(qr_counting[0], cr_aux[0])

print(qc)

num_shots=100
backend = BasicAer.get_backend('dm_simulator')
options={}
job = execute(qc, backend, shots=num_shots, **options)
result = job.result()
print(result)

qc= qc.decompose()#.decompose().decompose().decompose()
print(qc)

num_shots=100
backend = BasicAer.get_backend('dm_simulator')
options={}
job = execute(qc, backend, shots=num_shots, **options)
result = job.result()
print(result)

