import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk
import cv2
import numpy as np
import logging

logging.basicConfig(format='%(asctime)s:%(levelname)s:%(message)s', level=logging.INFO)

class App:
    def __init__(self, root):
        self.root = root
        self.img = None
        self.mask = None
        self.file_path = None
        self.need_refresh = False
        self.current_method = None
        self.create_widgets()
        self.bind_brush_action()
        self.prev_pts = []
        self.pen_strokes = []

    def create_widgets(self):
        self.load_button = tk.Button(self.root, text="载入图片", command=self.load_image)
        self.load_button.pack()

        self.image_frame = tk.Frame(self.root)
        self.image_frame.pack()

        self.canvas_width = 350
        self.canvas_height = int(self.canvas_width*0.75)
        self.canvas1 = tk.Canvas(self.image_frame, width=self.canvas_width, height=self.canvas_height)
        self.canvas1.pack(side=tk.LEFT)
        self.canvas2 = tk.Canvas(self.image_frame, width=self.canvas_width, height=self.canvas_height)
        self.canvas2.pack(side=tk.LEFT)

        self.button_frame = tk.Frame(self.root)
        self.button_frame.pack()

        self.fmm_button = tk.Button(self.button_frame, text="FMM", command=self.apply_fmm)
        self.fmm_button.pack(side=tk.LEFT)
        self.ns_button = tk.Button(self.button_frame, text="NS", command=self.apply_ns)
        self.ns_button.pack(side=tk.LEFT)

        self.undo_button = tk.Button(self.button_frame, text="Undo", command=self.undo)
        self.undo_button.pack(side=tk.LEFT)

        self.refresh_button = tk.Button(self.button_frame, text="Refresh", command=self.refresh)
        self.refresh_button.pack(side=tk.LEFT)
        
        self.slider = tk.Scale(self.root, from_=3, to=51, orient=tk.HORIZONTAL, length=200, resolution=2,
                            command=self.slider_update, label="窗口大小")
        self.slider.set(7)
        self.slider.pack()

    def bind_brush_action(self):
        self.canvas1.bind("<B1-Motion>", self.draw_line)
        self.canvas1.bind("<ButtonRelease-1>", self.reset_prev_point)

    def load_img_mask(self):
        logging.info(self.file_path)
        self.img = cv2.imread(self.file_path)
        if self.img is None:
            logging.error("Failed to load image. Please check the file path and ensure it is a valid image.")
            return
        
        logging.info(self.img.shape)
        self.mask = np.zeros(self.img.shape[:2], np.uint8)
        h, w = self.img.shape[:2]

        ratio = min(self.canvas_height / h, self.canvas_width / w)
        self.img = cv2.resize(self.img, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_AREA)
        self.mask = cv2.resize(self.mask, (int(w * ratio), int(h * ratio)), interpolation=cv2.INTER_AREA)

    def load_image(self):
        if not self.need_refresh:
            self.file_path = filedialog.askopenfilename()
            logging.info("载入图像: %s", self.file_path.split('/')[-1])
        else:
            self.need_refresh = False

        self.pen_strokes = []
        self.canvas2.delete("all")

        if self.file_path:
            self.load_img_mask()
            if self.img is not None:
                pil_img = Image.fromarray(cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB))
                self.photo = ImageTk.PhotoImage(pil_img)
                self.canvas1.create_image(0, 0, anchor=tk.NW, image=self.photo)

    def reset_prev_point(self, event):
        self.pen_strokes.append(self.prev_pts)
        logging.info("已经保存了: %s 笔画", len(self.pen_strokes))
        self.prev_pts = []

    def draw_line(self, event):
        pt = (event.x, event.y)
        if self.img is not None:
            if self.prev_pts:
                pen_width = 5
                cv2.line(self.img, self.prev_pts[-1], pt, (255, 255, 255), pen_width)
                cv2.line(self.mask, self.prev_pts[-1], pt, 255, pen_width)
                pil_img = Image.fromarray(cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB))
                self.photo = ImageTk.PhotoImage(pil_img)
                self.canvas1.create_image(0, 0, anchor=tk.NW, image=self.photo)
            self.prev_pts.append(pt)

    def apply_fmm(self):
        self.current_method = 'fmm'
        if self.img is not None and self.mask is not None:
            res = cv2.inpaint(self.img, self.mask, self.slider.get(), cv2.INPAINT_TELEA)
            pil_img = Image.fromarray(cv2.cvtColor(res, cv2.COLOR_BGR2RGB))
            after_photo = ImageTk.PhotoImage(pil_img)
            self.canvas2.create_image(0, 0, anchor=tk.NW, image=after_photo)
            self.canvas2.image = after_photo
            logging.info("在图像上应用 FMM 算法")

    def apply_ns(self):
        self.current_method = 'ns'
        if self.img is not None and self.mask is not None:
            res = cv2.inpaint(self.img, self.mask, self.slider.get(), cv2.INPAINT_NS)
            pil_img = Image.fromarray(cv2.cvtColor(res, cv2.COLOR_BGR2RGB))
            after_photo = ImageTk.PhotoImage(pil_img)
            self.canvas2.create_image(0, 0, anchor=tk.NW, image=after_photo)
            self.canvas2.image = after_photo
            logging.info("在图像上应用 NS 算法")

    def __init__(self, root):
        self.root = root
        self.canvas1 = tk.Canvas(root)
        self.canvas1.pack()

    def refresh(self):
        self.need_refresh = True
        self.current_method = None
        self.load_image()

    def undo(self):
        logging.info("Undo strokes")
        # 其他代码保持不变

    def slider_update(self, val):
        if self.current_method == 'fmm':
            self.apply_fmm()
        elif self.current_method == 'ns':
            self.apply_ns()
        else:
            pass

root = tk.Tk()
app = App(root)
root.mainloop()