__author__ = "Wei-Chung Huang"
__copyright__ = "Copyright 2018, The SJSU MSCS Master project"
__license__ = "MIT"
__version__ = "1.0.0"


class HammingCode74:
    """
    Hamming (7, 4) encoder and decoder.

    Author: Wei-Chung (Wayne) Huang
    Date: Feb 2018

    """
    def __init__(self):
        pass

    def extract_bit(self, data, pos):
        """
        An utility to extract the bit in a specific position.
        :param data: An integer
        :param pos: Left shift position
        :return: An integer of selected bit.
        """
        return (data >> pos) & 1

    def set_bit(self, data, pos, bit):
        if bit is 1:
            return data | (1 << pos)
        else:
            return data & ~(1 << pos)

    def encode_byte(self, data):
        """
        Encode a byte data (8 bits) and return 2 bytes.
        :param data: Byte data or integer [0, 255]
        :return: 2 encoded byte due to hamming (7, 4).
        """
        lo = data & 0xf
        hi = (data >> 4) & 0xf
        half = [lo, hi]
        encoded = list()
        for i in range(2):
            d = [0, 0, 0, 0]
            d[0] = self.extract_bit(half[i], 3)
            d[1] = self.extract_bit(half[i], 2)
            d[2] = self.extract_bit(half[i], 1)
            d[3] = self.extract_bit(half[i], 0)

            code = self.encode(d)
            encoded.append(code)

        return encoded

    def encode(self, d):
        """
        Encode 4 bits data.
        Encoded word: P1 P2 D1 P3 D2 D3 D4
        Example:
            encode([1, 0, 1, 1]) -> binary(0 1 1 0 0 1 1)
        :param d: List of binary data with size = 4
        :return: An encoded integer.
        """
        p = [0, 0, 0]
        p[0] = (d[0] + d[1] + d[3]) & 1
        p[1] = (d[0] + d[2] + d[3]) & 1
        p[2] = (d[1] + d[2] + d[3]) & 1

        code = p[0] << 6 | p[1] << 5 | d[0] << 4 | p[2] << 3 | d[1] << 2 | d[2] << 1 | d[3] << 0
        return code

    def decode(self, code):
        """
        Decode an one byte codeword to a half byte.
        Expected encoded word: 0 P1 P2 D1 P3 D2 D3 D4
        :param code: An integer codeword [0, 255]
        :return: An integer with only 4 lower bits affected
        """
        d1 = self.extract_bit(code, 4)
        d2 = self.extract_bit(code, 2)
        d3 = self.extract_bit(code, 1)
        d4 = self.extract_bit(code, 0)

        p1 = self.extract_bit(code, 6)
        p2 = self.extract_bit(code, 5)
        p3 = self.extract_bit(code, 3)

        y1 = (d1 + d2 + d4 + p1) & 1
        y2 = (d1 + d3 + d4 + p2) & 1
        y3 = (d2 + d3 + d4 + p3) & 1

        error_pos = (y3 * 4 + y2 * 2 + y1)
        if error_pos > 0:
            if error_pos == 3:
                d1 ^= 1
            elif error_pos == 5:
                d2 ^= 1
            elif error_pos == 6:
                d3 ^= 1
            elif error_pos == 7:
                d4 ^= 1

        return d1 << 3 | d2 << 2 | d3 << 1 | d4

    def get_syndrome(self, code):
        """
        Get the syndrome from an input hamming code
        :param code: An integer code with 7 bits affected
        :return: Triplet of the syndrome
        """
        d1 = self.extract_bit(code, 4)
        d2 = self.extract_bit(code, 2)
        d3 = self.extract_bit(code, 1)
        d4 = self.extract_bit(code, 0)

        p1 = self.extract_bit(code, 6)
        p2 = self.extract_bit(code, 5)
        p3 = self.extract_bit(code, 3)

        y1 = (d1 + d2 + d4 + p1) & 1
        y2 = (d1 + d3 + d4 + p2) & 1
        y3 = (d2 + d3 + d4 + p3) & 1
        return y1, y2, y3

    @staticmethod
    def hamming_distance(s1, s2):
        """
        Get a Hamming distance with 2 strings character by character.
        :param s1: An input string 1
        :param s2: An input string 2
        :return: Hamming distance
        """
        assert len(s1) == len(s2), 'Input string lengths are not equal!!'
        distance = 0
        for i in range(len(s1)):
            if s1[i] != s2[i]:
                distance += 1
        return distance


def test():
    # Example
    coder = HammingCode74()

    nums = [12, 67, 24, 255, 30]
    encoded = list()
    for num in nums:
        encoded += coder.encode_byte(num)

    print(encoded)

    print('Testing all 4 data bits...')
    for i in range(16):
        u_test = [coder.extract_bit(i, 3), coder.extract_bit(i, 2), coder.extract_bit(i, 1), coder.extract_bit(i, 0)]
        res = coder.encode(u_test)
        for j in range(8):
            corrupted = coder.set_bit(res, j, 0)
            corrected = coder.decode(corrupted)
            assert corrected == i, 'Incorrect decoding!'

            corrupted = coder.set_bit(res, j, 1)
            corrected = coder.decode(corrupted)
            assert corrected == i, 'Incorrect decoding!'
    print('Test has passed')

    res = coder.encode([1, 1, 1, 1])
    print(bin(res))

    res = coder.decode(res)
    print(bin(res))

    res = coder.encode_byte(51)
    print(bin(res[0]), bin(res[1]))


if __name__ == '__main__':
    test()
