from __future__ import nested_scopes
import math, numarray, random
from numarray import where


class QuState:
    
    def __init__(self, matrix):
        self.matrix = matrix
        
    def measure_away(self, dim):
        state = self.matrix
        dimensions = len(state.shape)
        norm = (state * state.conjugate()).astype(numarray.Float64)
        for i in range(dimensions-1, -1, -1):
            if i != dim:
                norm = numarray.sum(norm, i)
        assert len(norm.shape) == 1
        cumnorm = numarray.cumsum(norm)
        p = random.random() * cumnorm[-1]
        result = numarray.searchsorted(cumnorm, p)
        self.matrix = numarray.take(state, (result,), axis=dim)
        return result


class QuReg:
    def __init__(self, initialstate):
        assert len(initialstate.shape) == 1
        self.state = QuState(initialstate)
        self.sdim = 0
    def __int__(self):
        base = numarray.arange(self.state.matrix.shape[self.sdim])
        result = self.state.measure_away(self.sdim)
        self.state = QuState(base == result)
        self.sdim = 0
        return result


class QuOperator:
    
    def __init__(self, f):
        self.f = f
        self.matrixcache = {}
    
    def __call__(self, a):
        state = a.state.matrix
        size = state.shape[a.sdim]
        try:
            matrix = self.matrixcache[size]
        except KeyError:
            matrix = numarray.fromfunction(self.f, (size, size))
            self.matrixcache[size] = matrix
        assert a.sdim == 0
        state = numarray.innerproduct(state, matrix)
        a.state.matrix = state


# ____________________________________________________________
#
# Discrete logarithm:
#
#   pow(21, ?, 31) == 22


def log2(n):
    return math.log(n) / math.log(2)


def pow1(base, exponent, modulo):
    "Python version of pow().  Works with numarrays as well."
    
    # compute [n, n**2, n**4, n**8, n**16, ...]
    precision = int(log2(modulo))+1
    n = base
    powers = [n]
    for i in range(1, precision):
        n = (n**2) % modulo
        powers.append(n)
    
    # the result is the product of some items in that list,
    # as selected by the exponent
    result = 1
    for i in range(precision):
        factor = where(exponent & (2**i), powers[i], 1)
        result = (result * factor) % modulo
    return result


class DiscreteLogarithm:
    
    def __init__(self, base, prime, answer):
        self.base = base
        self.prime = prime
        self.answer = answer

    def is_correct_exponent(self, n):
        return pow1(self.base, n, self.prime) == self.answer


def size(matrix):
    return matrix.shape[0]

diffuse = QuOperator(lambda input, output:
                     (2.0 / size(input)) - (input == output))


def grover_search(predicate, upperbound):
    rotate = QuOperator(lambda input, output:
                        where(input == output,
                              where(predicate(input), 1, -1),
                              0))
    r = QuReg(numarray.ones((upperbound,)))
    optimalstepcount = int(math.pi/4 * math.sqrt(upperbound))
    for i in range(optimalstepcount):
        rotate(r)
        diffuse(r)
    return int(r)


print grover_search(DiscreteLogarithm(2, 1291, 1165).is_correct_exponent, 1291)