from pprint import pprint

# VERBOSE = True # Debug flag
VERBOSE = False

"""
Utility module provides common operations when working over Z_n, such
as:
- fast modular exponentiation
- finding generators of a group
- naiively computing the discrete logarithm
- applying Fermat's primality test
"""

def vprint(*args):
    if not VERBOSE:
        return
    if not args:
        print(); return
    for arg in args:
        print (arg)

def times_table(n):
    """
    Outputs a 2D matrix of modular multiplication in Zn
    """
    return [[(i * j) % n for i in range(n)] for j in range(n)]

def plus_table(n):
    """
    Outputs a 2D matrix of modular addition in Zn
    """
    return [[(i + j) % n for i in range(n)] for j in range(n)]

def zn_pow(x, y, n):
    """
    Computes x^y in z_n
    """
    if y < 0:
        y = abs(y)
        x = inverse_in_zn(x, n)
    product = 1
    for i in range(y):
        product = (product * x) % n
        vprint("{}^{}={}".format(x, i, product))
    return product

def zn_star(n):
    """
    Generates Z*_n = {x | 0 < x < n and gcd(x, n) = 1}.
    This structure is the sort of group used for "textbook RSA"
    """
    return [x for x in range(1, n) if coprime(x, n)]

def coprime(a,b):
  return gcd(a, b) == 1

def elts_generated_by_g_in_zn(g, n):
    """
    Produces the cyclic subgroup generated by the element g in Zn.
    Note: Performs checks to ensure that g is the representative
    element of its coset.
    """
    product = 1
    elts_generated = []
    g = g % n # Ensure that g is a representative element
    for i in range(n):
        vprint("$${}*{} = {} \equiv {}$$".format(
            product, g, product*g, (product*g) % n))
        product = (product * g) % n
        elts_generated.append(product)
        if product == 1:
            return elts_generated
    return elts_generated

def is_generator_in_zn(g, n):
    """Checks if g is a generator of z_n*"""
    elts_generated = elts_generated_by_g_in_zn(g, n)
    vprint("g:{}, elts:{}\n".format(g, elts_generated))
    return all([i in elts_generated for i in zn_star(n)])

def generators(n):
    """Returns a list containing the generator(s) of z_n"""
    return [g for g in zn_star(n)
            if is_generator_in_zn(g, n)]

def division_algo(a, b):
    """
    Computes a | b

    Finds q, r such that
    0 <= r < q
    a = bq + r
    """
    return a / b, a % b

def gcd(a, b):
    """
    Computes gcd(a, b)
    Precondition: a >= b
    """
    if (a < b):
      a, b = b, a

    if a == b:
        return a

    q, r = division_algo(a, b)
    vprint("{} = {}*{} + {};".format(a, q, b, r))

    if r == 0:
        return b
    return gcd(b, r)

def extended_euclidean(a, b):
    """
    Find the values which satisfy Bezout's identity for a and b, i.e.
    finds X, Y, d such that aX + bY = d, where d = gcd(a, b)
    """
    vprint("a: {}, b: {}".format(a, b))
    assert(a >= b)
    if a == b:
        x = 1
        y = 0
        d = a
        vprint("Base case: a == b")
        vprint("x: {}, y: {}, d: {}".format(x, y, d))
        return x, y, d

    q, r = division_algo(a, b)
    vprint(" a = q * b + r")
    vprint("{} = {} * {} + {}".format(a, q, b, r))

    if r == 0:
        x = 0
        y = 1
        d = b
        vprint("Base case r == 0")
        vprint("x: {}, y: {}, d: {}".format(x, y, d))
        return x, y, d

    vprint("Entering recursive call")
    vprint("extended_euclidean({}, {})".format(b, r))
    x1, y1, d1 = extended_euclidean(b, r)
    vprint("Returning from recursive call")
    vprint("x1: {}, y1: {}, d1: {}".format(x1, y1, d1))
    x = y1
    y = x1 - q*y1
    d = d1
    vprint("x: {}, y: {}, d: {}".format(x, y, d))
    return x, y, d

def is_prime_by_fermat_test(n, a):
    """
    Test if n is prime, with witness a.
    Output True if n is prime, False otherwise
    """
    output = zn_pow(a, n-1, n) # a^{n-1}, mod n
    if output == 1:
        vprint("prime")
        return True
    vprint("composite")
    return False

def inverse_in_zn(g, n):
    """
    Find inverse of g in z_n

    NOTE: extended_euclidean(n, g) returns X, Y, d such that
        x*n + y*g = gcd(n, g)

    If n is prime, then gcd(n, g) is guaranteed to be 1, so
        x*n + y*g = gcd(n, g) = 1

    Furthermore, since we're in Z_n, (x*n) % n = 0. So
        x*n + y*g = 1
         0  + y*g = 1
    i.e. y is the multiplicative inverse of g in Z_n
    """
    assert(n >= g)
    x, y, d = extended_euclidean(n, g)
    if y < 0:
        return y + n
    return y

def dlog(g, h, n):
    """
    Finds dlog(h) base-g in Z_n, i.e.
    given generator g, finds x such that g^x = h in Z_n
    """
    elts_generated = elts_generated_by_g_in_zn(g, n)
    # If this assert fails, then g wasn't actually a generator
    assert(h in elts_generated)
    return elts_generated.index(h)

def int2bin(i):
    """Converts i to its binary string representation"""
    if i == 0: return "0"
    s = ''
    while i:
        if i & 1 == 1:
            s = "1" + s
        else:
            s = "0" + s
        i /= 2
    return s

def fast_exp(a, x, n):
    """Speedily computes a^x (mod n)"""
    x_2 = int2bin(x)
    vprint("{} = [{}]_2".format(str(x), x_2))
    powers = [a % n]
    vprint("{}^(2^0) = {}^1 = {} \\equiv {}".format(a, a, a, (a % n)))
    i = 1
    while i < len(x_2):
        # This (hilariously ugly) print statement prints the
        # intermediary operations in a format that can be easily
        # exported to LaTeX. TODO: Split it up into sane chunks.
        vprint("{}^{{ {}^{} }} = {}^{{ {} }} = {}^{{ {} }} * {}^{{ {} }} = {}*{} = {} \\equiv {}".format(
            a, 2, i,
            a, pow(2, i),
            a, pow(2, i-1),
            a, pow(2, i-1),
            powers[-1], powers[-1],
            powers[-1] * powers[-1],
            (powers[-1] * powers[-1]) % n))
        next_power = (powers[-1] * powers[-1]) % n
        powers.append(next_power)
        i += 1

    vprint("{}^{{ {} }} = ...".format(a, x))
    rpowers = list(reversed(powers))
    prod = 1
    i = 0
    while i < len(x_2):
        bit = x_2[i]
        power = rpowers[i]
        if bit == "1":
            vprint("* {} \t== {}^{{ 2^{{ {} }} }}\n".format(power, a, len(x_2) - i - 1))
            prod *= power
        i += 1
    result = prod % n
    vprint("= {} \\equiv {}".format(prod, result))
    return result
