# This code has been adapted and modified from IBM Qiskit 2021 and also from https://github.com/ttlion/ShorAlgQiskit.
# It uses the implementation as contained in the work of Stephane Beauregard (https://arxiv.org/abs/quant-ph/0205095)
# Many thanks to IBM Qiskit team, Tiago Miguel (ttlion), Qubit by Qubit, Peter Shor and Stephane Beauregard.

from typing import Optional, Union, Tuple, List
import math
import array
import fractions
import logging
import numpy as np
from qiskit import ClassicalRegister, QuantumCircuit, QuantumRegister, execute, IBMQ, transpile,BasicAer, Aer, assemble
from qiskit.circuit import Gate, Instruction, ParameterVector
from qiskit.circuit.library import QFT
from qiskit.providers import BaseBackend, Backend
from qiskit.quantum_info import partial_trace
from qiskit.utils import summarize_circuits
from qiskit.utils.arithmetic import is_power
from qiskit.utils.validation import validate_min
from qiskit.utils.quantum_instance import QuantumInstance
import qiskit.visualization
from qiskit.providers.aer import QasmSimulator
from datetime import datetime
import csv
# provider = IBMQ.enable_account("PUT TOKEN HERE")
backend = QasmSimulator()
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"  #"last_expr" or "all"

# """ Function to check if N is of type q^p"""
def check_if_power(N):
#     """ Check if N is a perfect power in O(n^3) time, n=ceil(logN) """
    b=2
    while (2**b) <= N:
        a = 1
        c = N
        while (c-a) >= 2:
            m = int( (a+c)/2 )

            if (m**b) < (N+1):
                p = int( (m**b) )
            else:
                p = int(N+1)

            if int(p) == int(N):
                print('N is {0}^{1}'.format(int(m),int(b)) )
                return True

            if p<N:
                a = int(m)
            else:
                c = int(m)
        b=b+1

    return False

