# 写入jpeg格式的译码信息
# filename: 输出文件名
# h: 图片高度
# w: 图片宽度
def write_head(filename, h, w):
    # 二进制写入形式打开文件(覆盖)
    fp = open(filename, "wb")
 
    # SOI
    fp.write(pack(">H", 0xffd8))
    # APP0
    fp.write(pack(">H", 0xffe0))
    fp.write(pack(">H", 16))            # APP0字节数
    fp.write(pack(">L", 0x4a464946))    # JFIF
    fp.write(pack(">B", 0))                # 0
    fp.write(pack(">H", 0x0101))        # 版本号: 1.1
    fp.write(pack(">B", 0x01))            # 像素密度单位: 像素/英寸
    fp.write(pack(">L", 0x00480048))    # XY方向像素密度
    fp.write(pack(">H", 0x0000))        # 无缩略图信息
    # DQT_0
    fp.write(pack(">H", 0xffdb))
    fp.write(pack(">H", 64+3))            # 量化表字节数
    fp.write(pack(">B", 0x00))            # 量化表精度: 8bit(0)  量化表ID: 0
    tbl = block2zz(std_luminance_quant_tbl)
    for item in tbl:
        fp.write(pack(">B", item))    # 量化表0内容
    # DQT_1
    fp.write(pack(">H", 0xffdb))
    fp.write(pack(">H", 64+3))            # 量化表字节数
    fp.write(pack(">B", 0x01))            # 量化表精度: 8bit(0)  量化表ID: 1
    tbl = block2zz(std_chrominance_quant_tbl)
    for item in tbl:
        fp.write(pack(">B", item))    # 量化表1内容
    # SOF0
    fp.write(pack(">H", 0xffc0))
    fp.write(pack(">H", 17))            # 帧图像信息字节数
    fp.write(pack(">B", 8))                # 精度: 8bit
    fp.write(pack(">H", h))                # 图像高度
    fp.write(pack(">H", w))                # 图像宽度
    fp.write(pack(">B", 3))                # 颜色分量数: 3(YCrCb)
    fp.write(pack(">B", 1))                # 颜色分量ID: 1
    fp.write(pack(">H", 0x1100))        # 水平垂直采样因子: 1  使用的量化表ID: 0
    fp.write(pack(">B", 2))                # 颜色分量ID: 2
    fp.write(pack(">H", 0x1101))        # 水平垂直采样因子: 1  使用的量化表ID: 1
    fp.write(pack(">B", 3))                # 颜色分量ID: 3
    fp.write(pack(">H", 0x1101))        # 水平垂直采样因子: 1  使用的量化表ID: 1
    # DHT_DC0
    fp.write(pack(">H", 0xffc4))
    fp.write(pack(">H", len(std_huffman_DC0)+3))    # 哈夫曼表字节数
    fp.write(pack(">B", 0x00))                        # DC0
    for item in std_huffman_DC0:
        fp.write(pack(">B", item))                    # 哈夫曼表内容
    # DHT_AC0
    fp.write(pack(">H", 0xffc4))
    fp.write(pack(">H", len(std_huffman_AC0)+3))    # 哈夫曼表字节数
    fp.write(pack(">B", 0x10))                        # AC0
    for item in std_huffman_AC0:
        fp.write(pack(">B", item))                    # 哈夫曼表内容
    # DHT_DC1
    fp.write(pack(">H", 0xffc4))
    fp.write(pack(">H", len(std_huffman_DC1)+3))    # 哈夫曼表字节数
    fp.write(pack(">B", 0x01))                        # DC1
    for item in std_huffman_DC1:
        fp.write(pack(">B", item))                    # 哈夫曼表内容
    # DHT_AC1
    fp.write(pack(">H", 0xffc4))
    fp.write(pack(">H", len(std_huffman_AC1)+3))    # 哈夫曼表字节数
    fp.write(pack(">B", 0x11))                        # AC1
    for item in std_huffman_AC1:
        fp.write(pack(">B", item))                    # 哈夫曼表内容
    # SOS
    fp.write(pack(">H", 0xffda))
    fp.write(pack(">H", 12))            # 扫描开始信息字节数
    fp.write(pack(">B", 3))                # 颜色分量数: 3
    fp.write(pack(">H", 0x0100))        # 颜色分量1 DC、AC使用的哈夫曼表ID
    fp.write(pack(">H", 0x0211))        # 颜色分量2 DC、AC使用的哈夫曼表ID
    fp.write(pack(">H", 0x0311))        # 颜色分量3 DC、AC使用的哈夫曼表ID
    fp.write(pack(">B", 0x00))
    fp.write(pack(">B", 0x3f))
    fp.write(pack(">B", 0x00))            # 固定值
    fp.close()

# 转换图片大小，必须能被切分成8*8的小块
if((h % 8 == 0) and (w % 8 == 0)):
    nblock = int(h * w / 64)
