import os
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from PIL import Image, ImageTk
import numpy as np
import random

try:
    import cv2  # type: ignore
except Exception:
    cv2 = None

# --- Handle dirs.py persistence ---
DIRS_FILE = os.path.join(os.path.dirname(__file__), "dirs.py")

def ensure_dirs_file():
    """Create dirs.py if not exists with empty defaults."""
    if not os.path.exists(DIRS_FILE):
        with open(DIRS_FILE, "w") as f:
            f.write('image_dir = ""\n')
            f.write('mask_dir = ""\n')
            f.write('pred_dir = ""\n')
            f.write('resize_mode = ""  # "pred" | "mask" | "custom"\n')
            f.write("resize_width = 0\n")
            f.write("resize_height = 0\n")

def load_dirs():
    """Load folder paths from dirs.py safely."""
    ensure_dirs_file()
    try:
        import importlib.util

        def to_int(val, default=0):
            try:
                return int(val)
            except Exception:
                return default

        spec = importlib.util.spec_from_file_location("dirs", DIRS_FILE)
        dirs = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(dirs)
        image_dir = getattr(dirs, "image_dir", "")
        mask_dir = getattr(dirs, "mask_dir", "")
        pred_dir = getattr(dirs, "pred_dir", "")
        resize_mode = getattr(dirs, "resize_mode", "")
        resize_width = to_int(getattr(dirs, "resize_width", 0))
        resize_height = to_int(getattr(dirs, "resize_height", 0))

        # Backfill older dirs.py files that don't have resize settings yet
        if not all(hasattr(dirs, k) for k in ("resize_mode", "resize_width", "resize_height")):
            save_dirs(image_dir, mask_dir, pred_dir, resize_mode, resize_width, resize_height)

        return image_dir, mask_dir, pred_dir, resize_mode, resize_width, resize_height
    except Exception:
        return "", "", "", "", 0, 0

def save_dirs(image_dir, mask_dir, pred_dir, resize_mode="", resize_width=0, resize_height=0):
    """Save folder paths to dirs.py."""
    with open(DIRS_FILE, "w") as f:
        f.write(f'image_dir = r"""{image_dir}"""\n')
        f.write(f'mask_dir = r"""{mask_dir}"""\n')
        f.write(f'pred_dir = r"""{pred_dir}"""\n')
        f.write(f'resize_mode = r"""{resize_mode}"""\n')
        f.write(f"resize_width = {int(resize_width) if resize_width else 0}\n")
        f.write(f"resize_height = {int(resize_height) if resize_height else 0}\n")


