#!/usr/bin/python

# Linear Feedback Shift Register (LFSR) Random Number generator We should be
# able to go through all of the possible states from our initial one before we
# start repeating. This means that we can have 2^n-1 unique numbers, n being
# the number of bits in our state before we start repeating numbers.

# Warning: This algorithm is NOT cryptographically secure, because given enough
# outputted bytes, by solving a bunch of linear equations and recompute the
# LFSR generator bits.  If we do need something cryptographically secure, we
# could combine 2 LFSR streams XORed together, and they don't always shift at
# the same time, we could achieve something a lot more chaotic and very hard to
# predict.


def pcb(number: int):
    """
    A debug function for printing a number represented in a binary form, in a
    compressed manner, this is only really useful if the number is expected to
    have a lot of repeated bits (100000000001 would be 1x1+10x0+1x1)
    """
    compressed_binary = []
    binary_form = bin(number)[2:]
    last_bit = binary_form[0]  # To avoid 0x0 or 0x1 always use starting last bit
    bit_amount = 0

    for bit in binary_form:
        if bit == last_bit:
            bit_amount += 1
        else:
            compressed_binary.append(f"{bit_amount}x{last_bit}")
            last_bit = bit
            bit_amount = 1
    compressed_binary.append(f"{bit_amount}x{last_bit}")

    return "+".join(compressed_binary)


def lfsr(state: int, number_size: int, taps: list[int]) -> int:
    """
    LFSR random number generator algorithm.
    - `state` is the initial "seed" of our algorithm
    - `taps` are the bit positions which will be used in the XOR
    - `number_size` is the size for the number to be generated in bits

    The returned number is composed of `number_size` of bits, produced by the
    algorithm given the `state` and `taps`.
    """
    output = 0
    state_binary_length = len(bin(state)) - 2
    max_nonrepeating = 2 ** state_binary_length - 1
    if number_size > max_nonrepeating:
        print(
            f"Warning, with {state_binary_length} binary state length, you "
            f"only have {max_nonrepeating} non-repeating random numbers, you "
            f"requested {number_size} numbers, which means some of the bits in it "
            "will follow a repeated sequence."
        )
    for i in range(number_size):
        # Get the first bit, which will be our random output
        rand = state & 1
        # Print the state and bit output for debugging purposes
        print(f"#{i + 1} {state=:04b}, out={rand:01b}")

        # Leftshift the output to make space for a new ending bit and use OR to
        # set the bit depending on rand
        output = (output << 1) | rand
        # Take the XOR of the tapped bits by running XOR on the state and the
        # right-shifted state, and taking the last digit of that.
        # For example:
        # 1001 is our original state
        # 0100 is the right bitshifted state
        # 1101 is the XOR result for taps=[1] (works for state=0b1001)
        newbit = state
        for tap_bit_pos in taps:
            newbit = newbit ^ (state >> tap_bit_pos)
        # We can now use result & 1 to only obtain the last bit allowing us to
        # only run XOR between last 2 bits of current state
        newbit = newbit & 1

        # Bitshift the state to the right and set the first bit of it to the
        # computed newbit value
        state = (state >> 1) | (newbit << (state_binary_length - 1))

    return output



# Define some parameters to be used for the LFSR function

# The state used in the start, can't be only zeroes
#start_state = 0b1001
# For state of 1001, the taps (positions to get XORed) are just [1]
# taps = [1]

# Will produce 1000...0001 of 128 bits long guaranteeing 2^128 - 1
# non-repeating pseudo-random number sequence
start_state = (1 << 127) | 1
taps = [1, 2, 7]

# The desired size (in bits) of the generated random number
number_size = 24


if __name__ == "__main__":
    out = lfsr(start_state, number_size, taps=taps)
    binary_out = f"{out:b}".zfill(number_size)
    print(f"Obtained random number: {out} (binary form: {binary_out}")