Image compression and Information theory

def get_k(n):
    k = 0
    while 2 ** k < n + k + 1:
        k += 1
    return k

def get_check_groups(n, k):
    check_groups = {}
    for i in range(1, k + 1):
        index = 2 ** (i - 1)
        check_groups[index] = []
        for j in range(index + 1, n + k + 1):
            if j & (1 << (i - 1)):
                check_groups[index].append(j)
    return check_groups

def get_check_code(group_data, parity):
    """
    group_data: 校验组中的数据位
    parity: 奇校验或偶校验
    :return: 校验位的值
    """
    if parity == 'odd':
        return 1 - sum(group_data) % 2
    elif parity == 'even':
        return sum(group_data) % 2
    else:
        raise ValueError('parity must be odd or even')

class HammingCode:
    def __init__(self, data: (list, tuple, str, int)=None):
        if data is not None:
            self._process(data)

    def _process(self, data: (list, tuple, str, int)):
        if isinstance(data, str):
            data = [int(i) for i in data]
        elif isinstance(data, int):
            # 转为二进制列表
            data = [int(i) for i in bin(data)[2:]]
        self.data = data
        self.n = len(data)
        self.k = get_k(self.n)
        self.check_groups = get_check_groups(self.n, self.k)
    
    def _inverse_print(self, code):
        print("Hamming Code: {}".format("".join([str(i) for i in code[::-1]])))

    def _check_error(self, code=None):
        if code is None:
            code = self.code
        errors = []
        for i, group in self.check_groups.items():
            res = get_check_code([code[j - 1] for j in group + [i]], 'odd')
            if res: # 有错
                errors.append(1) # 补1
            else:
                errors.append(0) # 补0
        if sum(errors) == 0:
            return None
        else:
            return sum([2 ** i for i, e in enumerate(errors) if e])

        if len(errors) == 0:
            return None
        else:
            return errors

    def __call__(self, data=None):
        if data is not None:
            self._process(data)
        return self.encode()
    
    def encode(self):
        # 默认偶校验
        code = [0] * (self.n + self.k)
        data_index = 0
        for i in range(1, self.n + self.k + 1):
            if i in self.check_groups:
                code[i - 1] = None
            else:
                code[i - 1] = self.data[self.n - data_index - 1]
                data_index += 1
        for i, group in self.check_groups.items():
            code[i - 1] = get_check_code([code[j - 1] for j in group], 'even')
        # 倒序输出字符串
        self._inverse_print(code)
        self.code = code
        return code
    
    def decode(self, code=None):
        if code is None:
            code = self.code
        # 默认奇校验
        errors = self._check_error(code)
        if errors is not None:
            code[errors - 1] = 1 - code[errors - 1]
        data = []
        for i in range(1, self.n + self.k + 1):
            if i not in self.check_groups:
                data.append(code[i - 1])
        self.data = data[::-1]
        return "".join([str(i) for i in data[::-1]])

# 信源
class Node:
    def __init__(self, symbol, prob):
        self.symbol = symbol
        self.prob = prob
        self.left = None
        self.right = None
        self.code = None

# 信源按照概率从小到大排序
def sort_nodes(nodes):
    return sorted(nodes, key=lambda x: x.prob)

# 合并两个信源
def merge_nodes(node1, node2):
    node = Node(None, node1.prob + node2.prob)
    node.left = node1
    node.right = node2
    return node

# 生成霍夫曼编码
def huffman(nodes):
    # 信源按照概率从小到大排序
    nodes = sort_nodes(nodes)
    # 重复合并两个信源，直至合并到根节点
    while len(nodes) > 1:
        node1 = nodes.pop(0)
        node2 = nodes.pop(0)
        node = merge_nodes(node1, node2)
        for i in range(len(nodes)):
            if nodes[i].prob > node.prob:
                nodes.insert(i, node)
                break
        else:
            nodes.append(node)
    # 从根节点开始，左子树编码为0，右子树编码为1
    codes = {}
    root = nodes[0]
    root.code = ''
    stack = [root]
    while stack:
        node = stack.pop()
        if node.symbol:
            codes[node.symbol] = node.code
        if node.left:
            node.left.code = node.code + '0'
            stack.append(node.left)
        if node.right:
            node.right.code = node.code + '1'
            stack.append(node.right)
    # 输出编码
    return codes

probs = [0.1, 0.2, 0.3, 0.4]
symbols = ['A', 'B', 'C', 'D']
nodes = [Node(symbol, prob) for symbol, prob in zip(symbols, probs)]
codes = huffman(nodes)
print(codes)

from collections import Counter

string = "ABAACCAACB"

def float2bin(f):
    if f >= 1 or f <= 0:
        return "Error"
    res = "0."
    while f > 0:
        f *= 2
        if f >= 1:
            res += "1"
            f -= 1
        else:
            res += "0"
    return res

def arithmetic_encode(string):

    probs = Counter(string)
    probs = {k: v / len(string) for k, v in probs.items()}
    left = 0
    intervals = {}
    for k in probs:
        intervals[k] = (left, left + probs[k])
        left += probs[k]

    left, right = 0, 1
    for c in string:
        l, r = intervals[c] # 信源的编码区间
        interval = right - left # 区间长度
        right = left + interval * r # 更新右端点
        left = left + interval * l # 更新左端点
        # print(c, (left, right))
    midpoint = (left + right) / 2 # 区间的中点
    return float2bin(midpoint)[2:], probs, intervals # 转为二进制，去掉前面的0.

def arithmetic_decode(code, probs, intervals):
    goal = float("0." + code)
    string = ""
    l = 0
    r = 1
    while True:
        interval = r - l
        for k, v in intervals.items():
            l_new = l + interval * v[0]
            r_new = l + interval * v[1]
            if l_new <= goal < r_new:
                string += k
                l = l_new
                r = r_new
                break
        if len(string) == len(code):
            break
    return string