""" lfsr.py
    CJ Odenthal
    22/06/2005

    This module implements some functions dealing with Linear
    Feedback Shift Registers. For details see Section 5.6 of
    "Basic Methods of Cryptography" by J.C.A. van der Lubbe,
    Cambridge University Press (1994).
"""

from operator import xor

def lfsr(coefficients,state=[]):
    """ This implements a LFSR. The inputs are the 'coefficients'
        of the feedback function and the initial 'state' of the
        register.
    """
    # If 'state' is longer than 'coefficients' we'll truncate it.
    # If 'state' is shorter than 'coefficients' we'll pad it with 1s.
    m = len(coefficients)
    state = state[:m]
    state = state + [1]*(m-len(state))

    while 1:
        state.append(dot(coefficients,state))
        yield state.pop(0)

def fcsr(coefficients,state=[],carry=0):
    """ This implements a Feedback Carry Shift Register. The inputs
        are the 'coefficients' of the feedback function, the initial
        'state' of the register, and the initial 'carry' (an integer).
    """
    # If 'state' is longer than 'coefficients' we'll truncate it.
    # If 'state' is shorter than 'coefficients' we'll pad it with 1s.
    m = len(coefficients)
    state = state[:m]
    state = state + [1]*(m-len(state))

    while 1:
        comp = carry + sum([x&y for (x,y) in zip(coefficients,state)])
        state.append(comp % 2)
        carry = comp // 2
        yield state.pop(0)


def berlekamp_massey(S):
    """ This implements the Berlekamp-Massey algorithm whose purpose
        is to find a Linear Feedback Shift Register (with the smallest
        possible register length) that generates the given sequence 'S'.
        The assumption is that S is, in fact, generated by some LFSR.
        This can always be accomplished by a LFSR with register size
        equal to the length of S (initial register contents S, feedback
        function f(x_0,...,x_{m-1}) = x_0) but we hope to do better.

        A 'generator' is returned  that 'yields' the current LFSR at
        each call of 'next'. At the kth call to 'next' the 'yielded'
        LFSR will generate the first k bits of 'S'.

        The generator's output is a list [c_0, c_1,..., c_{m-1}] that
        contains the coefficients of the feedback function. The way
        the feedback function operates is as follows: if the internal
        state of the LFSR is (x_0,x_1,...,x_{m-1}) then the output bit
        is x_0, the register contents are shifted to the left, and the
        new x_{m-1} = c_0x_0 + c_1x_1 +...+c_{m-1}x_{m-1}.

        The generating function G(x)=s_0 + s_1x^1 + s_2x^2 + ....  of
        a LFSR is rational and (when written in lowest terms) the de-
        nominator f(x) is called the characteristic polynomial of the
        LFSR. Here we have f(x) = c_0x^m + c_1x^{m-1} +...+ c_{m-1}x + 1.
    """
    # Allow for an input string instead of a list or tuple.
    if type(S)==type("string"):
        S = map(int,tuple(S))

    # Initialize 'm': I thought that 'm' was going to be the degree
    # of the characteristic polynomial 'f' of the current LFSR, or
    # at least an upper bound on the degree of 'f'. But evidently I
    # was wrong. I don't know what 'm' represents. In any case, it
    # starts out as 0.
    m = 0

    # Initialize the variables 'N' and 'f': 'N' is the index of the
    # current element of the sequence 'S' that we are considering,
    # while 'f' is characteristic polynomial of the current LFSR.
    # 'N' starts out as the integer 0 while 'f' starts out as the
    # constant polynomial 'One' - a global variable defined as Poly([1]).
    N,f = 0,One

    # Initialize the variables 'N0' and 'f0': these are usually the
    # values that 'N' and 'f' had the last time that 'm' was changed.
    # The polynomial 'f0' starts out with the same value as 'f' (One);
    # however 'N0' starts out as the integer -1.
    N0,f0 = -1,One

    n = len(S)
    while N < n:
        # Does the current LFSR compute the next entry in the
        # sequence correctly? If not we need to update the LFSR.
        if S[N] != f.dot(S[:N]):

            # If N is small enough we can get away with just
            # updating the coefficients in the LFSR. Note that
            # the 'X' occuring below is a global variable and
            # is the indeterminant in the polynomials defined
            # by the class 'Poly'. That is, X=Poly([1,0]).
            if 2*m>N:
                f = f + f0 * X**(N-N0)

            # Otherwise we'll have to update everything.
            else:
                f0, f = f, f + f0 * X**(N-N0)
                N0 = N
                m = N+1-m

        # Go to the the next element in the sequence.
        N += 1

        # 'yield' the coefficients [c_0, ... , c_{m-1}] of the
        # current LFSR's feedback function.
        yield f[:-1]

    # We've processed the entire sequence 'S', so announce
    # that we're done.
    print "The LFSR is constructed."

    # 'yield' the final LFSR coefficients once again.
    yield f[:-1]

def run_berlekamp_massey(S):
    """ This runs the Berlekamp-Massey algorithm to completion,
        and returns the final LFSR.
    """
    L = berlekamp_massey(S)
    while 1:
        try:
            R=L.next()
        except StopIteration:
            return R

class Poly(list):
    """ A class for polynomials over GF(2). The coefficients are
        stored as a list [c_0,c_1,...,c_m] where the polynomial
        is c_0*x**m + c_1*x**(m-1) + ... + c_{m-1}*x + c_m.
    """

    def __str__(self):
        L = ["x^%s" % (k,) for (k,x) in enumerate(self[::-1]) if x==1 ]
        L.reverse()
        return "+".join(L)

    def __add__(self,other):
        if len(self)<len(other):
            return other + self

        else:
            k = len(self) - len(other)
            new = [a^b for (a,b) in zip(self,([0]*k)+other)]
            try:
                while new[0] == 0:
                    new.pop(0)
            except IndexError:
                new = []

            return Poly(new)

    def dot(self,S):
        """ This returns the sum of c_k*S_{N-k} where N+1 = len(S),
            and k runs from 0 to m-1 (m is the degree of self).
        """
        m = len(self)
        N = len(S)
        C = self[:-1]
        S = S[N-m+1:]
        return dot(C,S)

    def __mul__(self,other):
        deg = len(self) + len(other) - 2
        prod = [0]*(deg + 1)
        for k,x in enumerate(self):
            for j,y in enumerate(other):
                prod[k+j] ^= (x&y)
        return Poly(prod)

    def __pow__(self,ex):
        if ex==0:
            return Poly([1])
        else:
            return self*(self**(ex-1))

def dot(x,y):
    """ This returns the 'dot' product of the two binary lists
        'x' and 'y'. If either list is empty '0' is returned.
    """
    return reduce(xor,[a&b for (a,b) in zip(x,y)],0)

# Some stock polynomials to use.
X = Poly([1,0])
One =  Poly([1])

testS = [1,1,1,0,0,1,0,1,1,1]

def tunny(m,n):
    import random
    M=[]
    N=[]
    for k in range(m):
        M.append(random.randint(0,1))
    for k in range(n):
        N.append(random.randint(0,1))

    k = 0
    while 1:
        yield (M[k%m] ^ N[k%n])
        k += 1

def tunnytest(m,n,N=None):
    if not N:
        N = m*n
    T = tunny(m,n)
    S = []
    for k in range(N):
        S.append(T.next())
    return S