"""
calibrate.py
碗尺寸标定工具 - 交互式 GUI

使用流程：
  1. 分别为大/中/小碗加载若干样本图片
  2. 工具自动检测每张图的碗直径并显示预览
  3. 确认后计算阈值并写入 config.json
  4. 可选：输入各碗实际重量（克）
"""

import os
import sys
import json
import threading
import cv2
import numpy as np
import customtkinter as ctk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk

# 确保能找到项目根目录下的模块
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from bowl_size.bowl_detector import BowlDetector, BowlSize

ctk.set_appearance_mode("System")
ctk.set_default_color_theme("blue")

CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")
PREVIEW_SIZE = (300, 300)


class CalibrationApp:
    def __init__(self, root):
        self.root = root
        self.root.title("碗尺寸标定工具")
        self.root.geometry("1100x750")
        self.root.resizable(True, True)

        self.detector = BowlDetector(CONFIG_PATH)

        # 各尺寸的检测直径列表
        self.samples = {"small": [], "medium": [], "large": []}
        self.sample_images = {"small": [], "medium": [], "large": []}
        self.current_size = "small"

        self._build_ui()

    # ---------------------------------------------------------------- #
    #  UI 构建
    # ---------------------------------------------------------------- #

    def _build_ui(self):
        self.root.grid_rowconfigure(0, weight=1)
        self.root.grid_columnconfigure(0, weight=1)

        main = ctk.CTkFrame(self.root)
        main.grid(row=0, column=0, padx=15, pady=15, sticky="nsew")
        main.grid_rowconfigure(1, weight=1)
        main.grid_columnconfigure(0, weight=1)
        main.grid_columnconfigure(1, weight=2)

        # ── 标题 ────────────────────────────────────────────────────
        title = ctk.CTkLabel(main, text="碗尺寸标定工具",
                             font=ctk.CTkFont(size=22, weight="bold"))
        title.grid(row=0, column=0, columnspan=2, pady=(10, 15))

        # ── 左侧控制面板 ─────────────────────────────────────────────
        left = ctk.CTkFrame(main)
        left.grid(row=1, column=0, padx=(10, 5), pady=10, sticky="nsew")
        left.grid_columnconfigure(0, weight=1)
        self._build_left_panel(left)

        # ── 右侧预览区域 ─────────────────────────────────────────────
        right = ctk.CTkFrame(main)
        right.grid(row=1, column=1, padx=(5, 10), pady=10, sticky="nsew")
        right.grid_rowconfigure(1, weight=1)
        right.grid_columnconfigure(0, weight=1)
        self._build_right_panel(right)

    def _build_left_panel(self, parent):
        row = 0

        # 当前选择的碗尺寸
        ctk.CTkLabel(parent, text="当前标定碗型：",
                     font=ctk.CTkFont(size=14, weight="bold")).grid(
            row=row, column=0, padx=15, pady=(15, 5), sticky="w")
        row += 1

        self.size_var = ctk.StringVar(value="small")
        size_options = [("小碗 (Small)", "small"),
                        ("中碗 (Medium)", "medium"),
                        ("大碗 (Large)", "large")]
        for text, val in size_options:
            rb = ctk.CTkRadioButton(parent, text=text, variable=self.size_var,
                                    value=val, command=self._on_size_changed)
            rb.grid(row=row, column=0, padx=25, pady=3, sticky="w")
            row += 1

        ctk.CTkFrame(parent, height=1, fg_color="gray30").grid(
            row=row, column=0, padx=10, pady=10, sticky="ew")
        row += 1

        # 加载图片按钮
        ctk.CTkButton(parent, text="📂 加载该碗型图片", height=40,
                      command=self._load_images).grid(
            row=row, column=0, padx=15, pady=5, sticky="ew")
        row += 1

        ctk.CTkButton(parent, text="🗑 清除该碗型数据", height=35,
                      fg_color="gray40", hover_color="gray30",
                      command=self._clear_current).grid(
            row=row, column=0, padx=15, pady=5, sticky="ew")
        row += 1

        ctk.CTkFrame(parent, height=1, fg_color="gray30").grid(
            row=row, column=0, padx=10, pady=10, sticky="ew")
        row += 1

        # 各尺寸统计
        ctk.CTkLabel(parent, text="检测统计：",
                     font=ctk.CTkFont(size=13, weight="bold")).grid(
            row=row, column=0, padx=15, pady=(5, 3), sticky="w")
        row += 1

        self.stat_labels = {}
        for sz, name in [("small", "小碗"), ("medium", "中碗"), ("large", "大碗")]:
            lbl = ctk.CTkLabel(parent, text=f"{name}: 0 张，均值 0px",
                               font=ctk.CTkFont(size=12), text_color="gray60")
            lbl.grid(row=row, column=0, padx=25, pady=2, sticky="w")
            self.stat_labels[sz] = lbl
            row += 1

        ctk.CTkFrame(parent, height=1, fg_color="gray30").grid(
            row=row, column=0, padx=10, pady=10, sticky="ew")
        row += 1

        # 碗重量输入
        ctk.CTkLabel(parent, text="碗的重量（克）：",
                     font=ctk.CTkFont(size=13, weight="bold")).grid(
            row=row, column=0, padx=15, pady=(5, 3), sticky="w")
        row += 1

        self.weight_entries = {}
        for sz, name in [("small", "小碗"), ("medium", "中碗"), ("large", "大碗")]:
            frame = ctk.CTkFrame(parent, fg_color="transparent")
            frame.grid(row=row, column=0, padx=15, pady=2, sticky="ew")
            frame.grid_columnconfigure(1, weight=1)
            ctk.CTkLabel(frame, text=f"{name}: ", width=55).grid(row=0, column=0, sticky="w")
            entry = ctk.CTkEntry(frame, placeholder_text="0", width=80)
            entry.grid(row=0, column=1, sticky="ew")
            ctk.CTkLabel(frame, text=" 克").grid(row=0, column=2, sticky="w")
            self.weight_entries[sz] = entry
            row += 1

        ctk.CTkFrame(parent, height=1, fg_color="gray30").grid(
            row=row, column=0, padx=10, pady=10, sticky="ew")
        row += 1

        # 生成标定按钮
        self.calibrate_btn = ctk.CTkButton(
            parent, text="✅ 生成标定配置", height=45,
            font=ctk.CTkFont(size=14, weight="bold"),
            fg_color="#2B8A3E", hover_color="#1F6B2E",
            command=self._run_calibration)
        self.calibrate_btn.grid(row=row, column=0, padx=15, pady=8, sticky="ew")
        row += 1

        self.status_label = ctk.CTkLabel(parent, text="请先加载各碗型图片",
                                         text_color="gray60",
                                         font=ctk.CTkFont(size=11),
                                         wraplength=220)
        self.status_label.grid(row=row, column=0, padx=15, pady=5, sticky="w")

    def _build_right_panel(self, parent):
        ctk.CTkLabel(parent, text="图片预览与检测结果",
                     font=ctk.CTkFont(size=14, weight="bold")).grid(
            row=0, column=0, pady=(10, 5))

        self.preview_frame = ctk.CTkScrollableFrame(parent)
        self.preview_frame.grid(row=1, column=0, padx=10, pady=(0, 10), sticky="nsew")

        self.preview_widgets = []

    # ---------------------------------------------------------------- #
    #  事件处理
    # ---------------------------------------------------------------- #

    def _on_size_changed(self):
        self.current_size = self.size_var.get()
        self._refresh_preview()

    def _load_images(self):
        paths = filedialog.askopenfilenames(
            title="选择图片（可多选）",
            filetypes=[("图片文件", "*.jpg *.jpeg *.png *.bmp *.webp"), ("所有文件", "*.*")]
        )
        if not paths:
            return

        sz = self.size_var.get()
        self._set_status(f"正在检测 {len(paths)} 张图片...")

        def worker():
            for path in paths:
                result = self.detector.detect(path, draw_debug=True)
                # 标定阶段：只要检测到直径就算成功（不需要完成分类）
                if result.pixel_diameter > 0:
                    self.samples[sz].append(result.pixel_diameter)
                self.sample_images[sz].append((path, result))
            self.root.after(0, self._after_load)

        threading.Thread(target=worker, daemon=True).start()

    def _after_load(self):
        self._update_stats()
        self._refresh_preview()
        self._set_status("加载完成，请继续加载其他碗型或生成标定配置。")

    def _clear_current(self):
        sz = self.size_var.get()
        self.samples[sz].clear()
        self.sample_images[sz].clear()
        self._update_stats()
        self._refresh_preview()

    def _run_calibration(self):
        for sz in ["small", "medium", "large"]:
            if not self.samples[sz]:
                messagebox.showwarning("数据不足",
                                       f"{'小中大'[['small','medium','large'].index(sz)]}碗还没有有效样本！")
                return

        small_avg = np.mean(self.samples["small"])
        medium_avg = np.mean(self.samples["medium"])
        large_avg = np.mean(self.samples["large"])

        # 确保顺序正确
        sorted_avgs = sorted([(small_avg, "small"), (medium_avg, "medium"), (large_avg, "large")],
                             key=lambda x: x[0])
        names = [x[1] for x in sorted_avgs]
        avgs = [x[0] for x in sorted_avgs]

        if names != ["small", "medium", "large"]:
            messagebox.showwarning("数据异常",
                                   f"检测到的碗尺寸顺序异常：\n小碗均值={small_avg:.0f}px\n"
                                   f"中碗均值={medium_avg:.0f}px\n大碗均值={large_avg:.0f}px\n\n"
                                   "请检查图片是否正确分类。")
            return

        # 阈值取相邻两档的中点
        small_max = (avgs[0] + avgs[1]) / 2
        medium_max = (avgs[1] + avgs[2]) / 2

        # 读取碗重
        weights = {}
        for sz in ["small", "medium", "large"]:
            try:
                weights[sz] = float(self.weight_entries[sz].get() or "0")
            except ValueError:
                weights[sz] = 0.0

        # 加载并更新配置
        if os.path.exists(CONFIG_PATH):
            with open(CONFIG_PATH, "r", encoding="utf-8") as f:
                config = json.load(f)
        else:
            config = {}

        config["calibrated"] = True
        config["thresholds"] = {
            "small_max_diameter": round(small_max, 1),
            "medium_max_diameter": round(medium_max, 1),
            "description": "像素直径阈值：<=small_max为小碗，<=medium_max为中碗，其余为大碗",
        }
        config["calibration_samples"] = {
            "small": [round(d, 1) for d in self.samples["small"]],
            "medium": [round(d, 1) for d in self.samples["medium"]],
            "large": [round(d, 1) for d in self.samples["large"]],
        }
        config["calibration_averages"] = {
            "small": round(small_avg, 1),
            "medium": round(medium_avg, 1),
            "large": round(large_avg, 1),
        }
        config["bowl_weights_grams"] = weights

        with open(CONFIG_PATH, "w", encoding="utf-8") as f:
            json.dump(config, f, ensure_ascii=False, indent=2)

        messagebox.showinfo(
            "标定成功",
            f"标定完成！配置已保存到 config.json\n\n"
            f"小碗均值: {small_avg:.0f}px\n"
            f"中碗均值: {medium_avg:.0f}px\n"
            f"大碗均值: {large_avg:.0f}px\n\n"
            f"小碗阈值上限: {small_max:.0f}px\n"
            f"中碗阈值上限: {medium_max:.0f}px"
        )
        self._set_status("✅ 标定配置已保存！")

    # ---------------------------------------------------------------- #
    #  辅助方法
    # ---------------------------------------------------------------- #

    def _update_stats(self):
        for sz, name in [("small", "小碗"), ("medium", "中碗"), ("large", "大碗")]:
            diameters = self.samples[sz]
            n = len(diameters)
            if n > 0:
                avg = np.mean(diameters)
                text = f"{name}: {n} 张，均值 {avg:.0f}px"
                self.stat_labels[sz].configure(text=text, text_color="green")
            else:
                self.stat_labels[sz].configure(text=f"{name}: 0 张", text_color="gray60")

    def _refresh_preview(self):
        for w in self.preview_widgets:
            w.destroy()
        self.preview_widgets.clear()

        sz = self.size_var.get()
        items = self.sample_images.get(sz, [])

        if not items:
            lbl = ctk.CTkLabel(self.preview_frame, text="暂无图片，请点击左侧[加载]按钮",
                               text_color="gray60")
            lbl.grid(row=0, column=0, padx=20, pady=30)
            self.preview_widgets.append(lbl)
            return

        col_count = 3
        for i, (path, result) in enumerate(items):
            r, c = divmod(i, col_count)
            card = ctk.CTkFrame(self.preview_frame)
            card.grid(row=r, column=c, padx=5, pady=5)

            # 图片
            if result.debug_image is not None:
                display_img = result.debug_image
            else:
                display_img = cv2.imdecode(
                    np.fromfile(path, dtype=np.uint8), cv2.IMREAD_COLOR)

            if display_img is not None:
                rgb = cv2.cvtColor(display_img, cv2.COLOR_BGR2RGB)
                pil = Image.fromarray(rgb)
                pil.thumbnail(PREVIEW_SIZE)
                ctk_img = ctk.CTkImage(pil, size=pil.size)
                img_lbl = ctk.CTkLabel(card, image=ctk_img, text="")
                img_lbl.image = ctk_img
                img_lbl.pack(padx=5, pady=(5, 2))

            # 标注信息：只要直径 > 0 就算检测成功
            if result.pixel_diameter > 0:
                info = f"直径: {result.pixel_diameter:.0f}px"
                color = "green"
            else:
                info = "检测失败"
                color = "red"
            ctk.CTkLabel(card, text=info, text_color=color,
                         font=ctk.CTkFont(size=11)).pack(padx=5, pady=(0, 5))

            self.preview_widgets.append(card)

    def _set_status(self, text: str):
        self.status_label.configure(text=text)


def main():
    root = ctk.CTk()
    app = CalibrationApp(root)
    root.mainloop()


if __name__ == "__main__":
    main()