import sys
import pygame
import os
import shutil
from pygame.locals import *
import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageEnhance, ImageFilter

pygame.init()
pygame.mixer.init()

SIZE = WIDTH, HEIGHT = 500, 800
FPS = 240
TITLE = "观测应力双折射实验图像处理工具"
folder_name = "temp"
ICON = pygame.image.load('skd1.png')
BG_COLOR = 255, 255, 255
max_x = 2000
max_y = 1200
ini_image = None
bottom_image = None 
sub_image = None
image_needto_save = None
ini_image_path = None
bottom_image_path = None
sub_image_path = None

full_path = None 

index = 0

image1 = None 
image2 = None
image3 = None

filetypes_images = [("Image Files", "*.jpg *.jpeg *.png *.bmp")]

screen = pygame.display.set_mode(SIZE)
pygame.display.set_caption(TITLE)
pygame.display.set_icon(ICON)
clock = pygame.time.Clock()
font1 = pygame.font.SysFont("SimSun", 40, bold=False)
font2 = pygame.font.SysFont("SimSun", 25, bold=False)
font3 = pygame.font.SysFont("SimSun", 35, bold=True)

text_1 = font3.render("画图", True, (0, 0, 0))
text_2 = font3.render("功能", True, (0, 0, 0))

BG_image_initial = pygame.image.load('BG_pic.png')
BG_image = pygame.transform.scale(BG_image_initial, SIZE).convert_alpha()

def subtract_images(image1, image2):
    global sub_image, sub_image_path, image_needto_save, image_is_sub
    sub_image = pygame.Surface(image1.get_size())
    for x in range(image1.get_width()):
        for y in range(image1.get_height()):
            pixel1 = image1.get_at((x, y))
            pixel2 = image2.get_at((x, y))
            new_pixel = (
                #max(0, min(255 + pixel2[0] - pixel1[0], 255)),
                #max(0, min(255 + pixel2[0] - pixel1[0], 255)),
                #max(0, min(255 + pixel2[0] - pixel1[0], 255)),
                max(0, min(pixel1[0] - pixel2[0], 255)),
                max(0, min(pixel1[0] - pixel2[0], 255)),
                max(0, min(pixel1[0] - pixel2[0], 255)),
                pixel1[3]
            )
            sub_image.set_at((x, y), new_pixel)
    full_path = os.path.dirname(ini_image_path)
    sub_image_path = full_path + f"/{folder_name}/sub_image_temp.png"
    # print(f"测试:{sub_image_path}")
    if not os.path.exists(os.path.dirname(sub_image_path)):
        os.makedirs(os.path.dirname(sub_image_path))
    pygame.image.save(sub_image, sub_image_path)
    image_needto_save = sub_image_path
    image_is_sub = True
    
    return sub_image

def check_and_report_image_size(image1, image2):
    global image_size_warning_printed, enable_sub
    try:
        [ini_image_x, ini_image_y] = image1.get_size()
        [bottom_image_x, bottom_image_y] = image2.get_size()
        if ini_image_x != bottom_image_x or ini_image_y != bottom_image_y:
            if not image_size_warning_printed: 
                print("图像大小不一致，无法进行图像相减！")
                image_size_warning_printed = True 
        if ini_image_x == bottom_image_x and ini_image_y == bottom_image_y:
            button3.is_pressed = False  
            if not image_size_warning_printed: 
                print("图像匹配!")
            image_size_warning_printed = True 
            enable_sub = True
    except AttributeError:
        if not image_size_warning_printed: 
            print("图片尺寸读取失败，请确保所有图像均已正确加载。")
            image_size_warning_printed = True

