import numpy as np
from struct import unpack
from PIL import Image, ImageTk
import sys
import tkinter
import tkinter.filedialog

class BmpData:
    def __init__(self, file_path:str):
        with open(file_path, "rb") as file:
            self.file = file

            self.bfType = unpack("<H", file.read(2))[0] # 0x00 文件类型
            self.bfSize = unpack("<i", file.read(4))[0] # 0x02 文件大小
            self.bfReserved1 = unpack("<H", file.read(2))[0]    # 0x06 保留，必须设置为0
            self.bfReserved2 = unpack("<H", file.read(2))[0]    # 0x08 保留，必须设置为0
            self.bfOffBits = unpack("<i", file.read(4))[0]  # 0x0a 从头到位图数据的偏移
            self.biSize = unpack("<i", file.read(4))[0] # 0x0e 信息头的大小
            self.biWidth = unpack("<i", file.read(4))[0]    # 0x12 图像的宽度(以像素为单位)
            self.biHeight = unpack("<i", file.read(4))[0]   # 0x16 图像的高度(以像素为单位)(负说明图像是倒立的)
            self.biPlanes = unpack("<H", file.read(2))[0]   # 0x1a 颜色平面数
            self.biBitCount = unpack("<H", file.read(2))[0] # 0x1c 比特数/像素数
            self.biCompression = unpack("<i", file.read(4))[0]  # 0x1e 压缩类型
            self.biSizeImage = unpack("<i", file.read(4))[0]    # 0x22 位图数据的大小
            self.biXPelsPerMeter = unpack("<i", file.read(4))[0]    # 0x26 水平分辨率
            self.biYPelsPerMeter = unpack("<i", file.read(4))[0]    # 0x2a 垂直分辨率
            self.biClrUsed = unpack("<i", file.read(4))[0]  # 0x2e 位图使用的调色板中的颜色索引数
            self.biClrImportant = unpack("<i", file.read(4))[0] # 0x32 对图像显示有重要影响的颜色索引数(0说明都重要)

            self.color_palette = self.get_color_palette()
            self.img_np = self.get_numpy_img()
            self.gray = self.get_gray_img()
            file.close()

    def get_color_palette(self) -> np.ndarray:
        if (self.bfOffBits == 0x36):  # 16/24位图像不需要调色板，起始位置就等于0x36
            return None
        color_alette_size = 2 ** int(self.biBitCount) # 多少字节调色板颜色就有2^n个
        color_palette = np.zeros((color_alette_size, 3), dtype=np.int32)
        self.file.seek(0x36)
        for i in range(color_alette_size):
            b = unpack("B", self.file.read(1))[0]
            g = unpack("B", self.file.read(1))[0]
            r = unpack("B", self.file.read(1))[0]
            alpha = unpack("B", self.file.read(1))[0]
            color_palette[i][0] = b
            color_palette[i][1] = g
            color_palette[i][2] = r
        return color_palette

    def get_numpy_img(self) -> np.ndarray:
        biHeight = abs(self.biHeight)
        img_np = np.zeros((biHeight, self.biWidth, 3), dtype=np.int32)
        self.file.seek(self.bfOffBits)
        for x in range(biHeight):
            row_byte_count = ((self.biWidth * self.biBitCount + 31) >> 5) << 2
            row_bits = self.file.read(row_byte_count)
            row_bits = ''.join(format(byte, '08b') for byte in row_bits)
            for y in range(self.biWidth):
                pixel_data = row_bits[y * self.biBitCount: (y + 1) * self.biBitCount]
                if self.biHeight > 0:   # 图像倒立
                    img_np[biHeight - 1 - x][y] = self.get_RGB(pixel_data)
                else:
                    img_np[x][y] = self.get_RGB(pixel_data)
        return img_np
    
    def get_gray_img(self) -> np.ndarray:
        biHeight = abs(self.biHeight)
        gray_img = np.dot(self.img_np.reshape((biHeight * self.biWidth, 3)).astype(np.float32), 
                          [0.299, 0.587, 0.114]).astype(np.int32)
        gray_img = gray_img.reshape((biHeight, self.biWidth))
        return gray_img

    def get_RGB(self, pixel_data:str):
        if len(pixel_data) <= 8:
            color_index = int(pixel_data, 2)
            return self.color_palette[color_index]
        elif len(pixel_data) == 16:
            b = int(pixel_data[1:6], 2) * 8
            g = int(pixel_data[6:11], 2) * 8
            r = int(pixel_data[11:16], 2) * 8
            return [r, g, b]
        elif len(pixel_data) == 24:
            b = int(pixel_data[0:8], 2)
            g = int(pixel_data[8:16], 2)
            r = int(pixel_data[16:24], 2)
            return [r, g, b]
        elif len(pixel_data) == 32:
            b = int(pixel_data[0:8], 2)
            g = int(pixel_data[8:16], 2)
            r = int(pixel_data[16:24], 2)
            alpha = int(pixel_data[24:32], 2)
            return [r, g, b]
        

    def equalize(self, level:int):
        biHeight = abs(self.biHeight)
        self.hist = np.zeros(256, dtype=np.int32)
        max_value = self.gray.max()
        min_value = self.gray.min()
        gap = (max_value - min_value + 1) / level
        for x in range(biHeight):
            for y in range(self.biWidth):
                self.hist[self.gray[x, y]] += 1
        hist = np.zeros(level, dtype=np.float32)
        for i in range(level):
            hist[i] = np.sum(self.hist[min_value + int(i * gap) : min_value + int((i + 1) * gap)])
        hist /= biHeight * self.biWidth
        for i in range(1, level):
            hist[i] += hist[i - 1]
        hist *= level
        hist = np.around(hist)
        hist /= level
        hist = np.floor(hist * 255).astype(np.int32)
        self.equalized_img = np.zeros_like(self.gray)
        self.equalized_hist = np.zeros(256, dtype=np.int32)
        for x in range(biHeight):
            for y in range(self.biWidth):
                self.equalized_img[x, y] = hist[int((self.gray[x, y] - min_value) / gap)]
                self.equalized_hist[self.equalized_img[x, y]] += 1
        return self.equalized_img, self.hist, self.equalized_hist
    
    def save_equalized_img(self, save_path:str):
        self.save_img(image=self.equalized_img, save_path=save_path)
    
    def save_img(self, image:np.ndarray, save_path:str):
        with open(save_path, "wb") as file:
            file.write(int(self.bfType).to_bytes(2, byteorder='little'))    # 0x00 文件类型
            file.write(int(0x36 + 0x100 * 4 + self.biWidth * abs(self.biHeight)).to_bytes(4, byteorder='little'))   # 0x02 文件大小
            file.write(int(0).to_bytes(4, byteorder='little'))  # 0x06 保留，必须设置为0
            file.write(int(0x36 + 0x100 * 4).to_bytes(4, byteorder='little')) # 0x0a 从头到位图数据的偏移
            file.write(int(40).to_bytes(4, byteorder='little')) # 0x0e 信息头的大小
            file.write(int(self.biWidth).to_bytes(4, byteorder='little'))   # 0x12 图像的宽度
            file.write(int(self.biHeight).to_bytes(4, byteorder='little'))  # 0x16 图像的高度
            file.write(int(self.biPlanes).to_bytes(2, byteorder='little'))  # 0x1a 颜色平面数
            file.write(int(8).to_bytes(2, byteorder='little'))    # 0x1c 比特数/像素数
            file.write(int(self.biCompression).to_bytes(4, byteorder='little')) # 0x1e 压缩类型
            file.write(int(self.biSizeImage).to_bytes(4, byteorder='little'))   # 0x22 位图数据的大小
            file.write(int(self.biXPelsPerMeter).to_bytes(4, byteorder='little'))   # 0x26 水平分辨率
            file.write(int(self.biYPelsPerMeter).to_bytes(4, byteorder='little'))   # 0x2a 垂直分辨率
            file.write(int(0x100 * 4).to_bytes(4, byteorder='little'))  # 0x2e 位图使用的调色板中的颜色索引数
            file.write(int(0).to_bytes(4, byteorder='little'))  # 0x32 对图像显示有重要影响的颜色索引数

            for i in range(256):
                file.write(int(i).to_bytes(1, byteorder='little'))
                file.write(int(i).to_bytes(1, byteorder='little'))
                file.write(int(i).to_bytes(1, byteorder='little'))
                file.write(int(0).to_bytes(1, byteorder='little'))

            for x in range(abs(self.biHeight)):
                for y in range(self.biWidth):
                    if self.biHeight > 0:
                        file.write(int(image[self.biHeight - 1 - x][y]).to_bytes(1, byteorder='little'))
                    else:
                        file.write(int(image[x][y]).to_bytes(1, byteorder='little'))
                file.write(b'0' * ((((self.biWidth * 8 + 31) >> 5) << 2) - 8 * self.biWidth))

            file.close()


