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

# --- Mask helpers (binary vs multi-class) ---
def load_label_map_from_pil(im: Image.Image) -> np.ndarray:
    """
    Load a 2D label map from a PIL image without collapsing multi-class information.

    - Palette ("P") images return palette indices (class ids).
    - Grayscale ("L", 16-bit) images return pixel values as class ids.
    - RGB/RGBA images are packed into a single int per pixel: (R<<16 | G<<8 | B).
    """
    if im.mode in ("RGB", "RGBA"):
        rgb = np.array(im.convert("RGB"), dtype=np.uint32)
        return ((rgb[..., 0] << 16) | (rgb[..., 1] << 8) | rgb[..., 2]).astype(np.uint32)

    # Keep palette indices for "P". For most other modes, a direct array is fine.
    arr = np.array(im)
    if arr.ndim == 3:
        # Defensive: if we still got channels, fall back to first channel
        arr = arr[..., 0]
    return arr


def _unique_2d(arr: np.ndarray) -> np.ndarray:
    arr = np.asarray(arr)
    if arr.ndim != 2:
        raise ValueError("expected 2D array")
    return np.unique(arr)


def is_multiclass_gt_mask(label_map: np.ndarray) -> bool:
    """
    Ground-truth mask multi-class detector (safe for binary).

    If the GT mask has >=2 distinct non-zero labels and a small discrete label set,
    treat it as multi-class.
    """
    if label_map is None:
        return False
    arr = np.asarray(label_map)
    if arr.ndim != 2:
        return False
    uniq = _unique_2d(arr)
    if uniq.size <= 2:
        return False
    uniq_fg = uniq[uniq != 0]
    if uniq_fg.size < 2:
        return False
    if uniq.size > 64:
        return False

    try:
        fg_max = int(np.max(uniq_fg))
    except Exception:
        return False

    # Typical multi-class GT masks use small class ids (0..K-1).
    if fg_max <= 64:
        return True

    # RGB-packed GT masks produce large discrete ids.
    if fg_max > 255:
        return True

    # Values in 65..255 are ambiguous (often intensity artifacts or binary encodings like {0,1,255}).
    # Prefer binary behavior to avoid breaking previously-working binary masks.
    return False


def is_multiclass_pred_mask(label_map: np.ndarray) -> bool:
    """
    Prediction mask multi-class detector (conservative).

    Avoid treating grayscale/probability-like predictions as multi-class.
    """
    if label_map is None:
        return False
    arr = np.asarray(label_map)
    if arr.ndim != 2:
        return False
    uniq = _unique_2d(arr)
    if uniq.size <= 2:
        return False
    uniq_fg = uniq[uniq != 0]
    if uniq_fg.size < 2:
        return False
    if uniq.size > 64:
        return False

    try:
        fg_max = int(np.max(uniq_fg))
    except Exception:
        return False

    # Typical argmax label maps use small class ids.
    if fg_max <= 64:
        return True

    # RGB-packed ids are large but still discrete.
    if fg_max > 255:
        return True

    return False


def scale_to_uint8_for_display(arr: np.ndarray) -> np.ndarray:
    arr = np.asarray(arr)
    if arr.size == 0:
        return np.zeros((0, 0), dtype=np.uint8)
    mx = float(arr.max()) if arr.max() is not None else 0.0
    if mx <= 0:
        return np.zeros(arr.shape, dtype=np.uint8)
    if mx == 255.0 and arr.dtype == np.uint8:
        return arr
    out = (arr.astype(np.float32) / mx) * 255.0
    return out.clip(0, 255).astype(np.uint8)


def resize_labels_nearest(label_map: np.ndarray, target_size, cv2_mod=None) -> np.ndarray:
    tw, th = int(target_size[0]), int(target_size[1])
    if label_map is None:
        return None
    arr = np.asarray(label_map)
    if arr.ndim != 2:
        raise ValueError("label_map must be 2D")
    if cv2_mod is not None:
        if arr.dtype == np.uint32:
            arr = arr.astype(np.int32)
        return cv2_mod.resize(arr, (tw, th), interpolation=cv2_mod.INTER_NEAREST)
    if arr.dtype == np.uint32:
        arr = arr.astype(np.int32)
    pil = Image.fromarray(arr)
    return np.array(pil.resize((tw, th), Image.NEAREST))


