import numpy as np
from numpy import pi
from qiskit import QuantumCircuit, Aer, transpile
from qiskit.visualization import plot_histogram, plot_bloch_multivector
from qiskit.circuit import QuantumRegister, ClassicalRegister

sim = Aer.get_backend('statevector_simulator')
qc = QuantumCircuit(3)

init_state = '101'
qc.initialize(init_state)
qc.barrier()

qc.h(2)
qc.cp(pi/2, 1, 2)
qc.cp(pi/4, 0, 2)

qc.h(1)
qc.cp(pi/2, 0, 1)

qc.h(0)

qc.swap(0, 2)

display(qc.draw())

sv = sim.run(qc).result().get_statevector()
plot_bloch_multivector(sv)

from qiskit.circuit.library import QFT

qc2 = QuantumCircuit(3)
qc2.initialize(init_state)
qc2.append(QFT(num_qubits=3), [0, 1, 2])
display(qc2.draw())
sv2 = sim.run(transpile(qc2, sim)).result().get_statevector()
plot_bloch_multivector(sv2)


# implement general QFT

def myqft(n):
    qc = QuantumCircuit(n)
    for i in reversed(range(n)):
        qc.h(i)
        for e, j in enumerate(reversed(range(i))):
            qc.cp(pi/2**(e+1), j, i)
    for i in range(n//2):
        qc.swap(i, n-i-1)
    return qc.to_gate(label='myqft{}'.format(n))

qc3 = QuantumCircuit(3)
qc3.initialize(init_state)
qc3.append(myqft(3), [0, 1, 2])
display(qc3.draw())
sv3 = sim.run(transpile(qc3, sim)).result().get_statevector()
plot_bloch_multivector(sv3)

def create_carrier(name=None):
    qc = QuantumCircuit(4)
    qc.toffoli(1, 2, 3)
    qc.cx(1, 2)
    qc.toffoli(0, 2, 3)
    return qc.to_gate(label=name or 'carrier')

def create_carrier_inv(name=None):
    qc = QuantumCircuit(4)
    qc.append(create_carrier().inverse(), range(4))
    return qc.to_gate(label=name or 'carrier_inv')

def create_sum(name=None):
    qc = QuantumCircuit(3)
    qc.cx(1, 2)
    qc.cx(0, 2)
    return qc.to_gate(label=name or 'sum')

def create_sum_inv(name=None):
    qc = QuantumCircuit(3)
    qc.append(create_sum().inverse(), range(3))
    return qc.to_gate(label=name or 'sum_inv')

a = QuantumRegister(3, 'a')
b = QuantumRegister(4, 'b')
c = QuantumRegister(3, 'c')
creg = ClassicalRegister(4, 'creg')
qc = QuantumCircuit(a, b, c, creg)
qc.x(a[0])
qc.x(a[1])
qc.x(b[1])
qc.x(b[2])
qc.append(create_carrier(), [c[0], a[0], b[0], c[1]])
qc.append(create_carrier(), [c[1], a[1], b[1], c[2]])
qc.append(create_carrier(), [c[2], a[2], b[2], b[3]])
qc.cx(a[2], b[2])
qc.append(create_sum(), [c[2], a[2], b[2]])
qc.append(create_carrier_inv(), [c[1], a[1], b[1], c[2]])
qc.append(create_sum(), [c[1], a[1], b[1]])
qc.append(create_carrier_inv(), [c[0], a[0], b[0], c[1]])
qc.append(create_sum(), [c[0], a[0], b[0]])
qc.barrier()
qc.measure(b, creg)
display(qc.draw())
backend = Aer.get_backend('qasm_simulator')
job = backend.run(transpile(qc, backend), shots=1024)
counts = job.result().get_counts()
print(counts)

def create_adder(n):
    """create adder with n bits a, n+1 bits b and n bits carrier"""
    a = QuantumRegister(n, 'a')
    b = QuantumRegister(n+1, 'b')
    c = QuantumRegister(n, 'c')
    qc = QuantumCircuit(a, b, c)
    for i in range(n-1):
        qc.append(create_carrier(), [c[i], a[i], b[i], c[i+1]])
    qc.append(create_carrier(), [c[n-1], a[n-1], b[n-1], b[n]])
    qc.cx(a[n-1], b[n-1])
    qc.append(create_sum(), [c[n-1], a[n-1], b[n-1]])
    for i in reversed(range(n-1)):
        qc.append(create_carrier_inv(), [c[i], a[i], b[i], c[i+1]])
        qc.append(create_sum(), [c[i], a[i], b[i]])
    return qc.to_gate(label='adder')

def create_adder_inv(n):
    qc = QuantumCircuit(3*n+1)
    qc.append(create_adder(n).inverse(), range(3*n+1))
    return qc.to_gate(label='adder_inv')


qc = QuantumCircuit(10, 4)
qc.x(0)
qc.x(1)
qc.x(4)
qc.x(5)
qc.append(create_adder(3), range(10))
qc.measure([3, 4, 5, 6], range(4))
display(qc.draw())
backend = Aer.get_backend('qasm_simulator')
job = backend.run(transpile(qc, backend), shots=1024)
counts = job.result().get_counts()
print(counts)

def create_adder_mod(n, bigN):
    """
    :param n: number of bits for a, c and bigN; b has n+1 bits
    :param bigN: mod number
    :return: Gate of 4n+2 qubits
        0-2 qubits for a
        3-6 qubits for b
        7-9 qubits for c (init to 0)
        10-12 qubits for bigN (init to bigN when use)
        13 qubit for temp (init to 0)
    """
    a = QuantumRegister(n, 'a')
    b = QuantumRegister(n+1, 'b')
    c = QuantumRegister(n, 'c')
    bN = QuantumRegister(n, 'bigN')
    t = QuantumRegister(1, 't')
    qc = QuantumCircuit(a, b, c, bN, t)
    # block 1
    qc.append(create_adder(n), list(a)+list(b)+list(c))
    for i in range(n):
        qc.swap(a[i], bN[i])
    qc.append(create_adder_inv(n), list(a)+list(b)+list(c))
    qc.x(b[n])
    qc.cx(b[n], t[0])
    qc.x(b[n])
    tempN = bigN
    i = 0
    while tempN != 0:
        if tempN % 2 != 0:
            qc.cx(t[0], a[i])
        i = i + 1
        tempN = tempN // 2
    qc.append(create_adder(n), list(a)+list(b)+list(c))
    tempN = bigN
    i = 0
    while tempN != 0:
        if tempN % 2 != 0:
            qc.cx(t[0], a[i])
        i = i + 1
        tempN = tempN // 2
    for i in range(n):
        qc.swap(a[i], bN[i])
    # block 2
    qc.append(create_adder_inv(n), list(a)+list(b)+list(c))
    qc.cx(b[n], t[0])
    qc.append(create_adder(n), list(a)+list(b)+list(c))
    #display(qc.draw())
    return qc.to_gate(label='adder_mod')


n = 3
bigN = 5
#qc = QuantumCircuit(4*n+2)
a = QuantumRegister(n, 'a')
b = QuantumRegister(n+1, 'b')
c = QuantumRegister(n, 'c')
bN = QuantumRegister(n, 'bigN')
t = QuantumRegister(1, 't')
creg = ClassicalRegister(4, 'creg')
qc = QuantumCircuit(a, b, c, bN, t, creg)
# qc.x(a[0]) # a = 3
# qc.x(a[1])
qc.x(a[2]) # a = 4
qc.x(b[2]) # b = 4
qc.x(bN[0])
qc.x(bN[2]) # bN = 5
qc.append(create_adder_mod(n, bigN), list(a) + list(b) + list(c) + list(bN) + list(t))
qc.measure(b, creg)
display(qc.draw())
backend = Aer.get_backend('qasm_simulator')
job = backend.run(transpile(qc, backend), shots=1024)
counts = job.result().get_counts()
print(counts)

def create_ctrl_mult_mod(n, bigN, m):
    """
    Calculate zm mod N and output to b

    :param n: number of bits for a, c and bigN; b has n+1 bits
    :param bigN: mod number
    :param m: multiplier
    :return: Gate of 4n+2 qubits
        0 (1 qubit) for x
        1 to n (n qubits) for z
        n+1 to 2n (n qubits) for a
        2n+1 to 3n+1 (n+1 qubits) for b
        3n+2 to 4n+1 (n qubits) for c (init to 0)
        4n+2 to 5n+1 (n qubits) for bigN (init to bigN when use)
        5n+2 (1 qubit) for temp (init to 0)
    """
    x = QuantumRegister(1, 'x')
    z = QuantumRegister(n, 'z')
    a = QuantumRegister(n, 'a')
    b = QuantumRegister(n+1, 'b')
    c = QuantumRegister(n, 'c')
    bN = QuantumRegister(n, 'bN')
    t = QuantumRegister(1, 't')
    qc = QuantumCircuit(x, z, a, b, c, bN, t)
    next_mod = m # m 2^0 mod N
    for j in range(n):
        temp_mod = next_mod
        i = 0
        while temp_mod != 0:
            if temp_mod % 2 != 0:
                qc.ccx(x[0], z[j], a[i])
            i = i + 1
            temp_mod = temp_mod // 2

        qc.append(create_adder_mod(n, bigN), list(a) + list(b) + list(c) + list(bN) + list(t))

        temp_mod = next_mod
        i = 0
        while temp_mod != 0:
            if temp_mod % 2 != 0:
                qc.ccx(x[0], z[j], a[i])
            i = i + 1
            temp_mod = temp_mod // 2
    
        next_mod = (next_mod * 2) % bigN # update for m 2^i+1 mod N 
    
    qc.x(x[0])
    for j in range(n):
        qc.ccx(x[0], z[j], b[j])
    qc.x(x[0])

    # display(qc.draw())
    return qc.to_gate(label='ctrl_mult_mod')

def create_ctrl_mult_mod_inv(n, bigN, m):
    qc = QuantumCircuit(5*n+3)
    qc.append(create_ctrl_mult_mod(n, bigN, m).inverse(), range(5*n+3))
    return qc.to_gate(label='ctrl_mult_mod_inv')


n = 3
bigN = 5
m = 3
x = QuantumRegister(1, 'x')
z = QuantumRegister(n, 'z')
a = QuantumRegister(n, 'a')
b = QuantumRegister(n+1, 'b')
c = QuantumRegister(n, 'c')
bN = QuantumRegister(n, 'bN')
t = QuantumRegister(1, 't')
creg = ClassicalRegister(n, 'creg')
qc = QuantumCircuit(x, z, a, b, c, bN, t, creg)
qc.x(x[0]) # x = 1
qc.x(z[0]) # z = 3
qc.x(z[1])
qc.x(bN[0]) # bN = 5
qc.x(bN[2])
qc.append(create_ctrl_mult_mod(n, bigN, m), list(x) + z[:] + list(a) + list(b) + list(c) + list(bN) + list(t))
# for i in range(n):
#     qc.measure(b[i], creg[i])
qc.measure(b[0:n], creg)
display(qc.draw())
backend = Aer.get_backend('qasm_simulator')
job = backend.run(transpile(qc, backend), shots=1024)
counts = job.result().get_counts()
print(counts)

from sympy import mod_inverse

def create_mod_exp(n, bigN, y, nx):
    """
    Calculate y^x mod N and output to z

    :param n: number of bits for z, a, c and bigN; b has n+1 bits
    :param bigN: mod number
    :param y: multiplier
    :param nx: number of bits for x
    :return: Gate of 4n+2 qubits
        0 to nx-1 (n qubit) for x
        nx to nx+n-1 (n qubits) for z (init to 1 when use)
        nx+n to nx+2n-1 (n qubits) for a (init to 0)
        nx+2n to nx+3n (n+1 qubits) for b (init to 0)
        nx+3n+1 to nx+4n (n qubits) for c (init to 0)
        nx+4n+1 to nx+5n (n qubits) for bigN (init to bigN when use)
        nx+5n+1 (1 qubit) for temp (init to 0)
    """
    x = QuantumRegister(nx, 'x')
    z = QuantumRegister(n, 'z')
    a = QuantumRegister(n, 'a')
    b = QuantumRegister(n+1, 'b')
    c = QuantumRegister(n, 'c')
    bN = QuantumRegister(n, 'bN')
    t = QuantumRegister(1, 't')
    qc = QuantumCircuit(x, z, a, b, c, bN, t)
    m = y
    for i in range(nx):
        qc.append(create_ctrl_mult_mod(n, bigN, m), 
            [x[i]] + z[:] + a[:] + b[:] + c[:] + bN[:] + list(t))
        for j in range(n):
            qc.cswap(x[i], z[j], b[j])
        m_inv = mod_inverse(m, bigN)
        qc.append(create_ctrl_mult_mod_inv(n, bigN, m_inv), 
            [x[i]] + z[:] + a[:] + b[:] + c[:] + bN[:] + list(t))
        m = (m * m) % bigN # update for next cycle
    # display(qc.draw())
    return qc.to_gate(label='mod_exp')


n = 3
bigN = 5
y = 3
nx = 6
x = QuantumRegister(nx, 'x')
z = QuantumRegister(n, 'z')
a = QuantumRegister(n, 'a')
b = QuantumRegister(n+1, 'b')
c = QuantumRegister(n, 'c')
bN = QuantumRegister(n, 'bN')
t = QuantumRegister(1, 't')
creg = ClassicalRegister(n, 'creg')
qc = QuantumCircuit(x, z, a, b, c, bN, t, creg)
#qc.x(x[2]) # x = 4
# qc.x(x[0]) # x = 3
# qc.x(x[1])
qc.x(x[1]) # x = 42
qc.x(x[3])
qc.x(x[5])
qc.x(z[0]) # z = 1 init
qc.x(bN[0]) # bN = 5
qc.x(bN[2])
qc.append(create_mod_exp(n, bigN, y, nx), x[:] + z[:] + a[:] + b[:] + c[:] + bN[:] + list(t))
qc.measure(z, creg)
display(qc.draw())
backend = Aer.get_backend('qasm_simulator')
job = backend.run(transpile(qc, backend), shots=1)
counts = job.result().get_counts()
print(counts)

n = 5   # use 4 qubits may not work, do not know the reason
bigN = 15
y = 7
nx = 8
x = QuantumRegister(nx, 'x')
z = QuantumRegister(n, 'z')
a = QuantumRegister(n, 'a')
b = QuantumRegister(n+1, 'b')
c = QuantumRegister(n, 'c')
bN = QuantumRegister(n, 'bN')
t = QuantumRegister(1, 't')
# cregz = ClassicalRegister(n, 'cregz')
cregx = ClassicalRegister(nx, 'cregx')
# qc = QuantumCircuit(x, z, a, b, c, bN, t, cregz, cregx)
qc = QuantumCircuit(x, z, a, b, c, bN, t, cregx)
qc.h(x)
qc.x(z[0]) # z = 1 init
qc.x(bN[0]) # bN = 15 = 1111(2)
qc.x(bN[1]) 
qc.x(bN[2])
qc.x(bN[3])

qc.append(create_mod_exp(n, bigN, y, nx), x[:] + z[:] + a[:] + b[:] + c[:] + bN[:] + list(t))
qc.append(myqft(nx).inverse(), x)
# qc.measure(z, cregz)
qc.measure(x, cregx)
display(qc.draw())
backend = Aer.get_backend('qasm_simulator')
job = backend.run(transpile(qc, backend), shots=1024)
counts = job.result().get_counts()
print(counts)


from fractions import Fraction
Fraction(192/2**8).limit_denominator(15)

import math

print(math.gcd(7**2+1, 15))
print(math.gcd(7**2-1, 15))

n = 6
bigN = 55
y = 7
nx = 12
x = QuantumRegister(nx, 'x')
z = QuantumRegister(n, 'z')
a = QuantumRegister(n, 'a')
b = QuantumRegister(n+1, 'b')
c = QuantumRegister(n, 'c')
bN = QuantumRegister(n, 'bN')
t = QuantumRegister(1, 't')
# cregz = ClassicalRegister(n, 'cregz')
cregx = ClassicalRegister(nx, 'cregx')
# qc = QuantumCircuit(x, z, a, b, c, bN, t, cregz, cregx)
qc = QuantumCircuit(x, z, a, b, c, bN, t, cregx)
qc.h(x)
qc.x(z[0]) # z = 1 init
qc.x(bN[0]) # bN = 55 = 110111(2)
qc.x(bN[1])
qc.x(bN[2])
qc.x(bN[4])
qc.x(bN[5])
# qc.x(bN[7])

qc.append(create_mod_exp(n, bigN, y, nx), x[:] + z[:] + a[:] + b[:] + c[:] + bN[:] + list(t))
qc.append(myqft(nx).inverse(), x)
# qc.measure(z, cregz)
qc.measure(x, cregx)
display(qc.draw())
backend = Aer.get_backend('qasm_simulator')
job = backend.run(transpile(qc, backend), shots=1024)
counts = job.result().get_counts()
print(counts)


from fractions import Fraction
Fraction(205/2**12).limit_denominator(55)

import math

print(math.gcd(7**10+1, 55))
print(math.gcd(7**10-1, 55))

n = 8
bigN = 221
y = 7
nx = 16
x = QuantumRegister(nx, 'x')
z = QuantumRegister(n, 'z')
a = QuantumRegister(n, 'a')
b = QuantumRegister(n+1, 'b')
c = QuantumRegister(n, 'c')
bN = QuantumRegister(n, 'bN')
t = QuantumRegister(1, 't')
# cregz = ClassicalRegister(n, 'cregz')
cregx = ClassicalRegister(nx, 'cregx')
# qc = QuantumCircuit(x, z, a, b, c, bN, t, cregz, cregx)
qc = QuantumCircuit(x, z, a, b, c, bN, t, cregx)
qc.h(x)
qc.x(z[0]) # z = 1 init
qc.x(bN[0]) # bN = 221 = 11011101(2)
qc.x(bN[2])
qc.x(bN[3])
qc.x(bN[4])
qc.x(bN[6])
qc.x(bN[7])

qc.append(create_mod_exp(n, bigN, y, nx), x[:] + z[:] + a[:] + b[:] + c[:] + bN[:] + list(t))
qc.append(myqft(nx).inverse(), x)
# qc.measure(z, cregz)
qc.measure(x, cregx)
display(qc.draw())
backend = Aer.get_backend('qasm_simulator')
job = backend.run(transpile(qc, backend), shots=1)
counts = job.result().get_counts()
print(counts)


from fractions import Fraction
Fraction(53248/2**16).limit_denominator(221)

import math

print(math.gcd(7**8+1, 221))
print(math.gcd(7**8-1, 221))

17*13

n = 9
bigN = 437
y = 7
nx = 18
x = QuantumRegister(nx, 'x')
z = QuantumRegister(n, 'z')
a = QuantumRegister(n, 'a')
b = QuantumRegister(n+1, 'b')
c = QuantumRegister(n, 'c')
bN = QuantumRegister(n, 'bN')
t = QuantumRegister(1, 't')
# cregz = ClassicalRegister(n, 'cregz')
cregx = ClassicalRegister(nx, 'cregx')
# qc = QuantumCircuit(x, z, a, b, c, bN, t, cregz, cregx)
qc = QuantumCircuit(x, z, a, b, c, bN, t, cregx)
qc.h(x)
qc.x(z[0]) # z = 1 init
qc.x(bN[0]) # bN = 437 = 110110101(2)
qc.x(bN[2])
qc.x(bN[4])
qc.x(bN[5])
qc.x(bN[7])
qc.x(bN[8])

qc.append(create_mod_exp(n, bigN, y, nx), x[:] + z[:] + a[:] + b[:] + c[:] + bN[:] + list(t))
qc.append(myqft(nx).inverse(), x)
# qc.measure(z, cregz)
qc.measure(x, cregx)
display(qc.draw())
backend = Aer.get_backend('qasm_simulator')
job = backend.run(transpile(qc, backend), shots=1024)
counts = job.result().get_counts()
print(counts)


from fractions import Fraction
Fraction(170791/2**18).limit_denominator(437)

import math

print(math.gcd(7**33+1, 437))
print(math.gcd(7**33-1, 437))