#!/usr/bin/env python
# coding: utf-8

# # Shor's factoring algorithm

# In[1]:


from qsystem import QSystem, Gate
from random import randint
from math import log2, ceil, floor, gcd
from IPython.display import display, Latex, Math


# In[2]:


n = 15
Latex('Factoring $n = {}$'.format(n))


# ## step 1: Random select an $a$ less than $n$ and coprime with $n$

# In[3]:


a = 0
while gcd(n, a) != 1 or a == 1:
    a = randint(2, n)
Math('a = {}'.format(a))


# ## step 2: use the quantum period-finding to find the period $r$ of the function $f(x) = a^x \mod n$
# 
# ### Period-finding
# 
# It will be necessary 2 quantum registers with size of 
# $$s = \lceil\log_2(n+1)\rceil$$

# In[4]:


s = ceil(log2(n+1))
Math('s = {}'.format(s))


# and a quantum oracle 'POWN' that 
# $$\left|x\right>\left|0\right>
# \xrightarrow{\text{POWN}}
# \left|x\right>\left|a^x (\text{mod}\,n)\right>$$

# In[5]:


def pown(x):
    x = x >> s
    fx = pow(a, x, n)
    return (x << s) | fx

def it():
    for x in range(2**s):
        yield x << s
    
pown_gate = Gate.from_func(pown, 2*s, it())


# In[6]:


seed = randint(0,10000)
q = QSystem(s, seed)

print('init first register at')
print(q)


# ### step 1: Prepare a superposition
# $$\left|0\right>\left|0\right>
# \xrightarrow{H^{\otimes n}}
# {1\over\sqrt{2^s}}\sum_{x=0}^{2^{n}-1} \left|x\right>\left|0\right>
# $$

# In[7]:


q.evol(gate='H', qbit=0, count=s)
print(q)


# ### step 2: Prepare a periodic superposition
# $$
# {1\over\sqrt{2^s}}\sum_{x=0}^{2^{n}-1} \left|x\right>\left|0\right>
# \xrightarrow{\text{POWN}}
# {1\over\sqrt{2^s}}\sum_{x=0}^{2^{n}-1}
# \left|x\right>\left|a^x(\text{mod}\, n)\right>
# $$

# In[8]:


print('add the second register')
q.add_ancillas(s)
q.apply(gate=pown_gate, qbit=0)
print(q)


# ### step 3 (optional): measure the second register
# help to understand the algorithm
# 
# $$
# {1\over\sqrt{2^s}}\sum_{x=0}^{2^{s}-1}
# \left|x\right>\left|a^x(\text{mod}\, n)\right>
# \xrightarrow{\text{measure}[s:2s]}
# \sqrt{r\over{2^s}}\sum_{i=0}^{{2^{s}\over r}-1}
# \left|ir + x_0\right>\left|a^{x_0}(\text{mod}\, n)\right>
# $$
# 

# In[9]:


print('measure and remove the second register')
q.measure(s, s)
q.rm_ancillas() 
print(q)


# ### step 4: fourier transform of the first register
# $$
# \sqrt{r\over{2^s}}\sum_{i=0}^{{2^{s}\over r}-1}
# \left|ir + x_0\right>
# \xrightarrow{\text{QFT}_s}
# {1\over\sqrt{r}}\sum_{i=0}^{r-1}\left|i{2^s\over r}\right>e^{\phi_i}
# $$

# In[10]:


q.qft(qbit_begin=0, qbit_end=s)
print(q)


# ### step 5: measure the first register and repeat the algorithm to measure distincts multiples of $2^s\over r$

# In[11]:


q.measure_all()
c = q.bits()
c = sum([m*2**i for m, i in zip(c, reversed(range(len(c))))])
mea = [c]

for _ in range(s-1):
    seed = randint(0,10000)
    q = QSystem(s, seed)

    q.evol('H', 0, s)     # 1
    q.add_ancillas(s)
    q.apply(pown_gate, 0) # 2
    q.measure(s, s)
    q.rm_ancillas()
    q.qft(0, s)           # 4
    q.measure_all()       # 5

    c = q.bits()
    c = sum([m*2**i for m, i in zip(c, reversed(range(len(c))))])
    mea.append(c)
print('measurements results =', mea)


# ### step 6: with the measurements compute 
# $$
# r = {2^s\over\gcd(\text{measurements})}
# $$

# In[12]:


c = mea[0]
for m in mea:
    c = gcd(c, m)
if c == 0:
    print('repite the period-finding algorithm')
else:
    r = 2**s/c
    display(Latex('possible period $r = {}$'.format(r)))


# ## step 3: if $r$ is odd go to step 1 else compute the two nontrivial factors of $n$, $pq = n$
# $$
# p = \gcd(a^{r\over2}-1, n)
# $$
# $$
# q = \gcd(a^{r\over2}+1, n)
# $$

# In[13]:


if c == 0:
    print('repite the period-finding algorithm')
elif r % 2 == 1:
    print('go to step 1')
else: 
    p = gcd(int(a**(r/2)+1), n)
    q = gcd(int(a**(r/2)-1), n)
    display(Math('{}*{}={}'.format(p,q,p*q)))