# Encode and Decode Hamming Code for a binary string
#
# By: Chris Gregg (chg5w@virginia.edu), January 2009.
#
# Examples: hamEncode([1,1,0,1,1])
#             returns [0, 0, 1, 0, 1, 0, 1, 1, 1]
#           hamDecode([0,1,1,1,1,1])
#             returns [1,1,0]

import math

# hamEncode returns a binary list that has been encoded with a hamming code
# ex: hamEncode([1,1,0,1] returns [1,0,1,0,1,0,1])
def hamEncode(b):
    # copy b to a temporary list
    ham = list(b)
    # insert partity bits, but set them to 0 for now
    n = 0
    while len(ham) >= 2**n:
        ham.insert((2**n)-1,0)
        n = n + 1

    # determine the parity for each parity bit
    done = False
    for i in range(n):
        bit = 2**i # the start bit
        # To determine the parity, calculate as follows:
        # xor 'bit' number of bits, then skip 'bit' number, then
        # xor 'bit' number of bits, then skip 'bit' number, etc.
        # until the end of the list
        parity = 0
        offset = bit-1
        done = False
        while not done:
            for j in range(bit):
                if len(ham) > j+offset:
                    parity = parity ^ ham[j+offset]
                else:
                    done = True
                    continue
            offset = offset + 2*bit
        ham[bit-1]=parity
    return ham

# hamDecode(b) returns a binary list that corrects for an incorrect bit in a stream
# that has been encoded with Hamming Code
# Note: the original list is NOT changed, but the corrected list is returned,
#       without the check bits
def hamDecode(b):
    # copy b to a temporary list
    ham = list(b)
    totalCheck = 0
    for n in range(int(math.ceil(math.log(len(ham),2)))):
        check = 0
        done = False
        bit = 2**n
        offset = bit-1
        while not done:
            # check 'bit' bits and then skip 'bit' bits
            for i in range(bit):
                if i+offset < len(ham):
                    check = check ^ ham[i+offset]
                else:
                    done=True
            offset = offset + bit*2
        if check==1: # we have an error
            totalCheck = totalCheck + bit
    # The "totalCheck" bit is incorrect
    if totalCheck != 0:
        ham[totalCheck-1]=int(not ham[totalCheck-1])
    # We now have a corrected list, so we can strip out the check bits.
    # Strip the check bits from the right, to make it easier
    for n in range(int(math.ceil(math.log(len(ham),2))),0,-1):
        bit = 2**(n-1)-1
        ham.pop(bit)
    return ham