class ImageMaskViewer:
    def __init__(self, root):
        self.root = root
        self.root.title("Image & Mask Viewer")

        # File lists
        self.image_files = []
        self.mask_files = []
        self.pred_files = []
        self.items = []  # [{id, filename, image_path, mask_path, pred_path, iou}]
        self.index = 0

        # Right panel sorting state
        self.sort_column = "filename"
        self.sort_ascending = True
        self._suppress_table_select = False
        self._id_to_index = {}

        # Right panel search state
        self.search_var = tk.StringVar()
        self._table_filter = ""

        # Folder dirs
        (
            self.image_dir,
            self.mask_dir,
            self.pred_dir,
            self.resize_mode,
            self.resize_width,
            self.resize_height,
        ) = load_dirs()

        # If dirs exist, preload
        if self.image_dir and os.path.exists(self.image_dir):
            self.image_files = self.list_images(self.image_dir)
        if self.mask_dir and os.path.exists(self.mask_dir):
            self.mask_files = self.list_images(self.mask_dir)
        if self.pred_dir and os.path.exists(self.pred_dir):
            self.pred_files = self.list_images(self.pred_dir)

        # Zoom & pan state
        self.zoom = 1.0
        self.base_scale = 1.0
        self.pan_x = 0
        self.pan_y = 0
        self.dragging = False

        # Cached PIL images
        self.curr_image = None
        self.curr_mask = None
        self.curr_pred = None
        self.curr_overlay = None
        self.curr_pred_overlay = None

        # Colormap cache
        self.color_map = {}
        self.color_names = {}

        # Info text buffer
        self.pixel_info = "Pointer: NaN"
        self.iou_info = "IoU: NaN"
        self.dice_info = "Dice: NaN"
        self.info_text_base = ""

        # Layout: main grid on the left + file/IoU panel on the right
        self.root.grid_rowconfigure(0, weight=1)
        self.root.grid_columnconfigure(0, weight=1)
        self.root.grid_columnconfigure(1, weight=0, minsize=320)

        self.main_frame = tk.Frame(self.root, bg="black")
        self.main_frame.grid(row=0, column=0, sticky="nsew")
        self.side_frame = tk.Frame(self.root, bg="black")
        self.side_frame.grid(row=0, column=1, sticky="nsew")

        # Layout 3x2 grid (main area)
        for r in range(3):
            self.main_frame.grid_rowconfigure(r, weight=1)
        for c in range(2):
            self.main_frame.grid_columnconfigure(c, weight=1)

        # Panels (main area)
        self.info_text = self.create_panel(self.main_frame, "Info", 0, 0, text_widget=True)
        self.image_canvas = self.create_panel(self.main_frame, "Image", 0, 1)
        self.mask_canvas = self.create_panel(self.main_frame, "Mask", 1, 0)
        self.overlay_canvas = self.create_panel(self.main_frame, "Mask Overlay", 1, 1)
        self.pred_canvas = self.create_panel(self.main_frame, "Predicted Mask", 2, 0)
        self.pred_overlay_canvas = self.create_panel(self.main_frame, "Predicted Overlay", 2, 1)

        # Right-side File/IoU table
        self.file_tree = self.create_file_iou_panel(self.side_frame)
        self._clicked_row_id = None
        self.file_action_menu = tk.Menu(self.root, tearoff=0)
        self.file_action_menu.add_command(label="Save selected PNGs…", command=self.save_clicked_row_pngs)

        # Store references
        self.tk_img_image = None
        self.tk_img_mask = None
        self.tk_img_overlay = None
        self.tk_img_pred = None
        self.tk_img_pred_overlay = None

        # Menubar
        menubar = tk.Menu(root)
        filemenu = tk.Menu(menubar, tearoff=0)
        filemenu.add_command(label="Open Image Folder", command=self.load_images)
        filemenu.add_command(label="Open Mask Folder", command=self.load_masks)
        filemenu.add_command(label="Open Predicted Folder", command=self.load_preds)
        menubar.add_cascade(label="File", menu=filemenu)
        menubar.add_command(label="Clear Path", command=self.clear_paths)
        menubar.add_command(label="Reload", command=self.reload_all)
        root.config(menu=menubar)

        # Bindings
        root.bind("<Right>", self.next_image)
        root.bind("<Left>", self.prev_image)
        root.bind("r", lambda e: self.reset_view())
        root.bind("<Configure>", self.refresh_display)

        for canvas in (self.image_canvas, self.mask_canvas,
                       self.overlay_canvas, self.pred_canvas, self.pred_overlay_canvas):
            canvas.bind("<ButtonPress-1>", self.start_pan)
            canvas.bind("<B1-Motion>", self.do_pan)
            canvas.bind("<MouseWheel>", self.do_zoom)  # Windows
            canvas.bind("<Button-4>", self.do_zoom)    # Linux scroll up
            canvas.bind("<Button-5>", self.do_zoom)    # Linux scroll down
            canvas.bind("<Motion>", self.show_pixel_info)

        # If we loaded some dirs already → load first image
        self.rebuild_items()
        if self.items:
            self.index = 0
            self.reset_and_load()

    def list_images(self, folder):
        return sorted([os.path.join(folder, f)
                       for f in os.listdir(folder)
                       if f.lower().endswith(('.png', '.jpg', '.jpeg'))])

    def create_panel(self, parent, title, row, col, text_widget=False):
        frame = tk.Frame(parent, bg="black")
        frame.grid(row=row, column=col, sticky="nsew")
        frame.grid_rowconfigure(1, weight=1)
        frame.grid_columnconfigure(0, weight=1)

        label = tk.Label(frame, text=title, bg="gray20", fg="white",
                         font=("Consolas", 11, "bold"))
        label.grid(row=0, column=0, sticky="ew")

        if text_widget:
            text = tk.Text(frame, bg="black", fg="lime",
                           font=("Consolas", 10), wrap="word")
            text.grid(row=1, column=0, sticky="nsew")
            text.config(state="disabled")
            scrollbar = tk.Scrollbar(frame, command=text.yview)
            scrollbar.grid(row=1, column=1, sticky="ns")
            text.config(yscrollcommand=scrollbar.set)
            return text
        else:
            canvas = tk.Canvas(frame, bg="black")
            canvas.grid(row=1, column=0, sticky="nsew")
            return canvas

    def create_file_iou_panel(self, parent):
        frame = tk.Frame(parent, bg="black")
        frame.grid(row=0, column=0, sticky="nsew")
        parent.grid_rowconfigure(0, weight=1)
        parent.grid_columnconfigure(0, weight=1)

        header = tk.Frame(frame, bg="gray20")
        header.grid(row=0, column=0, columnspan=2, sticky="ew")
        header.grid_columnconfigure(0, weight=1)

        label = tk.Label(
            header,
            text="Files (left click opens, right click saves)",
            bg="gray20",
            fg="white",
            font=("Consolas", 11, "bold"),
            anchor="w",
        )
        label.grid(row=0, column=0, sticky="ew", padx=(6, 6), pady=3)

        reload_btn = ttk.Button(header, text="Reload", command=self.reload_all)
        reload_btn.grid(row=0, column=1, padx=(6, 2), pady=3)

        search_entry = ttk.Entry(header, textvariable=self.search_var, width=18)
        search_entry.grid(row=0, column=2, padx=2, pady=3)
        search_entry.bind("<Return>", lambda _e: self.apply_search())
        search_entry.bind("<Escape>", lambda _e: self.clear_search())

        search_btn = ttk.Button(header, text="Search", command=self.apply_search)
        search_btn.grid(row=0, column=3, padx=2, pady=3)

        files_btn = ttk.Button(header, text="Files", command=self.clear_search)
        files_btn.grid(row=0, column=4, padx=(2, 6), pady=3)

        tree = ttk.Treeview(frame, columns=("filename", "iou", "dice"), show="headings", selectmode="browse")
        tree.heading("filename", text="File Name", command=lambda: self.sort_items("filename"))
        tree.heading("iou", text="IoU", command=lambda: self.sort_items("iou"))
        tree.heading("dice", text="Dice", command=lambda: self.sort_items("dice"))
        tree.column("filename", width=220, anchor="w", stretch=True)
        tree.column("iou", width=80, anchor="e", stretch=False)
        tree.column("dice", width=80, anchor="e", stretch=False)
        tree.grid(row=1, column=0, sticky="nsew")

        yscroll = ttk.Scrollbar(frame, orient="vertical", command=tree.yview)
        yscroll.grid(row=1, column=1, sticky="ns")
        tree.configure(yscrollcommand=yscroll.set)

        frame.grid_rowconfigure(1, weight=1)
        frame.grid_columnconfigure(0, weight=1)

        tree.bind("<ButtonRelease-1>", self.on_table_click)
        tree.bind("<Button-3>", self.on_table_right_click)  # Windows/Linux
        tree.bind("<Button-2>", self.on_table_right_click)  # macOS fallback
        return tree

    def apply_search(self):
        self._table_filter = (self.search_var.get() or "").strip()
        self.refresh_file_table()
        self.select_current_in_table()

    def clear_search(self):
        self.search_var.set("")
        self._table_filter = ""
        self.refresh_file_table()
        self.select_current_in_table()

    def resize_bin01_nearest(self, bin01_np, target_size):
        tw, th = int(target_size[0]), int(target_size[1])
        if cv2 is not None:
            return cv2.resize(bin01_np, (tw, th), interpolation=cv2.INTER_NEAREST)
        return np.array(Image.fromarray(bin01_np).resize((tw, th), Image.NEAREST))

    def rebuild_items(self):
        def stem(path):
            return os.path.splitext(os.path.basename(path))[0]

        img_map = {stem(p): p for p in self.image_files}
        mask_map = {stem(p): p for p in self.mask_files}
        pred_map = {stem(p): p for p in self.pred_files}

        keys = sorted(set(img_map) | set(mask_map) | set(pred_map), key=lambda s: s.lower())
        self.items = []
        for key in keys:
            image_path = img_map.get(key)
            mask_path = mask_map.get(key)
            pred_path = pred_map.get(key)
            display_name = os.path.basename(image_path or mask_path or pred_path or key)
            iou_val, dice_val = self.compute_metrics_for_paths(mask_path, pred_path, image_path)
            self.items.append({
                "id": key,
                "filename": display_name,
                "image_path": image_path,
                "mask_path": mask_path,
                "pred_path": pred_path,
                "iou": iou_val,
                "dice": dice_val,
            })

        self.sort_items(self.sort_column, preserve_direction=True, refresh_only=True)

    def compute_metrics_for_paths(self, mask_path, pred_path, image_path=None):
        if not mask_path or not pred_path:
            return None, None
        try:
            with Image.open(pred_path) as pred_im:
                pred = pred_im.convert("L")
            pred_np = (np.array(pred) > 0).astype(np.uint8)
            pred = Image.fromarray(pred_np)
            target_size = pred.size

            with Image.open(mask_path) as mask_im:
                mask = mask_im.convert("L")
            mask_np = (np.array(mask) > 0).astype(np.uint8)
            if mask_np.shape[:2] != pred_np.shape[:2]:
                mask_np = self.resize_bin01_nearest(mask_np, target_size)

            mask_bin = mask_np.astype(np.uint8)
            pred_bin = pred_np.astype(np.uint8)
            inter = np.logical_and(mask_bin, pred_bin).sum()
            union = np.logical_or(mask_bin, pred_bin).sum()
            iou_val = float(inter / union) if union > 0 else 0.0

            denom = mask_bin.sum() + pred_bin.sum()
            dice_val = float((2.0 * inter) / denom) if denom > 0 else 0.0
            return iou_val, dice_val
        except Exception:
            return None, None

    def format_iou(self, iou_val):
        return "NaN" if iou_val is None else f"{iou_val:.4f}"

    def format_dice(self, dice_val):
        return "NaN" if dice_val is None else f"{dice_val:.4f}"

    def sort_items(self, column, preserve_direction=False, refresh_only=False):
        if not preserve_direction:
            if self.sort_column == column:
                self.sort_ascending = not self.sort_ascending
            else:
                self.sort_column = column
                self.sort_ascending = True
        else:
            self.sort_column = column

        current_id = None
        if self.items and 0 <= self.index < len(self.items):
            current_id = self.items[self.index]["id"]

        if column == "iou":
            if self.sort_ascending:
                self.items.sort(key=lambda it: (it["iou"] is None, it["iou"] if it["iou"] is not None else float("inf")))
            else:
                self.items.sort(key=lambda it: (it["iou"] is None, -(it["iou"] if it["iou"] is not None else 0.0)))
        elif column == "dice":
            if self.sort_ascending:
                self.items.sort(key=lambda it: (it.get("dice") is None, it.get("dice") if it.get("dice") is not None else float("inf")))
            else:
                self.items.sort(key=lambda it: (it.get("dice") is None, -(it.get("dice") if it.get("dice") is not None else 0.0)))
        else:
            self.items.sort(key=lambda it: it["filename"].lower(), reverse=not self.sort_ascending)

        self.refresh_file_table()

        if current_id in self._id_to_index:
            self.index = self._id_to_index[current_id]
            self.select_current_in_table()

        if not refresh_only and self.items:
            self.load_current_images()
            self.refresh_display()

    def refresh_file_table(self):
        self._id_to_index = {}
        children = self.file_tree.get_children()
        if children:
            self.file_tree.delete(*children)

        needle = (self._table_filter or "").strip().lower()
        for idx, it in enumerate(self.items):
            self._id_to_index[it["id"]] = idx
            if needle:
                hay = f"{it.get('filename', '')} {it.get('id', '')}".lower()
                if needle not in hay:
                    continue
            self.file_tree.insert(
                "",
                "end",
                iid=it["id"],
                values=(it["filename"], self.format_iou(it.get("iou")), self.format_dice(it.get("dice"))),
            )

    def select_current_in_table(self):
        if not self.items or not (0 <= self.index < len(self.items)):
            return
        iid = self.items[self.index]["id"]
        if iid not in self._id_to_index:
            return
        if not self.file_tree.exists(iid):
            return
        self._suppress_table_select = True
        try:
            self.file_tree.selection_set(iid)
            self.file_tree.see(iid)
        finally:
            self._suppress_table_select = False

    def on_table_click(self, event):
        if self._suppress_table_select:
            return
        region = self.file_tree.identify("region", event.x, event.y)
        if region != "cell":
            return
        col = self.file_tree.identify_column(event.x)
        if col != "#1":
            return
        row_id = self.file_tree.identify_row(event.y)
        if not row_id:
            return
        if row_id in self._id_to_index:
            self.index = self._id_to_index[row_id]
            self.reset_and_load()

    def on_table_right_click(self, event):
        if self._suppress_table_select:
            return
        region = self.file_tree.identify("region", event.x, event.y)
        if region != "cell":
            return
        col = self.file_tree.identify_column(event.x)
        if col != "#1":
            return
        row_id = self.file_tree.identify_row(event.y)
        if not row_id or row_id not in self._id_to_index:
            return

        self._clicked_row_id = row_id
        self._suppress_table_select = True
        try:
            self.file_tree.selection_set(row_id)
            self.file_tree.see(row_id)
        finally:
            self._suppress_table_select = False

        try:
            self.file_action_menu.tk_popup(event.x_root, event.y_root)
        finally:
            self.file_action_menu.grab_release()

    def save_clicked_row_pngs(self):
        row_id = self._clicked_row_id
        if not row_id or row_id not in self._id_to_index:
            return
        it = self.items[self._id_to_index[row_id]]

        base_path = filedialog.asksaveasfilename(
            title="Save selected as (base name)",
            defaultextension=".png",
            filetypes=[("PNG", "*.png")],
            initialfile=f"{it['id']}.png",
        )
        if not base_path:
            return

        out_dir = os.path.dirname(base_path)
        base_name = os.path.splitext(os.path.basename(base_path))[0]

        def save_as_png(src_path, out_path, mode):
            if not src_path or not os.path.exists(src_path):
                return False
            with Image.open(src_path) as im:
                im = im.convert(mode)
                im.save(out_path, format="PNG")
            return True

        saved_any = False
        errors = []

        try:
            img_out = os.path.join(out_dir, f"{base_name}_image.png")
            if save_as_png(it.get("image_path"), img_out, "RGB"):
                saved_any = True
        except Exception as e:
            errors.append(f"Image: {e}")

        try:
            mask_out = os.path.join(out_dir, f"{base_name}_mask.png")
            if save_as_png(it.get("mask_path"), mask_out, "L"):
                saved_any = True
        except Exception as e:
            errors.append(f"Mask: {e}")

        try:
            pred_out = os.path.join(out_dir, f"{base_name}_pred.png")
            if save_as_png(it.get("pred_path"), pred_out, "L"):
                saved_any = True
        except Exception as e:
            errors.append(f"Pred: {e}")

        if errors:
            messagebox.showerror("Save selected PNGs", "\n".join(errors), parent=self.root)
        elif saved_any:
            messagebox.showinfo("Save selected PNGs", "Saved PNGs successfully.", parent=self.root)
        else:
            messagebox.showinfo("Save selected PNGs", "Nothing to save (missing paths).", parent=self.root)

    def clear_paths(self):
        save_dirs("", "", "", self.resize_mode, self.resize_width, self.resize_height)
        self.clear_search()
        self.reload_all()
        messagebox.showinfo("Clear Path", "Cleared image/mask/pred folder paths.", parent=self.root)

    def reload_all(self):
        current_id = None
        if self.items and 0 <= self.index < len(self.items):
            current_id = self.items[self.index]["id"]

        (
            self.image_dir,
            self.mask_dir,
            self.pred_dir,
            self.resize_mode,
            self.resize_width,
            self.resize_height,
        ) = load_dirs()

        self.image_files = self.list_images(self.image_dir) if self.image_dir and os.path.exists(self.image_dir) else []
        self.mask_files = self.list_images(self.mask_dir) if self.mask_dir and os.path.exists(self.mask_dir) else []
        self.pred_files = self.list_images(self.pred_dir) if self.pred_dir and os.path.exists(self.pred_dir) else []

        self.rebuild_items()

        if current_id and current_id in self._id_to_index:
            self.index = self._id_to_index[current_id]
        else:
            self.index = 0

        if self.items:
            self.reset_and_load()
        else:
            self.curr_image = None
            self.curr_mask = None
            self.curr_pred = None
            self.curr_overlay = None
            self.curr_pred_overlay = None
            self.pixel_info = "Pointer: NaN"
            self.iou_info = "IoU: NaN"
            self.dice_info = "Dice: NaN"
            self.info_text_base = ""
            self.update_info_panel()
            self.refresh_display()

    def load_images(self):
        folder = filedialog.askdirectory(title="Select Image Folder")
        if folder:
            self.image_dir = folder
            self.image_files = self.list_images(folder)
            save_dirs(self.image_dir, self.mask_dir, self.pred_dir, self.resize_mode, self.resize_width, self.resize_height)
            self.index = 0
            self.rebuild_items()
            if self.items:
                self.reset_and_load()

    def load_masks(self):
        folder = filedialog.askdirectory(title="Select Mask Folder")
        if folder:
            self.mask_dir = folder
            self.mask_files = self.list_images(folder)
            save_dirs(self.image_dir, self.mask_dir, self.pred_dir, self.resize_mode, self.resize_width, self.resize_height)
            self.index = 0
            self.rebuild_items()
            if self.items:
                self.reset_and_load()

    def load_preds(self):
        folder = filedialog.askdirectory(title="Select Predicted Mask Folder")
        if folder:
            self.pred_dir = folder
            self.pred_files = self.list_images(folder)
            save_dirs(self.image_dir, self.mask_dir, self.pred_dir, self.resize_mode, self.resize_width, self.resize_height)
            self.index = 0
            self.rebuild_items()
            if self.items:
                self.reset_and_load()

    def reset_and_load(self):
        self.zoom = 1.0
        self.base_scale = 1.0
        self.pan_x = self.pan_y = 0
        self.load_current_images()
        self.refresh_display()
        self.select_current_in_table()

    def reset_view(self):
        self.zoom = 1.0
        self.base_scale = 1.0
        self.pan_x = 0
        self.pan_y = 0
        self.refresh_display()

    def load_current_images(self):
        if not self.items or not (0 <= self.index < len(self.items)):
            return

        item = self.items[self.index]

        pred_original_size = None
        mask_original_size = None
        image_original_size = None

        # Predicted Mask (never resized)
        self.curr_pred = None
        if item.get("pred_path") and os.path.exists(item["pred_path"]):
            with Image.open(item["pred_path"]) as pred_im:
                pred = pred_im.convert("L")
            pred_np = np.array(pred)
            if pred_np.max() > 0 and pred_np.max() != 255:
                pred_np = (pred_np.astype(np.float32) / pred_np.max()) * 255
                pred_np = pred_np.astype(np.uint8)
            self.curr_pred = Image.fromarray(pred_np)
            pred_original_size = self.curr_pred.size

        # Mask (may be resized to prediction size)
        self.curr_mask = None
        if item.get("mask_path") and os.path.exists(item["mask_path"]):
            with Image.open(item["mask_path"]) as mask_im:
                mask = mask_im.convert("L")
            mask_np = np.array(mask)
            if mask_np.max() > 0 and mask_np.max() != 255:
                mask_np = (mask_np.astype(np.float32) / mask_np.max()) * 255
                mask_np = mask_np.astype(np.uint8)
            self.curr_mask = Image.fromarray(mask_np)
            mask_original_size = self.curr_mask.size

        # Choose reference size:
        # 1) Prediction mask size (if available)
        # 2) Otherwise mask size
        ref_size = pred_original_size or mask_original_size
        ref_source = "prediction" if pred_original_size else ("mask" if mask_original_size else None)

        # Image (may be resized to prediction/mask reference size)
        self.curr_image = None
        if item.get("image_path") and os.path.exists(item["image_path"]):
            with Image.open(item["image_path"]) as img_im:
                self.curr_image = img_im.convert("RGB")
            image_original_size = self.curr_image.size

        image_resized = False
        mask_resized = False

        if ref_size:
            if self.curr_image and self.curr_image.size != ref_size:
                self.curr_image = self.curr_image.resize(ref_size, Image.BILINEAR)
                image_resized = True
            if self.curr_mask and self.curr_mask.size != ref_size:
                mask_bin01 = (np.array(self.curr_mask) > 0).astype(np.uint8)
                mask_bin01_resized = self.resize_bin01_nearest(mask_bin01, ref_size)
                self.curr_mask = Image.fromarray((mask_bin01_resized * 255).astype(np.uint8))
                mask_resized = True

        # Overlays
        self.curr_overlay = self.make_overlay(self.curr_image, self.curr_mask) if self.curr_image and self.curr_mask else None
        self.curr_pred_overlay = self.make_overlay(self.curr_image, self.curr_pred) if self.curr_image and self.curr_pred else None

        # IoU / Dice (threshold BEFORE resize -> to numpy -> binary 0/1 -> back to PIL -> resize NEAREST)
        if self.curr_mask is not None and self.curr_pred is not None:
            pred_bin = (np.array(self.curr_pred) > 0).astype(np.uint8)
            mask_bin = (np.array(self.curr_mask) > 0).astype(np.uint8)
            if mask_bin.shape[:2] != pred_bin.shape[:2]:
                mask_bin = self.resize_bin01_nearest(mask_bin, self.curr_pred.size).astype(np.uint8)

            inter = np.logical_and(mask_bin, pred_bin).sum()
            union = np.logical_or(mask_bin, pred_bin).sum()
            iou = inter / union if union > 0 else 0.0
            self.iou_info = f"IoU: {iou:.4f}"
            item["iou"] = float(iou)

            denom = mask_bin.sum() + pred_bin.sum()
            dice = (2.0 * inter) / denom if denom > 0 else 0.0
            self.dice_info = f"Dice: {dice:.4f}"
            item["dice"] = float(dice)
        else:
            self.iou_info = "IoU: NaN"
            item["iou"] = None
            self.dice_info = "Dice: NaN"
            item["dice"] = None

        # Info section
        info_lines = []
        if self.curr_image:
            img_size_kb = os.path.getsize(item["image_path"]) / 1024 if item.get("image_path") else 0.0
            img_orig = f"{image_original_size[0]}x{image_original_size[1]}" if image_original_size else "NaN"
            img_loaded = f"{self.curr_image.width}x{self.curr_image.height}"
            img_resize_note = "not resized"
            if image_resized:
                img_resize_note = f"resized to {ref_source} {ref_size[0]}x{ref_size[1]}" if ref_source and ref_size else "resized"
            info_lines.extend([
                f"Image: {os.path.basename(item['image_path']) if item.get('image_path') else item['filename']}",
                f" - Original: {img_orig}",
                f" - Loaded: {img_loaded} ({img_resize_note})",
                f" - Size: {img_size_kb:.1f} KB"
            ])
        if self.curr_mask:
            mask_size_kb = os.path.getsize(item["mask_path"]) / 1024 if item.get("mask_path") else 0.0
            mask_vals = np.unique(np.array(self.curr_mask))
            mask_orig = f"{mask_original_size[0]}x{mask_original_size[1]}" if mask_original_size else "NaN"
            mask_loaded = f"{self.curr_mask.width}x{self.curr_mask.height}"
            mask_resize_note = "not resized"
            if mask_resized:
                mask_resize_note = f"resized to {ref_source} {ref_size[0]}x{ref_size[1]}" if ref_source and ref_size else "resized"
            info_lines.extend([
                "",
                f"Mask: {os.path.basename(item['mask_path']) if item.get('mask_path') else item['filename']}",
                f" - Original: {mask_orig}",
                f" - Loaded: {mask_loaded} ({mask_resize_note})",
                f" - Size: {mask_size_kb:.1f} KB",
                f" - Unique values ({len(mask_vals)}):"
            ])
            for val in mask_vals:
                rgb, cname = self.get_color(val)
                info_lines.append(f"   {val}: {cname} RGB{rgb}")
        if self.curr_pred:
            pred_size_kb = os.path.getsize(item["pred_path"]) / 1024 if item.get("pred_path") else 0.0
            pred_vals = np.unique(np.array(self.curr_pred))
            info_lines.extend([
                "",
                f"Predicted Mask: {os.path.basename(item['pred_path']) if item.get('pred_path') else item['filename']}",
                f" - Resolution: {self.curr_pred.width}x{self.curr_pred.height} (not resized)",
                f" - Size: {pred_size_kb:.1f} KB",
                f" - Unique values ({len(pred_vals)}):"
            ])
            for val in pred_vals:
                rgb, cname = self.get_color(val)
                info_lines.append(f"   {val}: {cname} RGB{rgb}")

        self.info_text_base = "\n".join(info_lines)
        self.pixel_info = "Pointer: NaN"
        self.update_info_panel()

        # Keep the right panel's IoU values in sync (without re-rendering everything)
        if item["id"] in self._id_to_index:
            self.file_tree.item(
                item["id"],
                values=(item["filename"], self.format_iou(item.get("iou")), self.format_dice(item.get("dice"))),
            )

    def get_color(self, label_val):
        if label_val not in self.color_map:
            random.seed(int(label_val))
            r = random.randint(50, 255)
            g = random.randint(50, 255)
            b = random.randint(50, 255)
            self.color_map[label_val] = (r, g, b)
            self.color_names[label_val] = f"Class{label_val}"
        return self.color_map[label_val], self.color_names[label_val]

    def make_overlay(self, img_pil, mask_pil):
        if mask_pil is None or img_pil is None:
            return None
        img_np = np.array(img_pil).astype(np.float32)
        mask_np = np.array(mask_pil)

        overlay = img_np.copy()
        unique_vals = np.unique(mask_np)

        for val in unique_vals:
            if val == 0:
                continue
            (r, g, b), _ = self.get_color(val)
            region = mask_np == val
            overlay[region] = 0.8 * overlay[region] + 0.2 * np.array([r, g, b])

        edges = np.zeros_like(mask_np, dtype=bool)
        edges[:-1, :] |= mask_np[:-1, :] != mask_np[1:, :]
        edges[:, :-1] |= mask_np[:, :-1] != mask_np[:, 1:]
        overlay[edges] = [20, 20, 20]

        return Image.fromarray(overlay.astype(np.uint8))

    def start_pan(self, event):
        self.dragging = True
        self.start_x, self.start_y = event.x, event.y

    def do_pan(self, event):
        if self.dragging:
            dx, dy = event.x - self.start_x, event.y - self.start_y
            self.pan_x += dx
            self.pan_y += dy
            self.start_x, self.start_y = event.x, event.y
            self.refresh_display()

    def do_zoom(self, event):
        if event.num == 4 or event.delta > 0:
            self.zoom *= 1.1
        elif event.num == 5 or event.delta < 0:
            self.zoom /= 1.1
        self.refresh_display()

    def draw_on_canvas(self, canvas, img, which):
        if img is None:
            canvas.delete("all")
            return
        cw, ch = max(canvas.winfo_width(), 1), max(canvas.winfo_height(), 1)

        if self.base_scale == 1.0:
            self.base_scale = min(cw / img.width, ch / img.height)

        total_scale = self.base_scale * self.zoom
        new_w, new_h = int(img.width * total_scale), int(img.height * total_scale)
        if new_w < 2 or new_h < 2:
            return

        img_resized = img.resize((new_w, new_h), Image.NEAREST)
        tk_img = ImageTk.PhotoImage(img_resized)
        canvas.delete("all")
        x, y = cw // 2 + self.pan_x, ch // 2 + self.pan_y
        canvas.create_image(x, y, image=tk_img, anchor="center")

        if which == "image":
            self.tk_img_image = tk_img
        elif which == "mask":
            self.tk_img_mask = tk_img
        elif which == "overlay":
            self.tk_img_overlay = tk_img
        elif which == "pred":
            self.tk_img_pred = tk_img
        elif which == "pred_overlay":
            self.tk_img_pred_overlay = tk_img

    def show_pixel_info(self, event):
        if self.curr_image is None:
            return
        cw, ch = event.widget.winfo_width(), event.widget.winfo_height()
        total_scale = self.base_scale * self.zoom
        img_w, img_h = self.curr_image.size
        x_img = int((event.x - cw // 2 - self.pan_x) / total_scale + img_w // 2)
        y_img = int((event.y - ch // 2 - self.pan_y) / total_scale + img_h // 2)
        if 0 <= x_img < img_w and 0 <= y_img < img_h:
            rgb = np.array(self.curr_image)[y_img, x_img].tolist()
            if self.curr_mask is not None:
                mval = int(np.array(self.curr_mask)[y_img, x_img])
                (r, g, b), cname = self.get_color(mval)
                self.pixel_info = f"Pointer: ({x_img},{y_img}) Image RGB={rgb}, Mask={mval} ({cname} RGB({r},{g},{b}))"
            else:
                self.pixel_info = f"Pointer: ({x_img},{y_img}) Image RGB={rgb}, Mask=NaN"
        else:
            self.pixel_info = "Pointer: NaN"
        self.update_info_panel()

    def update_info_panel(self):
        self.info_text.config(state="normal")
        self.info_text.delete("1.0", tk.END)
        self.info_text.insert(
            tk.END,
            self.pixel_info + "\n" + self.iou_info + "\n" + self.dice_info + "\n\n" + self.info_text_base,
        )
        self.info_text.config(state="disabled")

    def refresh_display(self, event=None):
        self.draw_on_canvas(self.image_canvas, self.curr_image, "image")
        self.draw_on_canvas(self.mask_canvas, self.curr_mask, "mask")
        self.draw_on_canvas(self.overlay_canvas, self.curr_overlay, "overlay")
        self.draw_on_canvas(self.pred_canvas, self.curr_pred, "pred")
        self.draw_on_canvas(self.pred_overlay_canvas, self.curr_pred_overlay, "pred_overlay")

    def next_image(self, event=None):
        if self.index < len(self.items) - 1:
            self.index += 1
            self.reset_and_load()

    def prev_image(self, event=None):
        if self.index > 0:
            self.index -= 1
            self.reset_and_load()


def main():
    root = tk.Tk()
    app = ImageMaskViewer(root)
    root.geometry("1600x1200")
    root.mainloop()


if __name__ == "__main__":
    main()