else:
    h = h // 8 * 8
    w = w // 8 * 8
    YCrCb = cv2.resize(YCrCb, [h, w], cv2.INTER_CUBIC)
    nblock = int(h * w / 64)

YCrCb = cv2.cvtColor(BGR, cv2.COLOR_BGR2YCrCb)
npdata = np.array(YCrCb, np.int16)

for i in range(0, h, 8):
    for j in range(0, w, 8):
        ...

now_block = npdata[i:i+8, j:j+8, 0] - 128        # 取出一个8*8块并减去128 Y分量
now_block = npdata[i:i+8, j:j+8, 2] - 128        # 取出一个8*8块并减去128 Cb分量
now_block = npdata[i:i+8, j:j+8, 1] - 128        # 取出一个8*8块并减去128 Cr分量
now_block_dct = cv2.dct(np.float32(now_block))    # DCT变换

# 亮度量化表
std_luminance_quant_tbl = np.array(
    [
        [16, 11, 10, 16, 24, 40, 51, 61],
        [12, 12, 14, 19, 26, 58, 60, 55],
        [14, 13, 16, 24, 40, 57, 69, 56],
        [14, 17, 22, 29, 51, 87, 80, 62],
        [18, 22, 37, 56, 68,109,103, 77],
        [24, 35, 55, 64, 81,104,113, 92],
        [49, 64, 78, 87,103,121,120,101],
        [72, 92, 95, 98,112,100,103, 99]
    ],
    np.uint8
)
# 色度量化表
std_chrominance_quant_tbl = np.array(
    [
        [17, 18, 24, 47, 99, 99, 99, 99],
        [18, 21, 26, 66, 99, 99, 99, 99],
        [24, 26, 56, 99, 99, 99, 99, 99],
        [47, 66, 99, 99, 99, 99, 99, 99],
        [99, 99, 99, 99, 99, 99, 99, 99],
        [99, 99, 99, 99, 99, 99, 99, 99],
        [99, 99, 99, 99, 99, 99, 99, 99],
        [99, 99, 99, 99, 99, 99, 99, 99]
    ],
    np.uint8
)

now_block_qut = quantize(now_block_dct, 0)        # Y分量 量化
now_block_qut = quantize(now_block_dct, 2)        # Cb分量 量化
now_block_qut = quantize(now_block_dct, 1)        # Cr分量 量化

# 量化
# block: 当前8*8块的数据
# dim: 维度  0:Y  1:Cr  2:Cb
def quantize(block, dim):
    if(dim == 0):
        # 使用亮度量化表
        qarr = std_luminance_quant_tbl
    else:
        # 使用色度量化表
        qarr = std_chrominance_quant_tbl
    return (block / qarr).round().astype(np.int16)

now_block_zz = block2zz(now_block_qut)            # zigzag扫描

# zigzag扫描
# block: 当前8*8块的数据
def block2zz(block):
    re = np.empty(64, np.int16)
    # 当前在block的位置
    pos = np.array([0, 0])
    # 定义四个扫描方向
    R = np.array([0, 1])
    LD = np.array([1, -1])
    D = np.array([1, 0])
    RU = np.array([-1, 1])
    for i in range(0, 64):
        re[i] = block[pos[0], pos[1]]
        if(((pos[0] == 0) or (pos[0] == 7)) and (pos[1] % 2 == 0)):
            pos = pos + R
        elif(((pos[1] == 0) or (pos[1] == 7)) and (pos[0] % 2 == 1)):
            pos = pos + D
        elif((pos[0] + pos[1]) % 2 == 0):
            pos = pos + RU
        else:
            pos = pos + LD
    return re

last_block_ydc = 0
last_block_cbdc = 0
last_block_crdc = 0

now_block_dc = now_block_zz[0] - last_block_ydc # 直流分量差分形式记录
last_block_ydc = now_block_zz[0]                # 记录本次量

now_block_dc = now_block_zz[0] - last_block_cbdc
last_block_cbdc = now_block_zz[0]

now_block_dc = now_block_zz[0] - last_block_crdc
last_block_crdc = now_block_zz[0]

now_block_ac = RLE(now_block_zz[1:])

# 0的行程编码
# AClist: 要编码的交流数据
def RLE(AClist: np.ndarray) -> np.ndarray:
    re = []
    cnt = 0
    for i in range(0, 63):
        if(AClist[i] == 0 and cnt != 15):
            cnt += 1
        else:
            re.append(cnt)
            re.append(AClist[i])
            cnt = 0
    # 删除末尾的所有[15 0]
    while(re[-1] == 0):
        re.pop()
        re.pop()
        if(len(re) == 0):
            break
    # 在结尾添加两个0作为结束标记
    if(AClist[-1] == 0):
        re.extend([0, 0])
    return np.array(re, np.int16)

             数值               bit长度            实际保存值
              0                   0                   无
            -1,1                  1                  0,1
         -3,-2,2,3                2              00,01,10,11
   -7,-6,-5,-4,4,5,6,7            3    000,001,010,011,100,101,110,111
     -15,..,-8,8,..,15            4       0000,..,0111,1000,..,1111
    -31,..,-16,16,..,31           5     00000,..,01111,10000,..,11111
    -63,..,-32,32,..,63           6                  ...
   -127,..,-64,64,..,127          7                  ...
  -255,..,-128,128,..,255         8                  ...
  -511,..,-256,256,..,511         9                  ...
 -1023,..,-512,512,..,1023       10                  ...