def choosepic():
    global path_
    path_ = tkinter.filedialog.askopenfilename(title='请选择图片文件', filetypes=[('图片', '.bmp')])
    if path_ == '':
        return
    img_temp = Image.open(path_).resize((int(256 * 0.8), int(256 * 0.8)))  # 图片读取和加载
    img = ImageTk.PhotoImage(img_temp)
    label_image1.config(image=img)
    label_image1.image = img


def equalize():
    if path_ == '':
        return
    image = BmpData(path_)
    # img = Image.fromarray(image.img_np.astype(np.uint8))
    # img.show()

    equalized_img, hist, equalized_hist = image.equalize(8) # 分别为均衡化的图/直方图/均衡化后的直方图
    equalized_img = Image.fromarray(equalized_img.astype(np.uint8))
    # equalized_img.show()

    name_parts = path_.split('.')
    name_parts[-2] += "_equalized"
    new_file_name = '.'.join(name_parts)
    image.save_equalized_img(new_file_name)
    
    equalized_img = equalized_img.resize((int(256 * 0.8), int(256 * 0.8)))
    equalized_img = ImageTk.PhotoImage(equalized_img)
    label_image2.config(image=equalized_img)
    label_image2.image = equalized_img  # 处理后的图片的显示

                 
if __name__ == "__main__":
    root = tkinter.Tk()
    root.title('21281280柯劲帆')  # 标题
    width, height = 600, 400
    width_max, height_max = root.maxsize()
    s_center = '%dx%d+%d+%d' % (width, height, (width_max - width) / 2, (height_max - height) / 2)  # 将页面显示在正中间
    root.geometry(s_center)
    root.resizable(width=False, height=False)  # 窗口不可移动
    l = tkinter.Label(root, text='实验二', width=60, height=2, fg='black', font=("微软雅黑", 16), anchor=tkinter.CENTER)
    l.pack()

    label_image1 = tkinter.Label(root, width=int(256 * 0.8), height=int(256 * 0.8), bg='whitesmoke', anchor=tkinter.NE)
    label_image1.pack()
    label_image1.place(x=45, y=70, width=int(256 * 0.8), height=int(256 * 0.8))

    label_image2 = tkinter.Label(root, width=int(256 * 0.8), height=int(256 * 0.8), bg='whitesmoke', anchor=tkinter.NE)
    label_image2.place(x=350, y=70, width=int(256 * 0.8), height=int(256 * 0.8))

    # 文本按钮
    Image_Input = tkinter.Button(root, text='Choose', command=choosepic)
    Image_Input.place(x=105, y=300, width=80, height=30)

    # 直方图均衡化
    Fun1 = tkinter.Button(root, text='直方图均衡化', command=equalize)
    Fun1.place(x=265, y=300, width=80, height=30)

    # 退出
    Quit = tkinter.Button(root, text='Quit', command=sys.exit)
    Quit.place(x=415, y=300, width=80, height=30)

    end = tkinter.Label(root, text='21281280 柯劲帆', fg='silver', font=("微软雅黑", 10))
    end.place(x=215, y=360, width=200, height=20)
    root.mainloop()