def resize_gray_nearest(gray_map: np.ndarray, target_size, cv2_mod=None) -> np.ndarray:
    tw, th = int(target_size[0]), int(target_size[1])
    if gray_map is None:
        return None
    arr = np.asarray(gray_map)
    if arr.ndim != 2:
        raise ValueError("gray_map must be 2D")
    if cv2_mod is not None:
        return cv2_mod.resize(arr, (tw, th), interpolation=cv2_mod.INTER_NEAREST)
    pil = Image.fromarray(arr)
    return np.array(pil.resize((tw, th), Image.NEAREST))


def multiclass_mean_iou_dice(mask_labels: np.ndarray, pred_labels: np.ndarray, ignore_label=0):
    if mask_labels is None or pred_labels is None:
        return None, None
    if mask_labels.shape != pred_labels.shape:
        raise ValueError("mask_labels and pred_labels must have the same shape")

    classes = sorted(set(np.unique(mask_labels)).union(set(np.unique(pred_labels))))
    classes = [c for c in classes if c != ignore_label]
    if not classes:
        return 0.0, 0.0

    ious = []
    dices = []
    for c in classes:
        m = mask_labels == c
        p = pred_labels == c
        inter = np.logical_and(m, p).sum()
        union = np.logical_or(m, p).sum()
        iou = (inter / union) if union > 0 else 0.0
        denom = m.sum() + p.sum()
        dice = (2.0 * inter / denom) if denom > 0 else 0.0
        ious.append(float(iou))
        dices.append(float(dice))

    return float(np.mean(ious)), float(np.mean(dices))

