def genDLP(size):
    # Ensures p is safe-prime.
    p = next_prime(size)
    q = ZZ((p-1)/2)

    if 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 chooseB(p,Bparam):
    return Bparam

def indexCalculus(h,g,p=None, Bparam=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)
    
    orderRing = IntegerModRing(order)
    
    B = chooseB(p,Bparam)
    factorBasis = [i for i in primes(B+1)]
    k = len(factorBasis)
    #print k
    M = matrix(ZZ,generateRelation(g,factorBasis))

    # Grow matrix with extra relations till reach k*k matrix of full rank
    while M.rank() < k + 1:
        relation = generateRelation(g,factorBasis)
        if checkInvertibleRelation(relation[:-1], order):
            pass
        else:
            #print relation
            continue
        temp = matrix(ZZ,relation)
        tempM = M.stack(temp)
        if tempM.rank() > M.rank():
            M = tempM
        else:
            pass

    M.echelonize()
    M = M.change_ring(IntegerModRing(order))
    logP_values= M[:,:-1].solve_right(M[:,-1])
    
    #logP_values = M[:-1,-1]
    print "searching different relations"
    x = 0
    while g**x != h:
        dLogExponentRelation, randomNumber = searchRelation(g,h,factorBasis)
        x = (matrix(dLogExponentRelation)*logP_values)[0,0] - randomNumber
    
    return x

def checkInvertibleRelation(relations,order):
    #print relations
    if all([(gcd(ZZ(val),order) == 1) or val == 0 for val in relations[:]]):
        return True
    else:
        return False

def searchRelation(g,h,factorBasis):
    order = g.parent().order()
    smooth = False
    while smooth == False:
        k = ZZ.random_element(0,order)
        r = g**k * h
        smooth = smartFactor(ZZ(r),factorBasis)
    return smooth,k
    

def generateRelation(g,factorBasis):
    ans = False
    order = g.parent().order() - 1
    while ans == False:
        k = ZZ.random_element(1, order + 1)
        r = g**k
        r = ZZ(r)
        ans = smartFactor(r,factorBasis)
        
    return ans + [k]

def smartFactor(n, factorBasis):
    l = len(factorBasis)
    relationList = [0 for _ in range(l)]
    for i in range(0,l):
        while (n % factorBasis[i]) == 0:
            n = n // factorBasis[i]
            relationList[i] += 1
    if n == 1:
        return relationList
    else:
        return False

import time
#Testing 
# Create Dlog problem
h,g = genDLP(2**20)
# Set order of group
order = g.parent().order() - 1


# find Dlog for h = g^x with an upper bound on the size of the largest prime = 10
t0 = time.time()
x  = indexCalculus(h,g,Bparam = 10); g**x; h 
t1 = time.time() - t0
print "Found correct discrete log: ", g**x == h, ". Using a bound of: ", 10, " took ", t1, " seconds."


# find Dlog for h = g^x with an upper bound on the size of the largest prime = 25
t0 = time.time()
x  = indexCalculus(h,g,Bparam = 25); g**x; h 
t1 = time.time() - t0
print "Found correct discrete log: ", g**x == h, ". Using a bound of: ", 25, " took ", t1, " seconds."

# find Dlog for h = g^x with an upper bound on the size of the largest prime = 50
t0 = time.time()
x  = indexCalculus(h,g,Bparam = 50); g**x; h 
t1 = time.time() - t0
print "Found correct discrete log: ", g**x == h, ". Using a bound of: ", 50, " took ", t1, " seconds."

# find Dlog for h = g^x with an upper bound on the size of the largest prime = 100
t0 = time.time()
x  = indexCalculus(h,g,Bparam = 100); g**x; h 
t1 = time.time() - t0
print "Found correct discrete log: ", g**x == h, ". Using a bound of: ", 100, " took ", t1, " seconds."

# find Dlog for h = g^x with an upper bound on the size of the largest prime = 250
t0 = time.time()
x  = indexCalculus(h,g,Bparam = 250); g**x; h 
t1 = time.time() - t0
print "Found correct discrete log: ", g**x == h, ". Using a bound of: ", 250, " took ", t1, " seconds."


# find Dlog for h = g^x with an upper bound on the size of the largest prime = 500
t0 = time.time()
x  = indexCalculus(h,g,Bparam = 500); g**x; h 
t1 = time.time() - t0

print "Found correct discrete log: ", g**x == h, ". Using a bound of: ", 500, " took ", t1, " seconds."