# https://qiskit.org/documentation/tutorials/circuits/3_summary_of_quantum_operations.html
# https://github.com/olekssy/qiskit-shors/blob/master/shors.py
# https://quantumcomputinguk.org/tutorials/shors-algorithm-with-code

from math import pow,gcd,floor,isinf,ceil,log

from numpy import zeros,pi

from fractions import Fraction
from array import array

from qiskit import QuantumRegister,ClassicalRegister,QuantumCircuit
from qiskit import IBMQ

from quantum_runner import setup_IBM,quantum_simulate

def check_if_power(N):
    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 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 (pi)/(pow(2,(i-j))) > 0:
                circuit.crz( (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 (pi)/(pow(2,(j-y))) > 0:
                    circuit.crz( - (pi)/(pow(2,(j-y))) , up_reg[j] , up_reg[y] )


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 get_factors(x_value,t_upper,N,a):

    if x_value<=0:
        print('x_value is <= 0, there are no continued fractions\n')
        return False

#    print('Running continued fractions for this case\n')

    T = pow(2,t_upper)

    x_over_T = x_value/T

    i=0
    b = array('i')
    t = array('f')

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

    while i>=0:

        if i>0:
            b.append( floor( 1 / (t[i-1]) ) )
            t.append( ( 1 / (t[i-1]) ) - b[i] )


        aux = 0
        j=i
        while j>0:
            aux = 1 / ( b[j] + aux )
            j = j-1

        aux = aux + b[0]

        frac = Fraction(aux).limit_denominator()
        den=frac.denominator

        print('Approximation number {0} of continued fractions:'.format(i+1))
        print("Numerator:{0} \t\t Denominator: {1}\n".format(frac.numerator,frac.denominator))

        """ Increment i for next iteration """
        i=i+1

        if (den%2) == 1:
            if i>=15:
                print('Returning because have already done too much tries')
                return False
            print('Odd denominator, will try next iteration of continued fractions\n')
            continue

        """ If denominator even, try to get factors of N """

        """ Get the exponential a^(r/2) """

        exponential = 0

        if den<1000:
            exponential=pow(a , (den/2))

        """ Check if the value is too big or not """
        if isinf(exponential)==1 or exponential>1000000000:
            print('Denominator of continued fraction is too big!\n')
            aux_out = input('Input number 1 if you want to continue searching, other if you do not: ')
            if aux_out != '1':
                return False
            else:
                continue

        putting_plus = int(exponential + 1)

        putting_minus = int(exponential - 1)

        one_factor = gcd(putting_plus,N)
        other_factor = gcd(putting_minus,N)

        """ Check if the factors found are trivial factors or are the desired
        factors """

        if one_factor==1 or one_factor==N or other_factor==1 or other_factor==N:
            print('Found just trivial factors, not good enough\n')
            if t[i-1]==0:
                print('The continued fractions found exactly x_final/(2^(2n)) , leaving funtion\n')
                return False
            else:
                """ Return if already too much tries and numbers are huge """
                print('Returning because have already done too many tries\n')
                return False
        else:
            print('The factors of {0} are {1} and {2}\n'.format(N,one_factor,other_factor))
            print('Found the desired factors!\n')
            return True

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

def ccphase(circuit,angle,ctl1,ctl2,tgt):
    circuit.crz(angle/2,ctl1,tgt)
    circuit.cx(ctl2,ctl1)
    circuit.crz(-angle/2,ctl1,tgt)
    circuit.cx(ctl2,ctl1)
    circuit.crz(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.p(angle[i],q[i])
        else:
            circuit.p(-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.crz(angle[i],ctl,q[i])
        else:
            circuit.crz(-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)
    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)

    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)
    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)
    ccphiADD(circuit, q, ctl1, ctl2, a, n, 1)

def cMULTmodN(circuit, ctl, q, aux, a, N, n):
    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], pow(2,i)*a_inv % N, N, n+1)
        i -= 1
    create_inverse_QFT(circuit, aux, n+1, 0)

def get_value_a(N):
    a=2
    while gcd(a,N)!=1:
        a=a+1
    return a



def factorize_number(N):
    if N==1 or N==0:
        print('Please put an N different from 0 and from 1')
        exit()

    """ Check if N is even """

    if (N%2)==0:
        print('N is even, so does not make sense!')
        exit()

    """ Check if N can be put in N=p^q, p>1, q>=2 """

    """ Try all numbers for p: from 2 to sqrt(N) """
    if check_if_power(N)==True:
       exit()

#    print('Not an easy case, using the quantum circuit is necessary\n')

    setup_IBM()

    """ Get an integer a that is coprime with N """
    a = get_value_a(N)

#    """ If user wants to force some values, he can do that here, please make sure to update the print and that N and a are coprime"""
#    print('Forcing N=15 and a=4 because its the fastest case, please read top of source file for more info')
#    N=15
#    a=4

    print("help !!!")
    """ Get n value used in Shor's algorithm, to know how many qubits are used """
    n = ceil(log(N,2))

#    print('Total number of qubits used: {0}\n'.format(4*n+2))

    aux = QuantumRegister(n+2)
    up_reg = QuantumRegister(2*n)
    down_reg = QuantumRegister(n)
    up_classic = ClassicalRegister(2*n)
    circuit = QuantumCircuit(down_reg , up_reg , aux, up_classic)

    circuit.h(up_reg)
    circuit.x(down_reg[0])

    """ Apply the multiplication gates as showed in the report in order to create the exponentiation """
    for i in range(0, 2*n):
        cMULTmodN(circuit, up_reg[i], down_reg, aux, int(pow(a, pow(2, i))), N, n)

    """ Apply inverse QFT """
    create_inverse_QFT(circuit, up_reg, 2*n ,1)

    """ Measure the top qubits, to get x value"""
    circuit.measure(up_reg,up_classic)

    """ Simulate the created Quantum Circuit """
#    print(transpile(circuit,BasicAer.get_backend('qasm_simulator')))
    #simulation = execute(circuit, backend=BasicAer.get_backend('qasm_simulator'),shots=number_shots)
    provider = IBMQ.get_provider(hub='ibm-q', group='open', project='main')
    simulation = quantum_simulate(circuit,provider.get_backend('ibmq_lima'),True)

    i=0
    while i < len(simulation):
        print('Result \"{0}\" happened {1} times out of {2}'.format(list(simulation.keys())[i],list(simulation.values())[i]))
        i=i+1

    """ An empty print just to have a good display in terminal """
    print(' ')

    """ Initialize this variable """
    # prob_success=0

    """ For each simulation result, print proper info to user and try to calculate the factors of N"""
  #  i=0
  #  while i < len(counts_result):

        #""" Get the x_value from the final state qubits """
       # output_desired = list(sim_result.get_counts().keys())[i]
       # x_value = int(output_desired, 2)
       # prob_this_result = 100 * ( int( list(sim_result.get_counts().values())[i] ) ) / (number_shots)

        #print("------> Analysing result {0}. This result happened in {1:.4f} % of all cases\n".format(output_desired,prob_this_result))

        #""" Print the final x_value to user """
        # print('In decimal, x_final value for this result is: {0}\n'.format(x_value))

        #""" Get the factors using the x value obtained """
        # success=get_factors(int(x_value),int(2*n),int(N),int(a))

        #if success==True:
         #   prob_success = prob_success + prob_this_result

        #i=i+1

#    print("\nUsing a={0}, found the factors of N={1} in {2:.4f} % of the cases\n".format(a,N,prob_success))

N = 33
factorize_number(N)