# --- 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

        # Cached raw label maps (used for multi-class metrics/overlay)
        self.curr_is_multiclass = False
        self.curr_mask_labels = None
        self.curr_pred_labels = 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)
        self.file_action_menu.add_command(label="Save displayed PNGs…", command=self.save_clicked_row_displayed_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="Close and Save Path", command=self.close_and_save_path)
        menubar.add_command(label="Reload", command=self.reload_all)
        root.config(menu=menubar)
        root.protocol("WM_DELETE_WINDOW", self.on_app_close)

        # 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_labels = load_label_map_from_pil(pred_im)
                pred_np_gray = np.array(pred_im.convert("L"))
            pred_size = (pred_np_gray.shape[1], pred_np_gray.shape[0])

            with Image.open(mask_path) as mask_im:
                mask_labels = load_label_map_from_pil(mask_im)

            is_multi = is_multiclass_gt_mask(mask_labels) if mask_labels is not None else is_multiclass_pred_mask(pred_labels)

            if is_multi:
                if mask_labels.shape[:2] != pred_labels.shape[:2]:
                    mask_labels = resize_labels_nearest(mask_labels, pred_size, cv2)
                if pred_labels.shape[:2] != mask_labels.shape[:2]:
                    pred_labels = resize_labels_nearest(pred_labels, (mask_labels.shape[1], mask_labels.shape[0]), cv2)
                return multiclass_mean_iou_dice(mask_labels, pred_labels, ignore_label=0)

            # Binary (existing behavior): any non-zero value counts as foreground
            with Image.open(pred_path) as pred_im2:
                pred = pred_im2.convert("L")
            pred_bin = (np.array(pred) > 0).astype(np.uint8)
            target_size = (pred_bin.shape[1], pred_bin.shape[0])

            with Image.open(mask_path) as mask_im2:
                mask = mask_im2.convert("L")
            mask_bin = (np.array(mask) > 0).astype(np.uint8)
            if mask_bin.shape[:2] != pred_bin.shape[:2]:
                # Keep legacy binary behavior: threshold before resize.
                mask_bin = self.resize_bin01_nearest(mask_bin, target_size).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 _render_display_like_canvas(self, canvas, img_pil, base_scale, zoom, pan_x, pan_y):
        if img_pil is None:
            return None

        cw, ch = max(int(canvas.winfo_width()), 1), max(int(canvas.winfo_height()), 1)
        total_scale = float(base_scale) * float(zoom)
        new_w, new_h = int(img_pil.width * total_scale), int(img_pil.height * total_scale)

        mode = img_pil.mode if img_pil.mode in ("L", "RGB") else "RGB"
        bg = Image.new(mode, (cw, ch), 0 if mode == "L" else (0, 0, 0))
        if new_w < 2 or new_h < 2:
            return bg

        img_resized = img_pil.resize((new_w, new_h), Image.NEAREST)

        x_center = cw // 2 + int(pan_x)
        y_center = ch // 2 + int(pan_y)
        left = x_center - new_w // 2
        top = y_center - new_h // 2

        src_left = max(0, -left)
        src_top = max(0, -top)
        dst_left = max(0, left)
        dst_top = max(0, top)
        width = min(new_w - src_left, cw - dst_left)
        height = min(new_h - src_top, ch - dst_top)
        if width <= 0 or height <= 0:
            return bg

        crop = img_resized.crop((src_left, src_top, src_left + width, src_top + height))
        bg.paste(crop, (dst_left, dst_top))
        return bg

    def _build_display_bundle_for_item(self, item):
        pred_original_size = None
        mask_original_size = None

        pred_pil = None
        pred_labels = None
        if item.get("pred_path") and os.path.exists(item["pred_path"]):
            with Image.open(item["pred_path"]) as pred_im:
                pred_labels = load_label_map_from_pil(pred_im)
                pred_gray = np.array(pred_im.convert("L"))

            pred_is_multi = is_multiclass_pred_mask(pred_labels)
            if pred_is_multi:
                pred_disp = scale_to_uint8_for_display(pred_labels)
            else:
                pred_disp = pred_gray
                if pred_disp.max() > 0 and pred_disp.max() != 255:
                    pred_disp = (pred_disp.astype(np.float32) / pred_disp.max()) * 255
                    pred_disp = pred_disp.astype(np.uint8)

            pred_pil = Image.fromarray(pred_disp.astype(np.uint8))
            pred_original_size = pred_pil.size

        mask_pil = None
        mask_labels = None
        if item.get("mask_path") and os.path.exists(item["mask_path"]):
            with Image.open(item["mask_path"]) as mask_im:
                mask_labels = load_label_map_from_pil(mask_im)
                mask_gray = np.array(mask_im.convert("L"))

            mask_is_multi = is_multiclass_gt_mask(mask_labels)
            if mask_is_multi:
                mask_disp = scale_to_uint8_for_display(mask_labels)
            else:
                mask_disp = mask_gray
                if mask_disp.max() > 0 and mask_disp.max() != 255:
                    mask_disp = (mask_disp.astype(np.float32) / mask_disp.max()) * 255
                    mask_disp = mask_disp.astype(np.uint8)

            mask_pil = Image.fromarray(mask_disp.astype(np.uint8))
            mask_original_size = mask_pil.size

        is_multi = is_multiclass_gt_mask(mask_labels) if mask_labels is not None else is_multiclass_pred_mask(pred_labels)
        if is_multi and pred_labels is not None:
            pred_disp = scale_to_uint8_for_display(pred_labels)
            pred_pil = Image.fromarray(pred_disp.astype(np.uint8))

        ref_size = pred_original_size or mask_original_size

        image_pil = None
        if item.get("image_path") and os.path.exists(item["image_path"]):
            with Image.open(item["image_path"]) as img_im:
                image_pil = img_im.convert("RGB")
            if ref_size and image_pil.size != ref_size:
                image_pil = image_pil.resize(ref_size, Image.BILINEAR)

        if ref_size and mask_pil is not None and mask_pil.size != ref_size:
            if is_multi and mask_labels is not None:
                mask_labels = resize_labels_nearest(mask_labels, ref_size, cv2)
                mask_disp = scale_to_uint8_for_display(mask_labels)
                mask_pil = Image.fromarray(mask_disp.astype(np.uint8))
            else:
                mask_bin01 = (np.array(mask_pil) > 0).astype(np.uint8)
                mask_bin01_resized = self.resize_bin01_nearest(mask_bin01, ref_size)
                mask_pil = Image.fromarray((mask_bin01_resized * 255).astype(np.uint8))

        if is_multi and image_pil is not None:
            mask_for_overlay = Image.fromarray(np.asarray(mask_labels).astype(np.int32)) if mask_labels is not None else None
            pred_for_overlay = Image.fromarray(np.asarray(pred_labels).astype(np.int32)) if pred_labels is not None else None
            overlay_pil = self.make_overlay(image_pil, mask_for_overlay) if mask_for_overlay is not None else None
            pred_overlay_pil = self.make_overlay(image_pil, pred_for_overlay) if pred_for_overlay is not None else None
        else:
            overlay_pil = self.make_overlay(image_pil, mask_pil) if image_pil is not None and mask_pil is not None else None
            pred_overlay_pil = self.make_overlay(image_pil, pred_pil) if image_pil is not None and pred_pil is not None else None

        return {
            "image": image_pil,
            "mask": mask_pil,
            "overlay": overlay_pil,
            "pred": pred_pil,
            "pred_overlay": pred_overlay_pil,
        }

    def save_clicked_row_displayed_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 displayed as (base name)",
            defaultextension=".png",
            filetypes=[("PNG", "*.png")],
            initialfile=f"{it['id']}_display.png",
        )
        if not base_path:
            return

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

        # Save "displayed" content without canvas padding.
        # Everything is saved at prediction-size (if available), otherwise mask-size.
        current_it = self.items[self.index] if (self.items and 0 <= self.index < len(self.items)) else None
        if current_it and current_it.get("id") == it.get("id"):
            bundle = {
                "image": self.curr_image,
                "mask": self.curr_mask,
                "overlay": self.curr_overlay,
                "pred": self.curr_pred,
                "pred_overlay": self.curr_pred_overlay,
            }
        else:
            # Not currently displayed: save how it would display with default view.
            bundle = self._build_display_bundle_for_item(it)

        outputs = [
            ("image", os.path.join(out_dir, f"{base_name}_display_image.png")),
            ("mask", os.path.join(out_dir, f"{base_name}_display_mask.png")),
            ("overlay", os.path.join(out_dir, f"{base_name}_display_mask_overlay.png")),
            ("pred", os.path.join(out_dir, f"{base_name}_display_pred.png")),
            ("pred_overlay", os.path.join(out_dir, f"{base_name}_display_pred_overlay.png")),
        ]

        ref_img = bundle.get("pred") or bundle.get("mask") or bundle.get("image") or bundle.get("overlay") or bundle.get("pred_overlay")
        ref_size = ref_img.size if ref_img is not None else None

        def ensure_size(img, size):
            if img is None or size is None:
                return img
            if img.size == size:
                return img
            return img.resize(size, Image.NEAREST)

        saved_any = False
        errors = []

        for key, out_path in outputs:
            try:
                img = bundle.get(key)
                if img is None:
                    continue
                img = ensure_size(img, ref_size)
                img.save(out_path, format="PNG")
                saved_any = True
            except Exception as e:
                errors.append(f"{key}: {e}")

        if errors:
            messagebox.showerror("Save displayed PNGs", "\n".join(errors), parent=self.root)
        elif saved_any:
            messagebox.showinfo("Save displayed PNGs", "Saved displayed PNGs successfully.", parent=self.root)
        else:
            messagebox.showinfo("Save displayed PNGs", "Nothing to save (missing images).", 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 on_app_close(self):
        save_dirs("", "", "", self.resize_mode, self.resize_width, self.resize_height)
        self.root.destroy()

    def close_and_save_path(self):
        save_dirs(
            self.image_dir,
            self.mask_dir,
            self.pred_dir,
            self.resize_mode,
            self.resize_width,
            self.resize_height,
        )
        self.root.destroy()

    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_is_multiclass = False
        self.curr_mask_labels = None
        self.curr_pred_labels = None
        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:
                self.curr_pred_labels = load_label_map_from_pil(pred_im)
                pred_gray = np.array(pred_im.convert("L"))

            # For multi-class, display from labels to avoid losing palette/RGB class ids.
            pred_is_multi = is_multiclass_pred_mask(self.curr_pred_labels)
            if pred_is_multi:
                pred_disp = scale_to_uint8_for_display(self.curr_pred_labels)
            else:
                pred_disp = pred_gray
                if pred_disp.max() > 0 and pred_disp.max() != 255:
                    pred_disp = (pred_disp.astype(np.float32) / pred_disp.max()) * 255
                    pred_disp = pred_disp.astype(np.uint8)

            self.curr_pred = Image.fromarray(pred_disp.astype(np.uint8))
            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:
                self.curr_mask_labels = load_label_map_from_pil(mask_im)
                mask_gray = np.array(mask_im.convert("L"))

            mask_is_multi = is_multiclass_gt_mask(self.curr_mask_labels)

            if mask_is_multi:
                mask_disp = scale_to_uint8_for_display(self.curr_mask_labels)
            else:
                mask_disp = mask_gray
                if mask_disp.max() > 0 and mask_disp.max() != 255:
                    mask_disp = (mask_disp.astype(np.float32) / mask_disp.max()) * 255
                    mask_disp = mask_disp.astype(np.uint8)

            self.curr_mask = Image.fromarray(mask_disp.astype(np.uint8))
            mask_original_size = self.curr_mask.size

        # Decide multi-class primarily from the ground-truth mask to avoid breaking
        # binary workflows where predictions may be grayscale/probability-like.
        if self.curr_mask_labels is not None:
            self.curr_is_multiclass = is_multiclass_gt_mask(self.curr_mask_labels)
        else:
            self.curr_is_multiclass = is_multiclass_pred_mask(self.curr_pred_labels)

        # If it's a multi-class task, ensure prediction visualization is driven by labels
        # (avoids losing palette/RGB class ids via grayscale conversion).
        if self.curr_is_multiclass and self.curr_pred_labels is not None:
            pred_disp = scale_to_uint8_for_display(self.curr_pred_labels)
            self.curr_pred = Image.fromarray(pred_disp.astype(np.uint8))

        # 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:
                if self.curr_is_multiclass and self.curr_mask_labels is not None:
                    # Multi-class: preserve class ids while resizing (nearest).
                    self.curr_mask_labels = resize_labels_nearest(self.curr_mask_labels, ref_size, cv2)
                    mask_disp = scale_to_uint8_for_display(self.curr_mask_labels)
                    self.curr_mask = Image.fromarray(mask_disp.astype(np.uint8))
                else:
                    # Binary: keep legacy behavior (threshold before resize) to match previous-app.py.
                    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
        if self.curr_is_multiclass and self.curr_image is not None:
            mask_for_overlay = Image.fromarray(np.asarray(self.curr_mask_labels).astype(np.int32)) if self.curr_mask_labels is not None else None
            pred_for_overlay = Image.fromarray(np.asarray(self.curr_pred_labels).astype(np.int32)) if self.curr_pred_labels is not None else None
            self.curr_overlay = self.make_overlay(self.curr_image, mask_for_overlay) if mask_for_overlay is not None else None
            self.curr_pred_overlay = self.make_overlay(self.curr_image, pred_for_overlay) if pred_for_overlay is not None else None
        else:
            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
        if self.curr_mask is not None and self.curr_pred is not None:
            if self.curr_is_multiclass and self.curr_mask_labels is not None and self.curr_pred_labels is not None:
                # Multi-class: mean IoU/Dice over foreground classes (labels != 0)
                if self.curr_mask_labels.shape[:2] != self.curr_pred_labels.shape[:2]:
                    self.curr_mask_labels = resize_labels_nearest(self.curr_mask_labels, self.curr_pred.size, cv2)
                iou, dice = multiclass_mean_iou_dice(self.curr_mask_labels, self.curr_pred_labels, ignore_label=0)
                self.iou_info = f"IoU: {iou:.4f}"
                item["iou"] = float(iou)
                self.dice_info = f"Dice: {dice:.4f}"
                item["dice"] = float(dice)
            else:
                # Binary: keep legacy behavior (resize binary mask if needed)
                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.asarray(self.curr_mask_labels)) if self.curr_is_multiclass and self.curr_mask_labels is not None else 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.asarray(self.curr_pred_labels)) if self.curr_is_multiclass and self.curr_pred_labels is not None else 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:
                if self.curr_is_multiclass and self.curr_mask_labels is not None:
                    mval = int(np.asarray(self.curr_mask_labels)[y_img, x_img])
                else:
                    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()