class Button:
    def __init__(self, x, y, width, height, text, action=None):
        self.x, self.y = x, y
        self.width, self.height = width, height
        self.text = text
        self.action = action 
        self.color = (255, 255, 255)
        self.hover_color = (220, 220, 220) 
        self.is_hovered = False
        self.is_pressed = False
        self.clicked_once = False

    def draw(self, screen):
        if self.is_pressed:
            color = (180, 180, 180)
        elif self.is_hovered:
            color = self.hover_color
        else:
            color = self.color

        pygame.draw.rect(screen, (0, 0, 0), pygame.Rect(self.x, self.y, self.width, self.height), 2)

        pygame.draw.rect(screen, color, pygame.Rect(self.x + 1, self.y + 1, self.width - 2, self.height - 2))

        text_surface = font2.render(self.text, True, (0, 0, 0))
        screen.blit(text_surface, (self.x + (self.width - text_surface.get_width()) // 2,
                                self.y + (self.height - text_surface.get_height()) // 2))

    def update(self, mouse_pos, events):
        button_rect = pygame.Rect(self.x, self.y, self.width, self.height)
        
        self.is_hovered = button_rect.collidepoint(mouse_pos)
        
        for event in events:
            if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                if button_rect.collidepoint(event.pos) and not self.clicked_once and self.action is not None:
                    self.clicked_once = True
                    self.action()
                    return
                
def reszie_image(image_path):
    global full_path, index, image_needto_save
    image_temp = pygame.image.load(image_path)
    width, height = image_temp.get_size()
    scale = 1
    if width > max_x:
        scale = width / max_x
        if height / scale > max_y: 
            scale = height / max_y
    image_resized = pygame.transform.scale(image_temp, (int(width / scale), int(height / scale)))

    full_path = os.path.join(os.getcwd(), folder_name)
    if not os.path.exists(full_path):
        os.makedirs(full_path) 
        print(f"文件夹'{folder_name}'已成功创建。")
    else:
        print(f"文件夹'{folder_name}'已经存在。")

    save_path = os.path.join(full_path, f"temp{index}.png")
    index += 1
    print(f"该临时文件路径为:{save_path}")
    image_needto_save = save_path
    pygame.image.save(image_resized, save_path)
    return save_path

class Image_button:
    def __init__(self, image_path):
        self.image_path = image_path
        self.image = pygame.image.load(image_path)
        self.width, self.height = self.image.get_size()
    
    def resize(self):
        pygame.image.load(reszie_image(self.image_path))

def choose_file():
    root = tk.Tk()
    root.withdraw()
    file_path = filedialog.askopenfilename(filetypes=filetypes_images)
    root.destroy()  
    return file_path
def function1(): 
    global ini_image, image1, image_size_warning_printed, ini_image_path
    print("按钮：选择图片 clicked!")
    ini_image_path = choose_file()

    if ini_image_path:
        print(f"Selected file: {ini_image_path}")
        ini_image = pygame.image.load(ini_image_path)
        image_size_warning_printed = False 
    else:
        print("No initial image selected!")

    image1 = Image_button(ini_image_path)
    image1.resize()

def function2():
    global bottom_image, image_size_warning_printed, image2, bottom_image_path
    print("按钮：选择底片 clicked!")
    bottom_image_path = choose_file()
    
    if bottom_image_path:
        print(f"Selected file: {bottom_image_path}")
        bottom_image = pygame.image.load(bottom_image_path)
    else:
        print("No bottom image selected!")

    image2 = Image_button(bottom_image_path)
    image2.resize()

def function3(): 
    global ini_image, bottom_image, button3, sub_image, enable_sub, sub_image_path
    print("按钮：移除底片 clicked!")
    if ini_image is not None and bottom_image is not None:
        check_and_report_image_size(ini_image, bottom_image)
        if enable_sub:
            sub_image = subtract_images(ini_image, bottom_image)
            enable_sub = False
            print("图像相减完成！")
    else:
        print("缺少图像，无法进行图像相减！")

def function4():
    global ini_image
    image1 = Image_button(ini_image_path)
    image1.resize()
    
def function5():
    global bottom_image
    image2 = Image_button(bottom_image_path)
    image2.resize()

def function6(): 
    global sub_image
    image3 = Image_button(sub_image_path)
    image3.resize()

def function_save_image():
    global image_needto_save
    root = tk.Tk()
    root.withdraw()

    file_path = filedialog.asksaveasfilename(
        defaultextension = ".png",  
        filetypes = filetypes_images,  
        title = "保存当前显示的图片"  
    )
    if file_path:
        try:
            image_needto_save = pygame.image.load(image_needto_save)
            pygame.image.save(image_needto_save, file_path)
            print(f"Image saved successfully at {file_path}")
        except Exception as e:
            print(f"Error saving image: {e}")

def function7():
    global image_needto_save, index, image_needto_save, image_is_gray
    img = Image.open(sub_image_path)
    gray_image_temp = img.convert("L")
    gray_image_path = full_path + f"\\temp{index}.png"
    index += 1
    gray_image_temp.save(gray_image_path)
    print(f"灰度化图片保存路径为:{gray_image_path}")
    image_needto_save = gray_image_path
    image_is_gray = True

def function8():
    global image_needto_save, index, image_needto_save
    image = Image.open(sub_image_path)
    sharpened_image = image.filter(ImageFilter.UnsharpMask(radius=2, percent=150, threshold=3))
    sharpened_image_path = full_path + f"\\temp{index}.png"
    index += 1
    enhancer = ImageEnhance.Contrast(sharpened_image)
    contrast_enhanced_image = enhancer.enhance(1.5)
    contrast_enhanced_image.save(sharpened_image_path)
    print(f"锐化对比度图片保存路径为:{sharpened_image_path}")
    image_needto_save = sharpened_image_path


button1 = Button(30, 120, 150, 45, "选择图片", action=function1)
button2 = Button(30, 190, 150, 45, "选择底片", action=function2)
button3 = Button(30, 260, 150, 45, "图像相减", action=function3)
button4 = Button(WIDTH-150-30, 120, 150, 45, "原图", action=function4)
button5 = Button(WIDTH-150-30, 190, 150, 45, "底片", action=function5)
button6 = Button(WIDTH-150-30, 260, 150, 45, "相减", action=function6)
button_save_image = Button(WIDTH/2 - 75, HEIGHT - 45 - 30, 150, 45, "保存当前图片", action=function_save_image)
button7 = Button(WIDTH/2 - 75, HEIGHT - 215 + 75, 150, 45, "灰度化", action=function7)
button8 = Button(WIDTH/2 - 75, HEIGHT - 280 + 75, 150, 45, "对比度", action=function8)

running = True
image_size_warning_printed = False
enable_sub = False
image_is_sub = False 
image_is_gray = False
i = 0

while running:
    screen.fill(BG_COLOR)
    screen.blit(BG_image, (0, 0))
    screen.blit(text_2, (105 - text_2.get_width()/2, 50))
    screen.blit(text_1, (WIDTH - 75 - 30 - text_1.get_width()/2, 50))

    mouse_pos = pygame.mouse.get_pos()
    events = pygame.event.get()

    if not image_is_sub and full_path:
        sub_image_path = full_path + f"\\temp{0}.png"
    if image_is_gray and full_path:
        if i < 1:
            print(f"正常:{sub_image_path}")
        i += 1
        sub_image_path = image_needto_save

    buttons = [button1, button2, button3, button4, button5, button6, button_save_image, \
            button7, button8]
    for event in events:
        if event.type == QUIT:
            try:
                shutil.rmtree(full_path)
                print(f"Directory '{full_path}' has been successfully removed.")
            except FileNotFoundError:
                print(f"Directory '{full_path}' does not exist.")
            except PermissionError:
                print(f"Permission denied while trying to remove directory '{full_path}'.")
            except Exception as e:
                print(f"An unexpected error occurred: {str(e)}")
            pygame.quit()
            sys.exit()
        elif event.type == pygame.MOUSEBUTTONUP:
            for button in buttons:
                button.clicked_once = False

    for button in buttons:
        button.update(mouse_pos, events)
        button.draw(screen)
    
    if ini_image is not None and bottom_image is not None:
        check_and_report_image_size(ini_image, bottom_image)

    pygame.display.update()
    clock.tick(FPS)

pygame.quit()