#! /usr/bin/python

from modpower import *
from xgcd import *

def PohligHellman(beta, alpha, p):
    ''' PohligHellman(beta,alpha,p) 
     Given beta=alpha^x modp,
     this function solves for x.
     Fast if p-1 has only small prime factors
    '''
    pminus1=p-1
    smallprime=2
    congruenceList=[]
    while pminus1!=1:
        r=0  # will be the power of the smallPrime in p-1
        if (pminus1%smallprime==0):    # factor p-1        
            while pminus1%smallprime==0:
                pminus1=pminus1/smallprime
                r=r+1
            congruenceList.append(getX(beta,alpha,p,smallprime,r))
        smallprime=smallprime+1
    (x,m)=ChineseRemainder(congruenceList)
    print "Given", beta,"=", alpha,"^x mod",p, "\n","x=", x
    assert pow(alpha, x, p) == beta

def getX(beta, alpha, p, q, r):
    ''' return (x, q**r) with (p-1)/q**r = k, 0 <= x < q**r, os
    beta^(x*k) = alpha^k mod p
    '''
    modqpowerr=q**r
    pminus1=p-1
    bCurrent=beta
    xFinal=0  # returns x=x0+x1q+x2q^2+...+xiq^i with 0<=xi<q-1
    alphaRaisedModp=modpower(alpha, pminus1/q, p)
    for i in range(0,r):
        betaRaisedModp=modpower(bCurrent, pminus1/q**(i+1), p)
        xCurrent=0
        while (modpower(alphaRaisedModp, xCurrent, p)!=betaRaisedModp):
            xCurrent=xCurrent+1  # solve lower degree log problem
        xFinal=(xFinal+xCurrent*q**i)% modqpowerr
        #now we calculate the next beta
        bCurrent=((bCurrent*xgcd(modpower(alpha, xCurrent*q**i, p),p)[1])%p)
    return (xFinal,modqpowerr)

def ChineseRemainder(congruenceList): # neater than my original one! Why works?
    ''' takes a system of congruences a mod m or b mod n in the form
    [(a,m),(b,n),....] where m and n are different primes to some power
    returns solution x mod mn....'''
    (a, m)=congruenceList[0]
    for i in range(1, len(congruenceList)):
        (b,p)= congruenceList[i]
        k=((b-a)*xgcd(m,p)[1])%p #congruences have gcd(m,p) = 1 so inverse exists for mmodp
        a=(a+m*k)%(m*p)# joining a mod m and b mod p gives a mod(mp)
        m=m*p # mod mp
    return (a,m)

PohligHellman(3, 5, 2017)
PohligHellman(95, 37, 2017)
PohligHellman(19, 95, 3001)
PohligHellman(7531, 6, 8101)

def testGen(a, p):  # want generator base for discrete log (and Pohlig-Hellman)
    '''True if a generates Fp* for prime p.'''
    b = a
    for i in range(1, p-1):
        if b == 1:
            return False
        b = b*a % p
    assert b == 1
    return True