import bma
from sympy.polys.domains import ZZ
from sympy.polys.galoistools import gf_ddf_shoup, gf_edf_shoup

# This file provides an example of converting a raw bit sequence
# to a linear feedback shift register (LFSR).

def poly_to_str(x):
    """Convert list of 1/0 polynomial coefficients(x) to nicer format"""
    deg = len(x)-1
    coeffs = []
    for n,elem in enumerate(x):
        if elem != 0:
            coeffs.append(deg-n)
    out = ' + '.join(['x^%d'%c for c in coeffs]).replace('x^0','1')
    return out

def parity16bit(value):
    """1 if parity is odd, 0 if even"""
    value = value ^ (value >> 8)
    value = value ^ (value >> 4)
    value = value ^ (value >> 2)
    value = value ^ (value >> 1)
    return value & 1

def generate_b2b(sv_id,code):
    """Generate BeiDou B2b-Q or I code for given SV#.
    sv_id = satellite PRN 1-63
    code = 'Q' or 'I'. Q is only tested against raw memory codes for 19-48.
           'I' code should work for all ICD values.
    """
    B2B_XAStartState = 0x1fff
    B2B_XBiStartStateLUT = [
        0x1481, 0x0581, 0x16a1, 0x1e51, 0x1551, 0x0eb1, 0x0ef1, 0x1bf1,
        0x1299, 0x0b79, 0x1585, 0x0445, 0x1545, 0x1b45, 0x0745, 0x18a5,
        0x1de5, 0x1015, 0x0f95, 0x1ab5, 0x11b5, 0x194d, 0x08cd, 0x032d,
        0x0dad, 0x09ed, 0x1fed, 0x091d, 0x079d, 0x10bd, 0x027d, 0x057d,
        0x1afd, 0x19fd, 0x1143, 0x0523, 0x1da3, 0x1113, 0x1313, 0x1ab3,
        0x11b3, 0x0973, 0x154b, 0x05cb, 0x1a6b, 0x1d5b, 0x0587, 0x1827,
        0x1a27, 0x18a7, 0x02a7, 0x1b97, 0x1d37, 0x024f, 0x052f, 0x132f,
        0x0b6f, 0x03ef, 0x1fef, 0x15bf, 0x17bf, 0x143a, 0x1b9a]
    state_D1 = B2B_XAStartState
    state_D2 = B2B_XBiStartStateLUT[sv_id - 1]
    chip_cnt = 0
    out_bits = []
    if code == 'I':
        # From ICD:
        #  x^13 + x^10 + x^9 + x^1 + 1
        #  x^13 + x^12 + x^9 + x^6 + x^4 + x^3 + 1
        fb_D1 = 0x1301
        fb_D2 = 0x192C
    elif code == 'Q':
        # Tried out all possible values against memory codes until
        # the code matched.
        fb_D1 = 0x1c01
        fb_D2 = 0x1782

    while chip_cnt < 10230:
        wordI = 0
        for i in range(8):
            wordI <<= 1
            curr_bit = ((state_D1 ^ state_D2) >> 12) & 1
            wordI  |= curr_bit
            if chip_cnt < 10230:
                out_bits.append( curr_bit )

            bit1 = parity16bit( state_D1 & fb_D1 )
            bit2 = parity16bit( state_D2 & fb_D2 )
            state_D1 <<= 1
            state_D1 |= bit1;
            state_D2 <<= 1
            state_D2 |= bit2
            chip_cnt += 1
            if chip_cnt == 8190:
                state_D1 = B2B_XAStartState
    return out_bits

def print_lfsr_taps(seq):
    """Given a code sequence ([1,0,1,...]), try to generate the
    LFSR representation and a factored representation."""

    if len(seq)>2000:
        # Berlekamp-Massey seems to have problems if this gets too long...
        seq = seq[:2000]
        print("Only using first 2000 bits for Berlekamp-Massey")

    # Get full LFSR
    p_str, p_coeff, p_deg = bma.Berlekamp_Massey_algorithm(seq)
    if p_deg >= 100:
        print("The poly degree %d seems too high.  Aborting."%p_deg)
        return
    print("Full LFSR: %s"%p_str)

    if not p_str.endswith('+ 1'):
        print("Something is probably wrong if we don't have a +1 term in the")
        print("polynomial.  Aborting.")
        return

    # Convert to sympy format
    poly = [0]*(p_deg+1)
    for i in p_coeff:
        poly[i]=1

    # Try to factor into non-equal polys first.  Also gives degree.
    ddf = gf_ddf_shoup(poly, 2, ZZ)
    if len(ddf) > 1:
        print("gf_ddf_shoup() returned non-equal factors:")
        print('  ',ddf)
        print("This often happens if the input sequence was too short.  Aborting.")
        return

    # Now we have the degree, factor into equal polys.
    edf = gf_edf_shoup(ddf[0][0], ddf[0][1], 2, ZZ)

    if len(edf) == 2:
        print("Factored full LFSR into two equal-length LFSRs:")
        print(poly_to_str(edf[0]))
        print(poly_to_str(edf[1]))
    else:
        print("Couldn't factor into equal polynomials")


if __name__ == '__main__':
    # This is just an example input code where we know the answer.
    # You could have use any code here.
    for code in  ['I','Q']:
        print("********************")
        b2b_code = generate_b2b(48,code)
        print("Few bits of B2B-{} code PRN 48: {}".format(
            code,
            b2b_code[:40]))

        # Try to generate LFSR taps
        print_lfsr_taps(b2b_code)
        print("********************")