import numpy as np
# importing Qiskit
from qiskit import QuantumCircuit, transpile, Aer, IBMQ
from qiskit.providers.ibmq import least_busy
from qiskit.tools.monitor import job_monitor
from qiskit.visualization import plot_histogram, plot_bloch_multivector

# 3 qubit example

qc = QuantumCircuit(3)
qc.h(2)
qc.cp(np.pi/2, 1, 2) # CROT from qubit 1 to qubit 2
qc.cp(np.pi/4, 0, 2) # CROT from qubit 0 to qubit 2
qc.h(1)
qc.cp(np.pi/2, 0, 1) # CROT from qubit 0 to qubit 1
qc.h(0)
qc.swap(0, 2) #swaps qubit 0 and 2
qc.draw('mpl')

# General QFT function

def qft_rot(circuit, n):
    '''
        circuit: given quantum circuit
        n: number of qubit 
    '''
    if n==0:
        return circuit
    n -= 1   # we want the qubit indexing to go from 0 to n-2
    circuit.h(n)
    for qubit in range(n):
        circuit.cp(np.pi/2**(n-qubit), qubit, n)
    
    qft_rot(circuit, n) # recursively call the function to repeat the process for next n-1 qubits

def swap_registers(circuit, n):
    for qubit in range(n//2):
        circuit.swap(qubit, n-1-qubit)
    return circuit

def qft(circuit, n):
    ''' QFT on the n qubits
    '''
    qft_rot(circuit, n)
    swap_registers(circuit, n)
    return circuit

# Draw the circuit with n=4
qc = QuantumCircuit(4)
qft(qc, 4)
qc.draw('mpl')

bin(13)

# Checking the circuit with 13='1101'
#quantum circuit that encodes '1101'
qc = QuantumCircuit(4)
qc.x([0,2,3])
qc.draw('mpl')

sim = Aer.get_backend("aer_simulator")
qc_init = qc.copy()
qc_init.save_statevector()
statevector = sim.run(qc_init).result().get_statevector()
plot_bloch_multivector(statevector)

qft(qc, 4)
qc.draw('mpl')

qc.save_statevector()
statevector = sim.run(qc).result().get_statevector()
plot_bloch_multivector(statevector)

# Inverse QFT 

def inv_qft(circuit, n):
    """Does the inverse QFT on the first n qubits in circuit"""
    # First we create a QFT circuit of the correct size:
    qft_circ = qft(QuantumCircuit(n), n)
    # Then we take the inverse of this circuit
    invqft_circ = qft_circ.inverse()
    # And add it to the first n qubits in our existing circuit
    circuit.append(invqft_circ, circuit.qubits[:n])
    return circuit.decompose() # .decompose() allows us to see the individual gates

# Create the state |13>

state = 13
qc1 = QuantumCircuit(4)
for qubit in range(4):
    qc1.h(qubit)
qc1.p(state*np.pi/8, 0)
qc1.p(state*np.pi/4, 1)
qc1.p(state*np.pi/2, 2)
qc1.p(state*np.pi, 3)

qc1.draw('mpl')

# Check this creates the state |13>
qc1_init = qc1.copy()
qc1_init.save_statevector()
sim = Aer.get_backend("aer_simulator")
statevector = sim.run(qc1_init).result().get_statevector()
plot_bloch_multivector(statevector)

# Apply inverse QFT
inv_qc = inv_qft(qc1, 4)
inv_qc.measure_all()
inv_qc.draw('mpl')

# check on simulator
counts = sim.run(qc_init).result().get_counts()
plot_histogram(counts)

# initialize 

qpe = QuantumCircuit(4, 3)
qpe.x(3)
# apply the hadamard gates
for qubit in range(3):
    qpe.h(qubit)

j = 1
for qubit in range(3):
    for i in range(j):
        qpe.cp(np.pi/4, qubit, 3)
    j *= 2

qpe.barrier()
# apply inverse QFT, we use the in built QFT function
from qiskit.circuit.library import QFT
qpe = qpe.compose(QFT(3, inverse=True), [0,1,2])
# Measure
qpe.barrier()
for n in range(3):
    qpe.measure(n,n)

qpe.draw('mpl')

aer_sim = Aer.get_backend('aer_simulator')
shots = 2048
t_qpe = transpile(qpe, aer_sim)
results = aer_sim.run(t_qpe, shots=shots).result()
answer = results.get_counts()

plot_histogram(answer)

# Create and set up circuit
qpe2 = QuantumCircuit(4, 3)

# Apply H-Gates to counting qubits:
for qubit in range(3):
    qpe2.h(qubit)

# Prepare our eigenstate |psi>:
qpe2.x(3)

# Do the controlled-U operations:
angle = 2*np.pi/3
repetitions = 1
for counting_qubit in range(3):
    for i in range(repetitions):
        qpe2.cp(angle, counting_qubit, 3);
    repetitions *= 2

# Do the inverse QFT:
qpe2 = qpe2.compose(QFT(3, inverse=True), [0,1,2])

# Measure of course!
for n in range(3):
    qpe2.measure(n,n)

qpe2.draw('mpl')

# Let's see the results!
shots = 4096
t_qpe2 = transpile(qpe2, aer_sim)
results = aer_sim.run(t_qpe2, shots=shots).result()
answer = results.get_counts()

plot_histogram(answer)

# Examples of periods
N = 35
a = 3
import matplotlib.pyplot as plt
# Calculate the plotting data
xvals = np.arange(N)
yvals = [np.mod(a**x, N) for x in xvals]

# Use matplotlib to display it nicely
fig, ax = plt.subplots()
ax.plot(xvals, yvals, linewidth=1, linestyle='dotted', marker='x')
ax.set(xlabel='$x$', ylabel=f'${a}^x$ mod ${N}$',
       title="Example of Periodic Function in Shor's Algorithm")
try: # plot r on the graph
    r = yvals[1:].index(1) + 1
    plt.annotate('', xy=(0,1), xytext=(r,1),
                 arrowprops=dict(arrowstyle='<->'))
    plt.annotate(f'$r={r}$', xy=(r/3,1.5))
except ValueError:
    print('Could not find period, check a < N and have no common factors.')

# We solve period finding problem for N=15
# The function c_amod15 returns the controlled-U gate for a, repeated power times.
def c_amod15(a, power):
    """Controlled multiplication by a mod 15"""
    if a not in [2,4,7,8,11,13]:
        raise ValueError("'a' must be 2,4,7,8,11 or 13")
    U = QuantumCircuit(4)
    for _iteration in range(power):
        if a in [2,13]:
            U.swap(2,3)
            U.swap(1,2)
            U.swap(0,1)
        if a in [7,8]:
            U.swap(0,1)
            U.swap(1,2)
            U.swap(2,3)
        if a in [4, 11]:
            U.swap(1,3)
            U.swap(0,2)
        if a in [7,11,13]:
            for q in range(4):
                U.x(q)
    U = U.to_gate()
    U.name = f"{a}^{power} mod 15"
    c_U = U.control()
    return c_U

# Specify variables
N_COUNT = 8  # number of counting qubits
a = 8

# Create QuantumCircuit with N_COUNT counting qubits
# plus 4 qubits for U to act on
qc = QuantumCircuit(N_COUNT + 4, N_COUNT)

# Initialize counting qubits
# in state |+>
for qubit in range(N_COUNT):
    qc.h(qubit)

# And auxiliary register in state |1>
qc.x(N_COUNT)

# Do controlled-U operations
for qubit in range(N_COUNT):
    qc.append(c_amod15(a, 2**qubit), [qubit] + [i+N_COUNT for i in range(4)])

# Do inverse-QFT
qc.append(QFT(N_COUNT, inverse=True), range(N_COUNT))

# Measure circuit
qc.measure(range(N_COUNT), range(N_COUNT))
qc.draw('mpl', fold=-1)   

# measurement results
t_qc = transpile(qc, aer_sim)
counts = aer_sim.run(t_qc).result().get_counts()
plot_histogram(counts)

import pandas as pd
rows, measured_phases = [], []
for output in counts:
    decimal = int(output, 2)  # Convert (base 2) string to decimal
    phase = decimal/(2**N_COUNT)  # Find corresponding eigenvalue
    measured_phases.append(phase)
    # Add these values to the rows in our table:
    rows.append([f"{output}(bin) = {decimal:>3}(dec)",
                 f"{decimal}/{2**N_COUNT} = {phase:.2f}"])
# Print the rows in a table
headers=["Register Output", "Phase"]
df = pd.DataFrame(rows, columns=headers)
df

from fractions import Fraction
rows = []
for phase in measured_phases:
    frac = Fraction(phase).limit_denominator(15)
    rows.append([phase,
                 f"{frac.numerator}/{frac.denominator}",
                 frac.denominator])
# Print as a table
headers=["Phase", "Fraction", "Guess for r"]
df = pd.DataFrame(rows, columns=headers)
print(df)

def a_2jmodN(a, j, N):
    ''' Computes a^{2^j} mod N '''
    for _ in range(j):
        a = np.mod(a**2 , N)
    return a

# Factoring 15 using period finding
def qpe_amod15(a):
    """Performs quantum phase estimation on the operation a*r mod 15.
    Args:
        a (int): This is 'a' in a*r mod 15
    Returns:
        float: Estimate of the phase
    """
    N_COUNT = 8
    qc = QuantumCircuit(4+N_COUNT, N_COUNT)
    for q in range(N_COUNT):
        qc.h(q)     # Initialize counting qubits in state |+>
    qc.x(3+N_COUNT) # And auxiliary register in state |1>
    for q in range(N_COUNT): # Do controlled-U operations
        qc.append(c_amod15(a, 2**q), [q] + [i+N_COUNT for i in range(4)])
        
    qc.append(QFT(N_COUNT, inverse=True), range(N_COUNT)) # Do inverse-QFT
    qc.measure(range(N_COUNT), range(N_COUNT))
    # Simulate Results
    aer_sim = Aer.get_backend('aer_simulator')
    # `memory=True` tells the backend to save each measurement in a list
    job = aer_sim.run(transpile(qc, aer_sim), shots=1, memory=True)
    readings = job.result().get_memory()
    print("Register Reading: " + readings[0])
    phase = int(readings[0],2)/(2**N_COUNT)
    print(f"Corresponding Phase: {phase}")
    return phase

phase = qpe_amod15(a) # Phase = s/r
Fraction(phase).limit_denominator(15)
frac = Fraction(phase).limit_denominator(15)
s, r = frac.numerator, frac.denominator
print(r)

from math import gcd
# repeats the algorithm until at least one factor of 15 is found
a = 7
FACTOR_FOUND = False
ATTEMPT = 0
while not FACTOR_FOUND:
    ATTEMPT += 1
    print(f"\nATTEMPT {ATTEMPT}:")
    phase = qpe_amod15(a) # Phase = s/r
    frac = Fraction(phase).limit_denominator(N)
    r = frac.denominator
    print(f"Result: r = {r}")
    if phase != 0:
        # Guesses for factors are gcd(x^{r/2} ±1 , 15)
        guesses = [gcd(a**(r//2)-1, N), gcd(a**(r//2)+1, N)]
        print(f"Guessed Factors: {guesses[0]} and {guesses[1]}")
        for guess in guesses:
            if guess not in [1,N] and (N % guess) == 0:
                # Guess is a factor!
                print("*** Non-trivial factor found: {guess} ***")
                FACTOR_FOUND = True