def egcd(a, b):
    if a == 0:
        return (b, 0, 1)
    else:
        g, y, x = egcd(b % a, a)
        return (g, x - (b // a) * y, y)
    
def modinv(a, m):
    g, x, y = egcd(a, m)
    if g != 1:
        raise Exception('modular inverse does not exist')
    else:
        return x % m

def create_QFT(circuit,up_reg,n,with_swaps):
    i=n-1
    while i>=0:
        circuit.h(up_reg[i])        
        j=i-1  
        while j>=0:
            if (np.pi)/(pow(2,(i-j))) > 0:
                circuit.cu1( (np.pi)/(pow(2,(i-j))) , up_reg[i] , up_reg[j] )
                j=j-1   
        i=i-1  

    if with_swaps==1:
        i=0
        while i < ((n-1)/2):
            circuit.swap(up_reg[i], up_reg[n-1-i])
            i=i+1

def create_inverse_QFT(circuit,up_reg,n,with_swaps):
    if with_swaps==1:
        i=0
        while i < ((n-1)/2):
            circuit.swap(up_reg[i], up_reg[n-1-i])
            i=i+1
    i=0
    while i<n:
        circuit.h(up_reg[i])
        if i != n-1:
            j=i+1
            y=i
            while y>=0:
                 if (np.pi)/(pow(2,(j-y))) > 0:
                    circuit.cu1( - (np.pi)/(pow(2,(j-y))) , up_reg[j] , up_reg[y] )
                    y=y-1   
        i=i+1

def getAngle(a, N):
    s=bin(int(a))[2:].zfill(N) 
    angle = 0
    for i in range(0, N):
        if s[N-1-i] == '1': 
            angle += math.pow(2, -(N-i))
    angle *= np.pi
    return angle

def getAngles(a,N):
    s=bin(int(a))[2:].zfill(N) 
    angles=np.zeros([N])
    for i in range(0, N):
        for j in range(i,N):
            if s[j]=='1':
                angles[N-i-1]+=math.pow(2, -(j-i))
        angles[N-i-1]*=np.pi
    return angles

def ccphase(circuit, angle, ctl1, ctl2, tgt):
    circuit.cu1(angle/2,ctl1,tgt)
    circuit.cx(ctl2,ctl1)
    circuit.cu1(-angle/2,ctl1,tgt)
    circuit.cx(ctl2,ctl1)
    circuit.cu1(angle/2,ctl2,tgt)

def phiADD(circuit, q, a, N, inv):
    angle=getAngles(a,N)
    for i in range(0,N):
        if inv==0:
            circuit.u1(angle[i],q[i])
        else:
            circuit.u1(-angle[i],q[i])

def cphiADD(circuit, q, ctl, a, n, inv):
    angle=getAngles(a,n)
    for i in range(0,n):
        if inv==0:
            circuit.cu1(angle[i],ctl,q[i])
        else:
            circuit.cu1(-angle[i],ctl,q[i])
        
def ccphiADD(circuit,q,ctl1,ctl2,a,n,inv):
    angle=getAngles(a,n)
    for i in range(0,n):
        if inv==0:
            ccphase(circuit,angle[i],ctl1,ctl2,q[i])
        else:
            ccphase(circuit,-angle[i],ctl1,ctl2,q[i])
        
def ccphiADDmodN(circuit, q, ctl1, ctl2, aux, a, N, n):
    ccphiADD(circuit, q, ctl1, ctl2, a, n, 0)
    phiADD(circuit, q, N, n, 1)
#     phiADD(circuit, q, a,N, 1)
    create_inverse_QFT(circuit, q, n, 0)
    circuit.cx(q[n-1],aux)
    create_QFT(circuit,q,n,0)
    cphiADD(circuit, q, aux, N, n, 0)
#     cphiADD(circuit, q, aux, a, n, 0)
    
    ccphiADD(circuit, q, ctl1, ctl2, a, n, 1)
    create_inverse_QFT(circuit, q, n, 0)
    circuit.x(q[n-1])
    circuit.cx(q[n-1], aux)
    circuit.x(q[n-1])
    create_QFT(circuit,q,n,0)
    ccphiADD(circuit, q, ctl1, ctl2, a, n, 0)

def ccphiADDmodN_inv(circuit, q, ctl1, ctl2, aux, a, N, n):
    ccphiADD(circuit, q, ctl1, ctl2, a, n, 1)
    create_inverse_QFT(circuit, q, n, 0)
    circuit.x(q[n-1])
    circuit.cx(q[n-1],aux)
    circuit.x(q[n-1])
    create_QFT(circuit, q, n, 0)
    ccphiADD(circuit, q, ctl1, ctl2, a, n, 0)
    cphiADD(circuit, q, aux, N, n, 1)
#     cphiADD(circuit, q, aux, a, n, 1)
    create_inverse_QFT(circuit, q, n, 0)
    circuit.cx(q[n-1], aux)
    create_QFT(circuit, q, n, 0)
    phiADD(circuit, q, N, n, 0)
#     phiADD(circuit, q, a, N, 0)
    ccphiADD(circuit, q, ctl1, ctl2, a, n, 1)

def cMULTmodN(circuit, ctl, q, aux, a, N, n):
    

#     up_reg = QuantumRegister(1, name = "up_reg")
#     down_reg = QuantumRegister(n, name = "down_reg")
#     up_classic = ClassicalRegister(2*n, name="up_classic")
#     c_aux = ClassicalRegister(1, name = "aux_classic")
#     cMULTmodN_circuit = QuantumCircuit(
#         up_reg ,down_reg , aux,up_classic, c_aux, 
#         name=r"${0}^{{{1}^{{{2}}}}} mod{3}$".format(2,2,int(math.log(math.log(a,2),2)), N)
#     )
# create_QFT(cMULTmodN_circuit,aux,n+1,0)
#     for i in range(0, n):
#         ccphiADDmodN(cMULTmodN_circuit, aux, q[i], ctl, aux[n+1], (2**i)*a % N, N, n+1)
# create_inverse_QFT(cMULTmodN_circuit, aux, n+1, 0)
#     for i in range(0, n):
#         circuit.cswap(ctl,q[i],aux[i])
#         cMULTmodN_circuit.cswap(ctl,q[i],aux[i])
#     create_QFT(cMULTmodN_circuit, aux, n+1, 0)
#         ccphiADDmodN_inv(cMULTmodN_circuit, aux, q[i], ctl, aux[n+1], math.pow(2,i)*a_inv % N, N, n+1)
#     create_inverse_QFT(cMULTmodN_circuit, aux, n+1, 0)
#     cMULTmodN_circuit_instruction = cMULTmodN_circuit.to_instruction()
#     circuit.append(cMULTmodN_circuit_instruction, [ctl, *down_reg, *aux])
        
    create_QFT(circuit,aux,n+1,0)
    for i in range(0, n):
        ccphiADDmodN(circuit, aux, q[i], ctl, aux[n+1], (2**i)*a % N, N, n+1)
        
    create_inverse_QFT(circuit, aux, n+1, 0)
    
    for i in range(0, n):
        circuit.cswap(ctl,q[i],aux[i])

    a_inv = modinv(a, N)
    create_QFT(circuit, aux, n+1, 0)

    
    i = n-1
    while i >= 0:
        ccphiADDmodN_inv(circuit, aux, q[i], ctl, aux[n+1], math.pow(2,i)*a_inv % N, N, n+1)
        i -= 1
    create_inverse_QFT(circuit, aux, n+1, 0)
    

def calculate_continued_fraction(b: array.array) -> int:
    # """Calculate the continued fraction of x/T from the current terms of expansion b."""

    x_over_T = 0

    for i in reversed(range(len(b) - 1)):
        x_over_T = 1 / (b[i + 1] + x_over_T)

    x_over_T += b[0]

    frac = fractions.Fraction(x_over_T).limit_denominator()

    return frac.denominator

def get_factors(N: int, a: int, measurement: str) -> Optional[List[int]]:
    # """Apply the continued fractions to find r and the gcd to find the desired factors."""
    x_final = int(measurement, 2)
    #print('In decimal, x_final value for this result is: {}.'.format(x_final))

    if x_final <= 0:
        fail_reason = 'x_final value is <= 0, there are no continued fractions.'
    else:
        fail_reason = None
        #print('Running continued fractions for this case.')

    # Calculate T and x/T
    T_upper = len(measurement)
    T = pow(2, T_upper)
    x_over_T = x_final / T  ## this is our theta

    # Cycle in which each iteration corresponds to putting one more term in the
    # calculation of the Continued Fraction (CF) of x/T

    # Initialize the first values according to CF rule
    i = 0
    b = array.array('i')
    t = array.array('f')

    b.append(math.floor(x_over_T))
    t.append(x_over_T - b[i])

    exponential = 0.0
    while i < N and fail_reason is None:
        # From the 2nd iteration onwards, calculate the new terms of the CF based on the previous terms as the rule suggests
        if i > 0:
            try:
                b_temp = math.floor(1 / t[i - 1])
            except ZeroDivisionError as err:
                b_temp = 0
            b.append(b_temp)
            try:
                t_temp = (1 / t[i - 1]) - b[i]
            except ZeroDivisionError as err:
                t_temp = 0
            t.append(t_temp)  # type: ignore

        # Calculate the denominator of the CF using the known terms
        denominator = calculate_continued_fraction(b)

        # Increment i for next iteration
        i += 1

        if denominator % 2 == 1:
            #print('Odd denominator, will try next iteration of continued fractions.')
            continue

        # Denominator is even, try to get factors of N. Get the exponential a^(r/2)

        if denominator < 1000:
            try:
                exponential = pow(a, denominator / 2)
            except OverflowError as err:
                exponential = 999999999

        # Check if the value is too big or not
        if exponential > 1000000:
            if exponential == 999999999:
                fail_reason = 'OverflowError'
            else:
                fail_reason = 'denominator of continued fraction is too big (> 10^3).'
        else:
            # The value is not too big, get the right values and do the proper gcd()
            putting_plus = int(exponential + 1)
            putting_minus = int(exponential - 1)
            one_factor = math.gcd(putting_plus, N)
            other_factor = math.gcd(putting_minus, N)

            # Check if the factors found are trivial factors or are the desired factors
            if any(factor in {1, N} for factor in (one_factor, other_factor)):
                #print('Found just trivial factors, not good enough.')
                # Check if the number has already been found, (use i - 1 because i was already incremented)
                if t[i - 1] == 0:
                    fail_reason = 'the continued fractions found exactly x_final/(2^(2n)).'
            else:
                return sorted((one_factor, other_factor))

    return None


def process_results(sim_result, circuit, shots, N, a, n):
    counts_result = sim_result.get_counts(circuit)
    total_counts = len(counts_result)
    counts_result_sorted = sorted(counts_result.items(), key=lambda x: x[1], reverse=True)

    counts_result_keys = list(counts_result.keys())
    counts_result_values = list(counts_result.values())

    prob_success=0
    prob_failure=0
    result_successful_counts = 0
    result_failure_counts = 0
    

    for initial_undesired_measurement, frequency in counts_result_sorted:
        measurement = initial_undesired_measurement.split(" ")[1]
        x_value = int(measurement, 2)
        prob_this_result = 100 * frequency/shots
        factors = get_factors(N, a, measurement) 

        if factors:
            prob_success = prob_success + prob_this_result
            result_successful_counts = result_successful_counts + 1

            if factors not in result_factors:
                result_factors.append(factors)
        elif not factors:
            prob_failure = prob_failure + prob_this_result
            result_failure_counts = result_failure_counts + 1
            
    return [result_factors, prob_success, prob_failure, total_counts, result_successful_counts,result_failure_counts]

def my_shor(a,N,shots):
    start_time_number = datetime.now()
    start_time = start_time_number.strftime("%H:%M:%S")
    summary_result = dict()
    validate_min('N', N, 3)
    validate_min('a', a, 2)

    if N < 1 or N % 2 == 0:
        raise ValueError('The input needs to be an odd integer greater than 1.')

    if a >= N or math.gcd(a, N) != 1:
        raise ValueError('The integer a needs to satisfy a < N and gcd(a, N) = 1.')

    n = math.ceil(math.log(N,2))
    global result_factors
    result_factors = []
    tf, b, p = is_power(N, return_decomposition=True)
    if tf:
        print('The input integer is a power: {0}={1}^{2}.'.format(N, b, p))
        result_factors.append(b)



    # """auxilliary quantum register used in addition and multiplication"""
    aux = QuantumRegister(size = n+2, name="aux_reg")

    # """single qubit where the sequential QFT is performed"""
    up_reg = QuantumRegister(1, name = "up_reg")

    down_reg = QuantumRegister(n, name = "down_reg")

    # """classical register where the measured values of the sequential QFT are stored"""
    up_classic = ClassicalRegister(2*n, name="up_classic")

    # """classical bit used to reset the state of the top qubit to 0 if the previous measurement was 1"""
    c_aux = ClassicalRegister(1, name = "aux_classic")

    # """ Create Quantum Circuit """
    circuit = QuantumCircuit(up_reg ,down_reg , aux,up_classic, c_aux)

    circuit.x(down_reg[0])
#     circuit.draw(filename = "shor_semiclassical_QFT_initialization")


    for i in range(0, 2*n):
        circuit.x(up_reg).c_if(c_aux, 1)
        circuit.h(up_reg)
        cMULTmodN(circuit, up_reg[0], down_reg, aux, a**(2**(2*n-1-i)), N, n)
    #     later confirm if this should be up_reg[i] instead of up_reg[0]

        for j in range(0, 2**i):
            circuit.u1(getAngle(j, i), up_reg[0]).c_if(up_classic, j)
            circuit.h(up_reg)
            circuit.measure(up_reg[0], up_classic[i])
            circuit.measure(up_reg[0], c_aux[0])

    # circuit.draw(filename = "shor_semiclassical_QFT_final_circuit")
    circuit.draw()


    qc_compiled = transpile(circuit, backend, optimization_level = 3)
    job_sim_1 = backend.run(qc_compiled, shots=shots)
    sim_result=job_sim_1.result()

#     counts_result = sim_result.get_counts(circuit)
#     len(counts_result)

# measurement_plot = qiskit.visualization.plot_histogram(counts_result,figsize=(20, 12) ,number_to_keep = 30,bar_labels=True, title = "Measurement results from shor_standard_QFT circuit variant" )

# measurement_plot.savefig("shor_semiclassical_QFT_measurement_result")
# measurement_plot

    processed_result = process_results(sim_result, circuit, shots, N, a, n)

    end_time_number = datetime.now()
    end_time = end_time_number.strftime("%H:%M:%S")
    duration = end_time_number - start_time_number
    
    
    print("Current Start Time =", start_time)
    print(processed_result)
    print("Current End Time =", end_time)
    
    circuit_count_ops = circuit.count_ops()
    circuit_decomposed = circuit.decompose()
    circuit_decomposed_count_ops = circuit_decomposed.count_ops()
    qc_compiled_count_ops = qc_compiled.count_ops()
    
    summary_result["num_qubits"] = n
    summary_result["Number(N)"] = N
    summary_result["a"] = a
    summary_result["start_time"] = start_time
    summary_result["end_time"] = end_time
    summary_result["duration"] = duration
    summary_result["result_factors"] = processed_result[0]
    summary_result["prob_success"] = processed_result[1]
    summary_result["prob_failure"] = processed_result[2]
    summary_result["total_counts"] = processed_result[3]
    summary_result["result_successful_counts"] = processed_result[4]
    summary_result["result_failure_counts"] = processed_result[5]
    

    summary_result["circuit_width"] = circuit.width()
    summary_result["circuit_depth"] = circuit.depth()
    summary_result["circuit_size"] = circuit.size() 
    summary_result["circuit_num_nonlocal_gates"] = circuit.num_nonlocal_gates()
    summary_result["circuit_num_ancillas"] = circuit.num_ancillas
    summary_result["circuit_num_clbits"] = circuit.num_clbits
    summary_result["circuit_num_qubits"] = circuit.num_qubits
    summary_result["circuit_num_ancillas"] = circuit.num_ancillas
    summary_result["circuit_num_of_count_ops"] = len(circuit_count_ops)
    summary_result["circuit_num_of_x"] = circuit_count_ops.get('x')
    summary_result["circuit_num_of_measure"] = circuit_count_ops.get('measure')
    summary_result["circuit_num_of_h"] = circuit_count_ops.get('h')
    summary_result["circuit_num_of_cswap"] = circuit_count_ops.get('cswap')
    summary_result["circuit_num_of_swap"] = circuit_count_ops.get('swap')
    summary_result["circuit_num_of_cx"] = circuit_count_ops.get('cx')
    summary_result["circuit_num_of_toffoli"] = circuit_count_ops.get('toffoli')
    summary_result["circuit_num_of_p"] = circuit_count_ops.get('p')
    summary_result["circuit_num_of_t"] = circuit_count_ops.get('t')
    
    summary_result["circuit_decomposed_width"] = circuit_decomposed.width()
    summary_result["circuit_decomposed_depth"] = circuit_decomposed.depth()
    summary_result["circuit_decomposed_size"] = circuit_decomposed.size() 
    summary_result["circuit_decomposed_num_nonlocal_gates"] = circuit_decomposed.num_nonlocal_gates()
    summary_result["circuit_decomposed_num_ancillas"] = circuit_decomposed.num_ancillas
    summary_result["circuit_decomposed_num_clbits"] = circuit_decomposed.num_clbits
    summary_result["circuit_decomposed_num_qubits"] = circuit_decomposed.num_qubits
    summary_result["circuit_decomposed_num_ancillas"] = circuit_decomposed.num_ancillas
    summary_result["circuit_decomposed_num_of_count_ops"] = len(circuit_decomposed_count_ops)
    summary_result["circuit_decomposed_num_of_x"] = circuit_decomposed_count_ops.get('x')
    summary_result["circuit_decomposed_num_of_measure"] = circuit_decomposed_count_ops.get('measure')
    summary_result["circuit_decomposed_num_of_h"] = circuit_decomposed_count_ops.get('h')
    summary_result["circuit_decomposed_num_of_cswap"] = circuit_decomposed_count_ops.get('cswap')
    summary_result["circuit_decomposed_num_of_swap"] = circuit_decomposed_count_ops.get('swap')
    summary_result["circuit_decomposed_num_of_cx"] = circuit_decomposed_count_ops.get('cx')
    summary_result["circuit_decomposed_num_of_toffoli"] = circuit_decomposed_count_ops.get('toffoli')
    summary_result["circuit_decomposed_num_of_p"] = circuit_decomposed_count_ops.get('p')
    summary_result["circuit_decomposed_num_of_t"] = circuit_decomposed_count_ops.get('t')
    
    summary_result["qc_compiled_width"] = qc_compiled.width()
    summary_result["qc_compiled_depth"] = qc_compiled.depth()
    summary_result["qc_compiled_size"] = qc_compiled.size() 
    summary_result["qc_compiled_num_nonlocal_gates"] = qc_compiled.num_nonlocal_gates()
    summary_result["qc_compiled_num_ancillas"] = qc_compiled.num_ancillas
    summary_result["qc_compiled_num_clbits"] = qc_compiled.num_clbits
    summary_result["qc_compiled_num_qubits"] = qc_compiled.num_qubits
    summary_result["qc_compiled_num_ancillas"] = qc_compiled.num_ancillas
    summary_result["qc_compiled_num_of_count_ops"] = len(qc_compiled_count_ops)
    summary_result["qc_compiled_num_of_x"] = qc_compiled_count_ops.get('x')
    summary_result["qc_compiled_num_of_measure"] = qc_compiled_count_ops.get('measure')
    summary_result["qc_compiled_num_of_h"] = qc_compiled_count_ops.get('h')
    summary_result["qc_compiled_num_of_cswap"] = qc_compiled_count_ops.get('cswap')
    summary_result["qc_compiled_num_of_swap"] = qc_compiled_count_ops.get('swap')
    summary_result["qc_compiled_num_of_cx"] = qc_compiled_count_ops.get('cx')
    summary_result["qc_compiled_num_of_toffoli"] = qc_compiled_count_ops.get('toffoli')
    summary_result["qc_compiled_num_of_p"] = qc_compiled_count_ops.get('p')
    summary_result["qc_compiled_num_of_t"] = qc_compiled_count_ops.get('t')
    
    return summary_result

# Run for just a single number N
%%time
N = 21
shots = 1024
global result_factors

all_summary_result_temp = []
for random_a in range(2, N):
    if math.gcd(random_a,N) > 1:
        continue
    a = random_a
    summary_result = my_shor(a,N,shots)
    print("Finished running for a = {} and N = {}\n".format(a, N))
    all_summary_result_temp.append(summary_result)
    summary_result_list = []
    for key, value in summary_result.items():
        summary_result_list.append([key,value])
    summary_result_list
    
    with open("a({0})_N({1})_semiclassical.csv".format(a, N), 'a') as myfile:
        write = csv.writer(myfile)
        #write.writerow(fields)
        write.writerows(summary_result_list)

all_summary_result_temp

# Run for many numbers N.
%%time
shots = 1024
global result_factors
all_summary_result = []
for N in [15, 21, 33, 35, 39, 51, 55, 57]:
    for a in range(2, N):
        if math.gcd(a,N) > 1:
            continue
        print("Beginning running for a = {} and N = {}".format(a, N))
        summary_result = my_shor(a,N,shots)
        print("Finished running for a = {} and N = {}\n\n".format(a, N))
        all_summary_result.append(summary_result)

all_summary_result


%qiskit_copyright