import numpy as np

import copy

  

# 节点类

class Node:

    def __init__(self,name=None,value=None):

        self._name=name

        self._value=value

        self.l_child = []    # 子节点列表

        self.code=b""

    # 添加子节点

    def add_child(self,node):

        self.l_child.append(node)

  

# You can modify the code as what you want

class huffman_coding_in_byte:

    def __init__(self, text):

        self.text = text

        self.bits = 256

        # you can store your encode and decode dict here

        self.encode_dict = dict()

        self.decode_dict = dict()

        self.words=self.get_prob()[0]

        self.probs=self.get_prob()[1]

        self.L=len(self.words)

        self.k=0

        while self.k*(self.bits-1)+1 < self.L:

            self.k+=1

        self.first=self.bits-(self.k*(self.bits-1)+1 - self.L) #第一次合并的项数

        self.nodes=[Node(self.words[i],self.probs[i]) for i in range(self.L)]

    # you can use this to find all the unique source with sorted prob.

    def get_prob(self):

        unique = np.array(list(set(self.text)))

        prob = np.array([self.text.count(u)/len(self.text) for u in unique])

        sort_idx = np.argsort(prob)[::-1]

        return list(unique[sort_idx]), list(prob[sort_idx])

    def select_sort_nodes(self):

        if len(self.nodes) == 1:

            return

        else:

            # 选择概率最小的几个节点合并成新节点

            if self.first != 0:

                new_node=Node(

                    name=self.nodes[-self.first]._name+self.nodes[-1]._name,

                    value=sum([self.nodes[i]._value for i in range(-self.first,0)])

                )                

                for i in range(-self.first,0):

                    new_node.add_child(self.nodes[i])

                for i in range(0,self.first):

                    self.nodes.pop()        

                self.first = 0  

            else:

                new_node=Node(

                    name=self.nodes[-self.bits]._name+self.nodes[-1]._name,

                    value=sum([self.nodes[i]._value for i in range(-self.bits,0)])

                )

                for i in range(-self.bits,0):

                    new_node.add_child(self.nodes[i])

                for i in range(0,self.bits):

                    self.nodes.pop()

  

            self.nodes.append(new_node)

            # 节点排序

            values = []

            for i in range(len(self.nodes)):

                values.append(self.nodes[i]._value)

            nodes = []

            for i in range(len(self.nodes)):

                nodes.append(self.nodes[i]._name)

            idx = np.argsort(values)[::-1]

            n=[]

            for name in np.array(nodes)[idx]:

                for i in range(len(self.nodes)):

                    if self.nodes[i]._name ==  name:

                        n.append(self.nodes[i])

            self.nodes=copy.deepcopy(n) #这个没必要好像==================================

            self.select_sort_nodes()

            return

  

    def generate_encode_dict(self,current_node):

        for i in range(0,len(current_node.l_child)):

            if current_node.l_child[i].l_child == []: # 如果没有子节点就给一个编码            

                current_node.l_child[i].code = current_node.code+i.to_bytes(1,"big")

                self.encode_dict[current_node.l_child[i]._name]=current_node.l_child[i].code

                # with open("编码对照文件.txt","a+",encoding="utf-8") as f:

                #     f.write(current_node.l_child[i]._name+":")

                #     f.write(str(current_node.l_child[i]._value))

                #     f.write(str(current_node.l_child[i].code)+"\n")

            else:

                current_node.l_child[i].code = i.to_bytes(1,"big")

                self.generate_encode_dict(current_node.l_child[i])

        return

  

    def encode(self):

        #encode the text with huffman coding

        self.select_sort_nodes()

        self.generate_encode_dict(self.nodes[0]) #传入根节点

        print("编码字典做好了，编码的字符个数为：",len(self.encode_dict))

        encoded_text=b""

        # 编码

        for w in self.text:

            encoded_text += self.encode_dict[w]

        return encoded_text

    def decode(self, encoded_text):

        # decode the encoded text

        decoded_text = ""

        # 解码的字典

        self.decode_dict={v:k for k,v in self.encode_dict.items()}

        # 解码

        code = b""

        for w in encoded_text:

            code = code + w.to_bytes(1,"big")

            #print(code) #前缀码

            if code in self.decode_dict:

                decoded_text+= self.decode_dict[code]

                code = b""

        return decoded_text

  
  

## test 1

# decoded text should be the same as the original text

with open('孤星.txt', 'r', encoding="utf-8") as f:

    file_content = f.read()

  

my_huffman_coding = huffman_coding_in_byte(file_content)

encoded_text = my_huffman_coding.encode()

#print(encoded_text)

print(my_huffman_coding.encode_dict) # 打印编码字典

print(my_huffman_coding.decode(encoded_text))  # 打印解码结果

  

## Test 2

utf_encoded_text = file_content.encode("utf-8")

print("压缩前：",len(utf_encoded_text))

print("压缩后",len(encoded_text))

压缩前： 429992
压缩后 204812