-2047,..,-1024,1024,..,2047      11                  ...

# 特殊的二进制编码格式
# num: 待编码的数字
def tobin(num):
    s = ""
    if(num > 0):
        while(num != 0):
            s += '0' if(num % 2 == 0) else '1'
            num = int(num / 2)
        s = s[::-1]    # 反向
    elif(num < 0):
        num = -num
        while(num != 0):
            s += '1' if(num % 2 == 0) else '0'
            num = int(num / 2)
        s = s[::-1]
    return s

s = write_num(s, -1, now_block_dc, DC0)            # 根据编码方式写入直流数据
for l in range(0, len(now_block_ac), 2):        # 写入交流数据
    s = write_num(s, now_block_ac[l], now_block_ac[l+1], AC0)
    while(len(s) >= 8):                            # 记录数据太长会导致爆内存
        num = int(s[0:8], 2)                        # 运行速度变慢
        fp.write(pack(">B", num))
        if(num == 0xff):                        # 为防止标志冲突
            fp.write(pack(">B", 0))                # 数据中出现0xff需要在后面补两个0x00
        s = s[8:len(s)]

# 根据编码方式写入数据
# s: 未写入文件的二进制数据
# n: 数据前面0的个数(-1代表DC)
# num: 待写入的数据
# tbl: 范式哈夫曼编码字典
def write_num(s, n, num, tbl):
    bit = 0
    tnum = num
    while(tnum != 0):
        bit += 1
        tnum = int(tnum / 2)
    if(n == -1):                        # DC
        tnum = bit
        if(tnum > 11):
            print("Write DC data Error")
            exit()
    else:                                    # AC
        if((n > 15) or (bit > 11) or (((n != 0) and (n != 15)) and (bit == 0))):
            print("Write AC data Error")
            exit()
        tnum = n * 10 + bit + (0 if(n != 15) else 1)
    # 范式哈夫曼编码记录0的个数(AC)以及num的bit长度
    s += tbl[tnum].str_code
    # 特殊形式的数据存储num
    s += tobin(num)
    return s

# 亮度直流量范式哈夫曼编码表
std_huffman_DC0 = np.array(
    [0, 0, 7, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0,
     4, 5, 3, 2, 6, 1, 0, 7, 8, 9, 10, 11],
    np.uint8
)
...
# 换算出哈夫曼字典
DC0 = DHT2tbl(std_huffman_DC0)    # 亮度直流分量
DC1 = DHT2tbl(std_huffman_DC1)    # 色度直流分量
AC0 = DHT2tbl(std_huffman_AC0)    # 亮度交流分量
AC1 = DHT2tbl(std_huffman_AC1)    # 色度交流分量

# 记录哈夫曼字典的类
# symbol: 原始数据
# code: 对应的编码数据
# n_bit: 编码的二进制位数
# str_code: 编码的二进制数据
class Sym_Code():
    def __init__(self, symbol, code, n_bit):
        self.symbol = symbol
        self.code = code
        str_code=''
        mask = 1 << (n_bit - 1)
        for i in range(0, n_bit):
            if(mask & code):
                str_code += '1'
            else:
                str_code += '0'
            mask >>= 1
        self.str_code = str_code
    """定义输出形式"""
    def __str__(self):
        return "0x{:0>2x}    |  {}".format(self.symbol, self.str_code)
    """定义排序依据"""
    def __eq__(self, other):
        return self.symbol == other.symbol
    def __le__(self, other):
        return self.symbol < other.symbol
    def __gt__(self, other):
        return self.symbol > other.symbol
 
 
# 将范式哈夫曼编码表转换为哈夫曼字典
# data: 定义的范式哈夫曼编码表
def DHT2tbl(data):
    numbers = data[0:16]                # 1~16bit长度的编码对应的个数
    symbols = data[16:len(data)]        # 原数据
    if(sum(numbers) != len(symbols)):    # 判断是否为正确的范式哈夫曼编码表
        print("Wrong DHT!")
        exit()
    code = 0
    SC = []                                # 记录字典的列表
    for n_bit in range(1, 17):
        # 按范式哈夫曼编码规则换算出字典
        for symbol in symbols[sum(numbers[0:n_bit-1]):sum(numbers[0:n_bit])]:
            SC.append(Sym_Code(symbol, code, n_bit))
            code += 1
        code <<= 1
    return sorted(SC)