main.py

import csv
import json
import random
from pathlib import Path
import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk


# 作者信息
__author__ = "Euplectellashen"
__version__ = "1.0.0"


class ImageReviewer:
    def __init__(self, master, image_dir=None):
        self.master = master
        master.title(f"植物记忆练习 - Made by {__author__}")

        # 默认图片目录（workspace 下的 Pic）
        default_dir = Path(__file__).parent / "Pic"
        self.image_dir = Path(image_dir) if image_dir else default_dir

        self.supported_ext = ('.jpg', '.jpeg', '.png', '.bmp', '.gif')
        self.images = []
        self.load_images()

        # persistent counts file
        self.counts_path = Path(__file__).parent / 'counts.json'
        self.counts = self.load_counts()

        # persistent drag offsets file
        self.drag_offsets_path = Path(__file__).parent / 'drag_offsets.json'
        self.drag_offsets = self.load_drag_offsets()

        # UI state
        self.index = 0
        self.zoomed = True  # start zoomed to hide标签
        self.zoom_factor = 1.4
        self.label_position = 'bottom-right'  # 标签位置
        self.focus_unremembered = tk.BooleanVar(value=False)

        # 拖曳相关状态
        self.drag_start_x = 0
        self.drag_start_y = 0
        self.drag_offset_x = 0  # 当前拖曳偏移量
        self.drag_offset_y = 0
        # per-image drag offsets: {filename: [x, y]}
        self.is_dragging = False
        self.drag_step = 20  # 键盘微调步长（像素）

        # Build UI
        self.canvas_w = 900
        self.canvas_h = 600
        self.canvas = tk.Canvas(master, width=self.canvas_w, height=self.canvas_h, bg='black')
        self.canvas.pack(fill=tk.BOTH, expand=True)

        ctrl_frame = tk.Frame(master)
        ctrl_frame.pack(fill=tk.X)

        tk.Button(ctrl_frame, text='上一张 ◀', command=self.prev_image).pack(side=tk.LEFT)
        tk.Button(ctrl_frame, text='下一张 ▶', command=self.next_image).pack(side=tk.LEFT)
        tk.Button(ctrl_frame, text='随机打乱', command=self.shuffle_images).pack(side=tk.LEFT)

        tk.Button(ctrl_frame, text='记住了', command=self.mark_remembered).pack(side=tk.LEFT, padx=6)
        tk.Button(ctrl_frame, text='没记住', command=self.mark_unremembered).pack(side=tk.LEFT)

        tk.Checkbutton(ctrl_frame, text='只练没记住的', variable=self.focus_unremembered, command=self.toggle_focus).pack(side=tk.LEFT, padx=8)

        tk.Label(ctrl_frame, text='放大倍数:').pack(side=tk.LEFT)
        self.zoom_entry = tk.Entry(ctrl_frame, width=6)
        self.zoom_entry.insert(0, str(self.zoom_factor))
        self.zoom_entry.pack(side=tk.LEFT)
        tk.Label(ctrl_frame, text='标签位置:').pack(side=tk.LEFT, padx=6)
        self.pos_var = tk.StringVar(value=self.label_position)
        tk.OptionMenu(ctrl_frame, self.pos_var, 'bottom-right', 'bottom-left', 'top-right', 'top-left', command=self.change_label_pos).pack(side=tk.LEFT)

        # 第二行控制栏：导出和重置
        ctrl_frame2 = tk.Frame(master)
        ctrl_frame2.pack(fill=tk.X)
        tk.Button(ctrl_frame2, text='导出 CSV', command=self.export_csv).pack(side=tk.LEFT, padx=4)
        tk.Button(ctrl_frame2, text='重置所有计数', command=self.reset_counts).pack(side=tk.LEFT, padx=4)
        tk.Button(ctrl_frame2, text='选择图片目录', command=self.choose_directory).pack(side=tk.LEFT, padx=4)
        tk.Button(ctrl_frame2, text='重置拖曳位置', command=self.reset_drag_position).pack(side=tk.LEFT, padx=4)

        # 作者署名（显眼）
        author_label = tk.Label(ctrl_frame2, text=f'Made by {__author__}', fg='#0066CC', font=('Arial', 10, 'bold'))
        author_label.pack(side=tk.RIGHT, padx=10)

        # status bar
        self.status = tk.Label(master, text='', anchor='w')
        self.status.pack(fill=tk.X)

        # Events
        self.canvas.bind('<Button-1>', self.on_left_click)
        self.canvas.bind('<B1-Motion>', self.on_drag)
        self.canvas.bind('<ButtonRelease-1>', self.on_drag_end)
        master.bind('<Left>', self.on_key_left)
        master.bind('<Right>', self.on_key_right)
        master.bind('<Up>', self.on_key_up)
        master.bind('<Down>', self.on_key_down)
        master.bind('<Shift-Left>', lambda e: self.prev_image())
        master.bind('<Shift-Right>', lambda e: self.next_image())

        # prepare current filtered list
        self.filtered_indices = list(range(len(self.images)))
        self.update_filtered_indices()

        # show first
        if self.images:
            self.show_image()
        else:
            messagebox.showinfo('提示', f'目录 {self.image_dir} 中没有找到图片。')

    def load_images(self):
        if not self.image_dir.exists():
            self.images = []
            return
        files = [p for p in sorted(self.image_dir.iterdir()) if p.suffix.lower() in self.supported_ext]
        self.images = files
        random.shuffle(self.images)

    def load_counts(self):
        if self.counts_path.exists():
            try:
                return json.loads(self.counts_path.read_text(encoding='utf-8'))
            except Exception:
                return {}
        return {}

    def save_counts(self):
        try:
            self.counts_path.write_text(json.dumps(self.counts, ensure_ascii=False, indent=2), encoding='utf-8')
        except Exception as e:
            print('保存 counts 失败:', e)

    def load_drag_offsets(self):
        """从文件加载拖曳偏移量"""
        if self.drag_offsets_path.exists():
            try:
                data = json.loads(self.drag_offsets_path.read_text(encoding='utf-8'))
                # 转换为 {filename: [x, y]} 格式
                return {k: list(v) for k, v in data.items()}
            except Exception:
                return {}
        return {}

    def save_drag_offsets(self):
        """保存拖曳偏移量到文件"""
        try:
            self.drag_offsets_path.write_text(
                json.dumps(self.drag_offsets, ensure_ascii=False, indent=2), 
                encoding='utf-8'
            )
        except Exception as e:
            print('保存 drag_offsets 失败:', e)

    def get_current_path(self):
        if not self.filtered_indices:
            return None
        return self.images[self.filtered_indices[self.index]]

    def show_image(self):
        path = self.get_current_path()
        if path is None:
            self.canvas.delete('all')
            self.status.config(text='没有符合条件的图片。')
            return

        self.current_img = Image.open(path).convert('RGBA')
        # fit or zoom
        try:
            self.zoom_factor = float(self.zoom_entry.get())
        except Exception:
            self.zoom_factor = 1.4

        if self.zoomed:
            # 计算放大后的图片尺寸，用于限制拖曳范围
            self.scaled_w = int(self.current_img.width * self.zoom_factor)
            self.scaled_h = int(self.current_img.height * self.zoom_factor)
            disp = self.get_zoomed_image(self.current_img, path, self.canvas_w, self.canvas_h, 
                                         self.zoom_factor, self.pos_var.get(),
                                         self.drag_offset_x, self.drag_offset_y)
        else:
            disp = self.get_fit_image(self.current_img, self.canvas_w, self.canvas_h)

        self.tkimg = ImageTk.PhotoImage(disp)
        self.canvas.delete('all')
        self.canvas.create_image(self.canvas_w // 2, self.canvas_h // 2, image=self.tkimg)

        # status
        filename = path.name
        cnt = self.counts.get(filename, 0)
        drag_hint = " (可拖曳)" if self.zoomed else ""
        self.status.config(text=f'{filename}    记住次数: {cnt}    ({self.index+1}/{len(self.filtered_indices)}){drag_hint}')

    def get_fit_image(self, img: Image.Image, w, h):
        # fit to canvas
        img_ratio = img.width / img.height
        canvas_ratio = w / h
        if img_ratio > canvas_ratio:
            new_w = w
            new_h = int(w / img_ratio)
        else:
            new_h = h
            new_w = int(h * img_ratio)
        return img.resize((new_w, new_h), Image.Resampling.LANCZOS)

    def get_zoomed_image(self, img: Image.Image, path, w, h, zoom, label_pos, drag_offset_x=0, drag_offset_y=0):
        """放大图片并裁剪以遮挡标签区域，支持拖曳偏移"""
        scaled_w = int(img.width * zoom)
        scaled_h = int(img.height * zoom)
        scaled = img.resize((scaled_w, scaled_h), Image.Resampling.LANCZOS)

        # 中心裁剪基准
        cx = (scaled_w - w) // 2
        cy = (scaled_h - h) // 2

        # 使用手动指定的标签位置
        shift_x = max(0, (scaled_w - w) // 4)
        shift_y = max(0, (scaled_h - h) // 4)

        if label_pos == 'bottom-right':
            crop_x = cx - shift_x
            crop_y = cy - shift_y
        elif label_pos == 'bottom-left':
            crop_x = cx + shift_x
            crop_y = cy - shift_y
        elif label_pos == 'top-right':
            crop_x = cx - shift_x
            crop_y = cy + shift_y
        else:  # top-left
            crop_x = cx + shift_x
            crop_y = cy + shift_y

        # 应用拖曳偏移
        crop_x += drag_offset_x
        crop_y += drag_offset_y

        # clamp
        crop_x = max(0, min(int(crop_x), scaled_w - w))
        crop_y = max(0, min(int(crop_y), scaled_h - h))

        box = (crop_x, crop_y, crop_x + w, crop_y + h)
        return scaled.crop(box)

    def on_left_click(self, event):
        """鼠标左键按下 - 记录拖曳起始位置"""
        if self.zoomed:
            self.drag_start_x = event.x
            self.drag_start_y = event.y
            self.is_dragging = False

    def on_drag(self, event):
        """鼠标拖曳 - 实时更新拖曳偏移"""
        if not self.zoomed:
            return
        
        dx = self.drag_start_x - event.x
        dy = self.drag_start_y - event.y
        
        # 只有移动超过一定距离才认为是拖曳
        if abs(dx) > 5 or abs(dy) > 5:
            self.is_dragging = True
        
        if self.is_dragging:
            # 更新拖曳偏移量（基于当前图片的已保存偏移）
            cur = self.get_current_path()
            saved_x, saved_y = (0, 0)
            if cur is not None:
                saved_x, saved_y = self.drag_offsets.get(cur.name, (0, 0))
            self.drag_offset_x = saved_x + dx
            self.drag_offset_y = saved_y + dy
            self.show_image()

    def on_drag_end(self, event):
        """鼠标释放 - 如果没有拖曳则切换缩放状态"""
        if not self.zoomed:
            # 非放大模式下点击切换到放大模式
            self.zoomed = True
            # 恢复该图片已保存的偏移（如果有）
            cur = self.get_current_path()
            if cur is not None:
                self.drag_offset_x, self.drag_offset_y = self.drag_offsets.get(cur.name, (0, 0))
            else:
                self.drag_offset_x = 0
                self.drag_offset_y = 0
            self.show_image()
        elif not self.is_dragging:
            # 放大模式下点击（无拖曳）切换到还原模式
            self.zoomed = False
            self.show_image()
        else:
            # 拖曳结束，保存偏移量到当前图片
            cur = self.get_current_path()
            if cur is not None:
                self.drag_offsets[cur.name] = [self.drag_offset_x, self.drag_offset_y]
                self.save_drag_offsets()
        
        self.is_dragging = False

    def next_image(self):
        if not self.filtered_indices:
            return
        self.index = (self.index + 1) % len(self.filtered_indices)
        self.zoomed = True
        # 使用该图片保存的拖曳偏移量（如果有）
        cur = self.get_current_path()
        if cur is not None:
            self.drag_offset_x, self.drag_offset_y = self.drag_offsets.get(cur.name, (0, 0))
        self.show_image()

    def prev_image(self):
        if not self.filtered_indices:
            return
        self.index = (self.index - 1) % len(self.filtered_indices)
        self.zoomed = True
        # 使用该图片保存的拖曳偏移量（如果有）
        cur = self.get_current_path()
        if cur is not None:
            self.drag_offset_x, self.drag_offset_y = self.drag_offsets.get(cur.name, (0, 0))
        self.show_image()

    def mark_remembered(self):
        path = self.get_current_path()
        if path is None:
            return
        name = path.name
        self.counts[name] = self.counts.get(name, 0) + 1
        self.save_counts()
        self.show_image()

    def mark_unremembered(self):
        path = self.get_current_path()
        if path is None:
            return
        name = path.name
        self.counts[name] = max(0, self.counts.get(name, 0) - 1)
        self.save_counts()
        self.show_image()

    def toggle_focus(self):
        self.update_filtered_indices()
        self.index = 0
        self.show_image()

    def update_filtered_indices(self):
        if self.focus_unremembered.get():
            unrem = [i for i, p in enumerate(self.images) if self.counts.get(p.name, 0) == 0]
            if not unrem:
                messagebox.showinfo('提示', '所有图片都已被标记为记住。')
                self.focus_unremembered.set(False)
                self.filtered_indices = list(range(len(self.images)))
            else:
                self.filtered_indices = unrem
        else:
            self.filtered_indices = list(range(len(self.images)))

    def change_label_pos(self, val):
        self.label_position = val
        self.zoomed = True
        self.show_image()

    def shuffle_images(self):
        random.shuffle(self.images)
        self.update_filtered_indices()
        self.index = 0
        self.show_image()

    # 键盘微调偏移方法
    def on_key_left(self, event):
        """左键：放大模式下左移视图，否则上一张图片"""
        if self.zoomed:
            self.adjust_drag_offset(-self.drag_step, 0)
        else:
            self.prev_image()

    def on_key_right(self, event):
        """右键：放大模式下右移视图，否则下一张图片"""
        if self.zoomed:
            self.adjust_drag_offset(self.drag_step, 0)
        else:
            self.next_image()

    def on_key_up(self, event):
        """上键：放大模式下上移视图"""
        if self.zoomed:
            self.adjust_drag_offset(0, -self.drag_step)

    def on_key_down(self, event):
        """下键：放大模式下下移视图"""
        if self.zoomed:
            self.adjust_drag_offset(0, self.drag_step)

    def adjust_drag_offset(self, dx, dy):
        """调整拖曳偏移并保存"""
        self.drag_offset_x += dx
        self.drag_offset_y += dy
        # 保存到当前图片
        cur = self.get_current_path()
        if cur is not None:
            self.drag_offsets[cur.name] = [self.drag_offset_x, self.drag_offset_y]
            self.save_drag_offsets()
        self.show_image()

    def reset_drag_position(self):
        """重置拖曳位置到默认"""
        # 清除所有保存的拖曳位置
        self.drag_offsets.clear()
        self.save_drag_offsets()
        self.drag_offset_x = 0
        self.drag_offset_y = 0
        self.show_image()

    def export_csv(self):
        """导出记忆统计到 CSV 文件"""
        file_path = filedialog.asksaveasfilename(
            defaultextension='.csv',
            filetypes=[('CSV 文件', '*.csv'), ('所有文件', '*.*')],
            title='导出记忆统计'
        )
        if not file_path:
            return
        try:
            with open(file_path, 'w', newline='', encoding='utf-8-sig') as f:
                writer = csv.writer(f)
                writer.writerow(['图片文件名', '记住次数'])
                for img_path in self.images:
                    name = img_path.name
                    cnt = self.counts.get(name, 0)
                    writer.writerow([name, cnt])
            messagebox.showinfo('导出成功', f'已导出到 {file_path}')
        except Exception as e:
            messagebox.showerror('导出失败', str(e))

    def reset_counts(self):
        """重置所有图片的记忆计数"""
        if not messagebox.askyesno('确认重置', '确定要重置所有图片的记忆计数吗？此操作不可撤销。'):
            return
        self.counts = {}
        self.save_counts()
        self.update_filtered_indices()
        self.show_image()
        messagebox.showinfo('重置完成', '所有记忆计数已重置为 0。')

    def choose_directory(self):
        """选择新的图片目录"""
        new_dir = filedialog.askdirectory(title='选择图片目录')
        if not new_dir:
            return
        self.image_dir = Path(new_dir)
        self.load_images()
        self.drag_offsets.clear()  # 清除拖曳位置缓存
        self.update_filtered_indices()
        self.index = 0
        if self.images:
            self.show_image()
        else:
            messagebox.showinfo('提示', f'目录 {self.image_dir} 中没有找到图片。')


def main():
    root = tk.Tk()
    # make canvas responsive
    root.geometry('1000x720')
    app = ImageReviewer(root)
    root.mainloop()


if __name__ == '__main__':
    main()