def genDLP(size, opt=0):
    p = next_prime(2**size)
    q = ZZ((p-1)/2)

    while not q.is_prime():
        p = next_prime(p)
        q = ZZ((p-1)/2)

    R = IntegerModRing(p)
    g = R.multiplicative_generator()
    x = ZZ.random_element(2,p-2)
    h = g**x
    
    return h, g


def chooseN(num,param=2):
    """
    Helper function for BSGS. Chooses N
    """
    
    return ceil(RR(num)**(1/param))

def BSGS(h,g,p=None,Nfunc=None,NfuncParam=None):
    # Parse optional value for p and extract data if none exists
    if p == None:
        R = g.parent()
        order = R.order() - 1
    else:
        R = IntegerModRing(p)
        order = p-1
        g = R(g)
        h = R(h)

    # Find value for N'. Default is standard ceil(sqrt(order)), but allows
    # different choices for time/memory trade-off
    if Nfunc == None:
        Nprime1 = ceil(sqrt(order))
        Nprime2 = Nprime1
    else:
        Nprime1 = Nfunc(order, NfuncParam)
        Nprime2 = ceil(order / Nprime1)
    


    print "Order of group\t\t: ", order, "\nSize of list\t\t: ", Nprime1, "\nSize of loop\t\t: ", Nprime2, "\nTotal coverage of loops\t: ", Nprime1*Nprime2, "\nCovers whole group\t: ", Nprime1*Nprime2 >= order 
    # Note that we use a dictionary as this is stored as a hashtable,
    # hence constant time lookup
    generatorPowers = {g**j : j for j in xrange(Nprime1)}

    
    
    d = g**(-Nprime1)
    c = h
    
    for i in xrange(0,Nprime2):
        if c in generatorPowers:
            j = generatorPowers[c]
            return Nprime1*i + j
        else:
            c *= d

# Note that if we write compactly, we only use 11 lines of code
def BSGS_compact(h,g):
    Nprime = ceil(sqrt(g.parent().order() - 1))
    generatorPowers = {g**j : j for j in xrange(Nprime)}
    d = g**(-Nprime)
    c = h   
    for i in xrange(0,Nprime):
        if c in generatorPowers:
            j = generatorPowers[c]
            return Nprime*i + j
        else:
            c *= d

import time 

h,g = genDLP(30)

t1 = time.time()
answer =  BSGS(h,g,Nfunc=chooseN, NfuncParam=2)
t2 = time.time()
print "\n"
print "BSGS algorithm answer:", answer

print "Compact BSGS answer  :", BSGS_compact(h,g) 
t3 = time.time()
print "Actual DLog of h is  :", discrete_log(h,g)
t4 = time.time()
print ""
print "BSGS time              :", t2 - t1
print "DLog function time     :", t4 - t3

time.time()

# We compare the running times of the Baby-Step Giant-Step algorithm for different sizes of N'

# Note that we essentially have a trade-off between the size of the list and the size of the loop

# Choosing N' to be either extreme essentially collapses the algorithm to a brute-force search via a creation of a list or brute force via a loop search

# We skip N' = N as this is brute force create of a list for SageMath cloud to efficiently run

#t1 = time.time()
#answer =  BSGS(h,g,Nfunc=chooseN, NfuncParam=1)
t2 = time.time()
#print ""

#print "Time taken for BSGS with N' = N:", t2-t1
print ""
answer =  BSGS(h,g,Nfunc=chooseN, NfuncParam=1.5)
t3 = time.time()
print ""

print "BSGS time with N' = N^(2/3):", t3-t2
print ""

answer =  BSGS(h,g,Nfunc=chooseN, NfuncParam=2)
t4 = time.time()
print ""

print "Time taken for BSGS with N' = N^(1/2):", t4-t3
print ""

answer =  BSGS(h,g,Nfunc=chooseN, NfuncParam=3)
t5 = time.time()
print "Time taken for BSGS with N' = N^(1/3):", t5-t4
print ""

answer =  BSGS(h,g,Nfunc=chooseN, NfuncParam=5)
t6 = time.time()
print ""

print "Time taken for BSGS with N' = N^(1/5):", t6-t5
print ""

answer =  BSGS(h,g,Nfunc=chooseN, NfuncParam=10)
t7 = time.time()
print ""

print "Time taken for BSGS with N' = N^(1/10):", t7-t6
print ""