import gradio as gr
from gradio_bbox_annotator import BBoxAnnotator
from PIL import Image
import numpy as np
import torch
import os
import shutil
import time
import json
import uuid
from pathlib import Path
import tempfile
import zipfile
import html
from skimage import measure
from matplotlib import cm
from glob import glob
from natsort import natsorted
from huggingface_hub import HfApi, upload_file
import spaces

from inference_seg import load_model as load_seg_model, run as run_seg
from inference_count import load_model as load_count_model, run as run_count
from inference_track import load_model as load_track_model, run as run_track
from _utils.image_io import (
    standardize_image, inspect_image, image_size, array_nbytes, pixel_stats, bit_depth,
    RECOMMENDED_SIZE, WARN_SIZE, MAX_SIZE, MIN_SIZE, MAX_READ_BYTES, TIFF_EXTENSIONS,
)

HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_REPO = "VisionLanguageGroup/feedback"


print("===== clearing cache =====")
cache_path = os.path.expanduser("~/.cache/huggingface/gradio")
if os.path.exists(cache_path):
    try:
        shutil.rmtree(cache_path)
        print("✅ Deleted ~/.cache/huggingface/gradio")
    except:
        pass

SEG_MODEL = None
SEG_DEVICE = torch.device("cpu")

COUNT_MODEL = None
COUNT_DEVICE = torch.device("cpu")

TRACK_MODEL = None
TRACK_DEVICE = torch.device("cpu")

def load_all_models():
    global SEG_MODEL, SEG_DEVICE
    global COUNT_MODEL, COUNT_DEVICE
    global TRACK_MODEL, TRACK_DEVICE
    
    print("\n" + "="*60)
    print("📦 Loading Segmentation Model")
    print("="*60)
    SEG_MODEL, SEG_DEVICE = load_seg_model(use_box=False)
    
    print("\n" + "="*60)
    print("📦 Loading Counting Model")
    print("="*60)
    COUNT_MODEL, COUNT_DEVICE = load_count_model(use_box=False)
    
    print("\n" + "="*60)
    print("📦 Loading Tracking Model")
    print("="*60)
    TRACK_MODEL, TRACK_DEVICE = load_track_model(use_box=False)
    
    print("\n" + "="*60)
    print("✅ All Models Loaded Successfully")
    print("="*60)

load_all_models()

DATASET_DIR = Path("solver_cache")
DATASET_DIR.mkdir(parents=True, exist_ok=True)

def save_feedback_to_hf(query_id, feedback_type, feedback_text=None, img_path=None, bboxes=None):
    """Save feedback to Hugging Face Dataset"""
    
    if not HF_TOKEN:
        print("⚠️ No HF_TOKEN found, using local storage")
        save_feedback(query_id, feedback_type, feedback_text, img_path, bboxes)
        return
    
    # Shared by the record and its image so the two pair up; query_id alone is not unique (a session can submit more than once).
    stem = f"feedback_{query_id}_{int(time.time())}"

    try:
        api = HfApi()

        image_in_repo = None
        if img_path and os.path.exists(img_path):
            try:
                image_in_repo = f"images/{stem}.png"
                api.upload_file(
                    path_or_fileobj=img_path,
                    path_in_repo=image_in_repo,
                    repo_id=DATASET_REPO,
                    repo_type="dataset",
                    token=HF_TOKEN
                )
            except Exception as e:
                print(f"⚠️ Failed to upload image: {e}")
                image_in_repo = None  

        feedback_data = {
            "query_id": query_id,
            "feedback_type": feedback_type,
            "feedback_text": feedback_text,
            "image_path": image_in_repo, 
            "bboxes": str(bboxes), 
            "datetime": time.strftime("%Y-%m-%d %H:%M:%S"),
            "timestamp": time.time()
        }

        filename = f"{stem}.json"

        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(feedback_data, f, indent=2, ensure_ascii=False)

        api.upload_file(
            path_or_fileobj=filename,
            path_in_repo=f"data/{filename}",
            repo_id=DATASET_REPO,
            repo_type="dataset",
            token=HF_TOKEN
        )

        os.remove(filename)

        print(f"✅ Feedback saved to HF Dataset: {DATASET_REPO} ({stem})")

    except Exception as e:
        print(f"⚠️ Failed to save to HF Dataset: {e}")
        save_feedback(query_id, feedback_type, feedback_text, img_path, bboxes)


def save_feedback(query_id, feedback_type, feedback_text=None, img_path=None, bboxes=None):
    """Save feedback to local JSON file"""
    feedback_data = {
        "query_id": query_id,
        "feedback_type": feedback_type,
        "feedback_text": feedback_text,
        "image": img_path,
        "bboxes": bboxes,
        "datetime": time.strftime("%Y%m%d_%H%M%S")
    }
    feedback_file = DATASET_DIR / query_id / "feedback.json"
    feedback_file.parent.mkdir(parents=True, exist_ok=True)
    
    if feedback_file.exists():
        with feedback_file.open("r") as f:
            existing = json.load(f)
            if not isinstance(existing, list):
                existing = [existing]
            existing.append(feedback_data)
            feedback_data = existing
    else:
        feedback_data = [feedback_data]
    
    with feedback_file.open("w") as f:
        json.dump(feedback_data, f, indent=4, ensure_ascii=False)

def parse_first_bbox(bboxes):
    """Parse the first bounding box from the annotation input, supports dict or list format"""
    if not bboxes:
        return None
    b = bboxes[0]
    if isinstance(b, dict):
        x, y = float(b.get("x", 0)), float(b.get("y", 0))
        w, h = float(b.get("width", 0)), float(b.get("height", 0))
        return x, y, x + w, y + h
    if isinstance(b, (list, tuple)) and len(b) >= 4:
        return float(b[0]), float(b[1]), float(b[2]), float(b[3])
    return None

def parse_bboxes(bboxes):
    """Parse all bounding boxes from the annotation input"""
    if not bboxes:
        return None
    
    result = []
    for b in bboxes:
        if isinstance(b, dict):
            x, y = float(b.get("x", 0)), float(b.get("y", 0))
            w, h = float(b.get("width", 0)), float(b.get("height", 0))
            result.append([x, y, x + w, y + h])
        elif isinstance(b, (list, tuple)) and len(b) >= 4:
            result.append([float(b[0]), float(b[1]), float(b[2]), float(b[3])])
    
    return result

def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
    """Convert a 2D mask of instance IDs to a color image for visualization."""
    def hsv_to_rgb(h, s, v):
        i = int(h * 6.0)
        f = h * 6.0 - i
        i = i % 6
        p = v * (1 - s)
        q = v * (1 - f * s)
        t = v * (1 - (1 - f) * s)
        if i == 0: r, g, b = v, t, p
        elif i == 1: r, g, b = q, v, p
        elif i == 2: r, g, b = p, v, t
        elif i == 3: r, g, b = p, q, v
        elif i == 4: r, g, b = t, p, v
        else: r, g, b = v, p, q
        return int(r * 255), int(g * 255), int(b * 255)

    palette = [(0, 0, 0)]
    for i in range(1, num_colors):
        h = (i % num_colors) / float(num_colors)
        palette.append(hsv_to_rgb(h, 1.0, 0.95))

    palette_arr = np.array(palette, dtype=np.uint8)
    color_idx = mask % num_colors
    return palette_arr[color_idx]


def render_seg_overlay(img_np, inst_mask, overlay_alpha):
    """Render segmentation overlay from cached image/mask."""
    if img_np is None or inst_mask is None:
        return None

    overlay = img_np.copy()
    alpha = float(np.clip(overlay_alpha, 0.0, 1.0))

    for inst_id in np.unique(inst_mask):
        if inst_id == 0:
            continue
        binary_mask = (inst_mask == inst_id).astype(np.uint8)
        color = get_well_spaced_color(inst_id)
        overlay[binary_mask == 1] = (1 - alpha) * overlay[binary_mask == 1] + alpha * color

        contours = measure.find_contours(binary_mask, 0.5)
        for contour in contours:
            contour = contour.astype(np.int32)
            valid_y = np.clip(contour[:, 0], 0, overlay.shape[0] - 1)
            valid_x = np.clip(contour[:, 1], 0, overlay.shape[1] - 1)
            overlay[valid_y, valid_x] = [1.0, 1.0, 0.0]

    overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
    return Image.fromarray(overlay)


def render_count_overlay(img_np, density_normalized, overlay_alpha):
    """Render counting heatmap overlay from cached image/density."""
    if img_np is None or density_normalized is None:
        return None

    alpha = float(np.clip(overlay_alpha, 0.0, 1.0))
    cmap = cm.get_cmap("jet")
    density_colored = cmap(density_normalized)[:, :, :3]

    overlay = img_np.copy()
    threshold = 0.01
    significant_mask = density_normalized > threshold
    overlay[significant_mask] = (1 - alpha) * overlay[significant_mask] + alpha * density_colored[significant_mask]
    overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
    return Image.fromarray(overlay)


def update_seg_overlay_alpha(overlay_alpha, seg_vis_cache):
    """Live update segmentation visualization without rerunning inference."""
    if not seg_vis_cache:
        return None
    return render_seg_overlay(seg_vis_cache.get("img_np"), seg_vis_cache.get("inst_mask"), overlay_alpha)


def update_count_overlay_alpha(overlay_alpha, count_vis_cache):
    """Live update counting visualization without rerunning inference."""
    if not count_vis_cache:
        return None
    return render_count_overlay(count_vis_cache.get("img_np"), count_vis_cache.get("density_normalized"), overlay_alpha)


def update_tracking_overlay_alpha(overlay_alpha, track_vis_cache):
    """Regenerate tracking visualization at new opacity using cached outputs."""
    if not track_vis_cache:
        return None

    tif_dir = track_vis_cache.get("tif_dir")
    output_dir = track_vis_cache.get("output_dir")
    valid_tif_files = track_vis_cache.get("valid_tif_files")
    if not tif_dir or not output_dir or not valid_tif_files:
        return None

    try:
        return create_tracking_visualization(
            tif_dir=tif_dir,
            output_dir=output_dir,
            valid_tif_files=valid_tif_files,
            overlay_alpha=overlay_alpha
        )
    except Exception as e:
        print(f"⚠️ Failed to update tracking opacity: {e}")
        return None


def cleanup_tracking_cache(track_vis_cache):
    """Delete cached tracking temp directories from the previous run."""
    if not track_vis_cache:
        return
    for key in ["input_temp_dir", "output_dir"]:
        path = track_vis_cache.get(key)
        if path and os.path.isdir(path):
            try:
                shutil.rmtree(path)
            except Exception:
                pass


# Per session. Refused rather than evicting the oldest, since there is no way to
# remove a single entry.
MAX_GALLERY_UPLOADS = 10

_GALLERY_RAW_SOURCE = {}
_RAW_SOURCE_CAP = 500


def _remember_gallery_source(thumb_path, raw_path):
    """Map a gallery thumbnail back to the original file it was made from.
    """
    _GALLERY_RAW_SOURCE[thumb_path] = raw_path
    while len(_GALLERY_RAW_SOURCE) > _RAW_SOURCE_CAP:  # dicts keep insertion order
        _GALLERY_RAW_SOURCE.pop(next(iter(_GALLERY_RAW_SOURCE)))


def _annot_path(annot_value):
    """Extract the image path from a BBoxAnnotator value (path or (path, boxes))."""
    if not annot_value:
        return None
    if isinstance(annot_value, (list, tuple)):
        return annot_value[0] if len(annot_value) > 0 else None
    return annot_value


def _human_bytes(n):
    """Format a byte count as MB or GB, whichever reads better."""
    gb = n / 1024 ** 3
    return f"{gb:.1f} GB" if gb >= 1 else f"{n / 1024 ** 2:.0f} MB"


def check_image(img_path):
    """Validate an uploaded file, cheapest checks first.

    Returns a rejection message if the file cannot be used, otherwise None.
    Reject: empty file; unreadable / unsupported format; dimensions over
    MAX_SIZE; full array over MAX_READ_BYTES (a small time-lapse can still be
    huge); corrupt pixels; non-finite (NaN/inf) pixels. Advisory ``gr.Warning``
    (does not reject): larger than WARN_SIZE; smaller than MIN_SIZE; blank
    (uniform) image; very narrow dynamic range.

    Only the header is touched until the size guards pass; pixel checks decode
    image data afterwards, when doing so is bounded by the memory guard.
    """
    if os.path.exists(img_path) and os.path.getsize(img_path) == 0:
        return "This file is empty (0 bytes). Please upload a valid image file."

    w, h = image_size(img_path)
    if w <= 0 or h <= 0:
        # Neither reader could parse the header, so say so rather than failing
        # silently further down.
        ext = os.path.splitext(img_path)[1].lower() or "(no extension)"
        return (f"Could not read {ext} as an image. Supported formats: "
                f"TIFF / OME-TIFF (8/16/32-bit, stacks, multi-channel), PNG, JPG, and other formats supported by tifffile. "
                f"Please export to TIFF/PNG/JPG (e.g. from ImageJ/Fiji) first.")

    if w > MAX_SIZE or h > MAX_SIZE:
        return (f"Image is {w}×{h}, larger than the {MAX_SIZE}×{MAX_SIZE} limit. "
                f"Please crop or downsample it before uploading.")

    nbytes = array_nbytes(img_path)
    if nbytes > MAX_READ_BYTES:
        info = inspect_image(img_path)
        detail = f"{w}×{h}"
        if info.frames > 1:
            detail += f" × {info.frames} frames"
        if info.channels > 1:
            detail += f" × {info.channels} channels"
        return (f"This file needs {_human_bytes(nbytes)} of memory to open "
                f"({detail}), over the {_human_bytes(MAX_READ_BYTES)} limit. "
                f"Please crop it, or split out the frames you need.")

    if w > WARN_SIZE or h > WARN_SIZE:
        gr.Warning(f"Image is {w}×{h} and will be resized to "
                   f"{RECOMMENDED_SIZE}×{RECOMMENDED_SIZE}, so fine detail could be lost. "
                   f"You can try cropping the image for better results.",
                   duration=None, title="⚠️ Large image")

    if 0 < min(w, h) < MIN_SIZE:
        gr.Warning(f"Image is only {w}×{h} and will be upscaled to "
                   f"{RECOMMENDED_SIZE}×{RECOMMENDED_SIZE}, so results may be unreliable. "
                   f"A larger image may work better.",
                   duration=None, title="⚠️ Very small image")

    # Size guards passed, so pixel decoding is bounded. Check problems that
    # the header cannot reveal.
    rep = pixel_stats(img_path)
    if not rep.decoded:
        return ("This image appears to be corrupted or truncated and could not be read. "
                "Please re-export it (e.g. from ImageJ/Fiji) and try again.")
    if not rep.finite:
        return ("This image contains invalid pixel values (NaN or infinity). "
                "Please clean or re-export it before uploading.")
    if rep.vmin == rep.vmax:
        if rep.vmax == 0:
            kind = "completely black"
        elif rep.dtype_max and rep.vmax >= rep.dtype_max:
            kind = "completely white"
        else:
            kind = "a single uniform value"
        gr.Warning(f"This image is {kind}, it has no visible content, so results will not be meaningful.",
                   duration=None, title="⚠️ Blank image")
    elif rep.low_range:
        gr.Warning("This image may use a very narrow intensity range (low contrast). Signal may be too weak for reliable results; consider adjusting acquisition or contrast before uploading.",
                   duration=None, title="⚠️ Low dynamic range")
    return None


def _describe_read(info):
    """Plain-language summary of how a file's dimensions were interpreted.

    Deliberately avoids array jargon ("axis", "size-4"): the reader thinks in
    channels / z-slices / timepoints, not numpy axes.
    """
    # lines = [f"Detected shape: {info.shape}"]
    if info.frames >1:
        lines = [f"Detected {info.channels}-channel {info.frames}-frame image stack of size {info.width}×{info.height}."]
    else:
        lines = [f"Detected {info.channels}-channel image of size {info.width}×{info.height}."]
    if info.channels >= 3:
        lines.append(f"Loaded as 3-channel RGB image.")
    elif info.channels == 2:
        lines.append(f"Loaded as 2-channel image (red, green).")
    else:
        lines.append(f"Loaded as single-channel grayscale image.")

    if info.frames > 1:
        lines.append(f"Showing first frame by default (you can use the stack frame slider to pick another).")
        if info.sub_frames > 1:
            lines.append(f"Showing the first of {info.sub_frames} {info.sub_label}s by default (use the {info.sub_label} slider to pick another).")

    lines.append(
        "If wrong, please convert your image to a 8-bit RGB/Grayscale image file before uploading (there are available tools such as ImageJ/Fiji)."
    )
    return "<br>".join(lines)


def prepare_uploaded_image(annot_value):
    """On annotator upload: standardize frame 0 for preview and configure the
    stack sliders.

    Browsers cannot render 16-bit / float / multi-page TIFFs, so we standardize
    to an 8-bit RGB PNG. If the file is a stack, notify the user and reveal a
    frame slider (default frame 1); a file with a second stack axis (e.g. Z in a
    time+Z series) also gets a z-plane slider (0 = max-intensity projection).
    Sliders that don't apply stay hidden.

    Returns (annotator_value, raw_path, frame_slider_update, zplane_slider_update).
    """
    hidden = gr.update(visible=False, value=1)
    img_path = _annot_path(annot_value)
    if not img_path:
        return annot_value, None, hidden, hidden

    # Clear the annotator on refusal, so a rejected file cannot reach inference.
    rejected = check_image(img_path)
    if rejected:
        gr.Warning(rejected, duration=None, title="❌ Cannot use this file")
        return None, None, hidden, hidden

    info = inspect_image(img_path)
    display = standardize_image(img_path, frame=0,
                                sub_frame=0 if info.sub_frames > 1 else None)

    # Stay quiet unless something non-obvious happened; a plain RGB or grayscale
    # image needs no explanation.
    if info.guessed or info.frames > 1 or info.channels > 3:
        gr.Info(_describe_read(info), duration=None, title="📚 Image Loading Info")

    slider = (gr.update(visible=True, maximum=info.frames, value=1) if info.frames > 1
              else hidden)
    zslider = (gr.update(visible=True, maximum=info.sub_frames, value=1,
                         label=f"🔬 {info.sub_label.capitalize()} (choose plane to use)")
               if info.sub_frames > 1 else hidden)
    return display, img_path, slider, zslider


def select_frame(raw_path, frame_num, zplane=1):
    """Re-render the annotator preview for the chosen frame and z-plane (both
    1-based; z-plane is ignored for files with only one stack axis)."""
    if not raw_path:
        return gr.update()
    return standardize_image(raw_path, frame=int(frame_num) - 1, sub_frame=int(zplane) - 1)


@spaces.GPU
def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
    """Segmentation handler - supports bounding box, returns colorized overlay and original mask path"""
    if annot_value is None or len(annot_value) < 1:
        print("❌ No annotation input")
        return None, None, {}

    img_path = annot_value[0]
    bboxes = annot_value[1] if len(annot_value) > 1 else []

    print(f"🖼️ Image path: {img_path}")
    # One standardized image for both the model and the overlay (WYSIWYG).
    img_path = standardize_image(img_path)
    print(f"🧪 Standardized image: {img_path}")
    box_array = None
    if use_box_choice == "Yes" and bboxes:
        box = parse_bboxes(bboxes)
        if box:
            box_array = box
            print(f"📦 Using bounding boxes: {box_array}")


    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        mask = run_seg(SEG_MODEL, img_path, box=box_array, device=device)
        print("📏 mask shape:", mask.shape, "dtype:", mask.dtype)
    except Exception as e:
        print(f"❌ Inference failed: {str(e)}")
        return None, None, {}

    temp_mask_file = tempfile.NamedTemporaryFile(delete=False, suffix=".tif")
    mask_img = Image.fromarray(mask.astype(np.uint16))
    mask_img.save(temp_mask_file.name)
    print(f"💾 Original mask saved to: {temp_mask_file.name}")

    try:
        img = Image.open(img_path)
        print("📷 Image mode:", img.mode, "size:", img.size)
    except Exception as e:
        print(f"❌ Failed to open image: {e}")
        return None, None, {}

    try:
        img_rgb = img.convert("RGB").resize(mask.shape[::-1], resample=Image.BILINEAR)
        img_np = np.array(img_rgb, dtype=np.float32)
        if img_np.max() > 1.5:
            img_np = img_np / 255.0
    except Exception as e:
        print(f"❌ Error in image conversion/resizing: {e}")
        return None, None, {}

    mask_np = np.array(mask)
    inst_mask = mask_np.astype(np.int32)
    unique_ids = np.unique(inst_mask)
    num_instances = len(unique_ids[unique_ids != 0])
    if num_instances == 0:
        print("⚠️ No instance found, returning dummy red image")
        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), None, {}

    overlay_img = render_seg_overlay(img_np, inst_mask, overlay_alpha)
    seg_vis_cache = {"img_np": img_np, "inst_mask": inst_mask}
    return overlay_img, temp_mask_file.name, seg_vis_cache


@spaces.GPU
def count_cells_handler(use_box_choice, annot_value, overlay_alpha):
    """Counting handler - supports bounding box, returns only density map"""
    if annot_value is None or len(annot_value) < 1:
        return None, None, "⚠️ Please provide an image.", {}

    image_path = annot_value[0]
    bboxes = annot_value[1] if len(annot_value) > 1 else []

    print(f"🖼️ Image path: {image_path}")
    # One standardized image for both the model and the overlay (WYSIWYG).
    image_path = standardize_image(image_path)
    print(f"🧪 Standardized image: {image_path}")
    box_array = None
    if use_box_choice == "Yes" and bboxes:
        box = parse_bboxes(bboxes)
        if box:
            box_array = box
            print(f"📦 Using bounding boxes: {box_array}")

    try:
        print(f"🔢 Counting - Image: {image_path}")

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        result = run_count(
            COUNT_MODEL,
            image_path,
            box=box_array,
            device=device,
            visualize=True
        )
        
        if 'error' in result:
            return None, None, f"❌ Counting failed: {result['error']}", {}
        
        count = result['count']
        density_map = result['density_map']
        temp_density_file = tempfile.NamedTemporaryFile(delete=False, suffix=".npy")
        np.save(temp_density_file.name, density_map)
        print(f"💾 Density map saved to {temp_density_file.name}")
        

        try:
            img = Image.open(image_path)
            print("📷 Image mode:", img.mode, "size:", img.size)
        except Exception as e:
            print(f"❌ Failed to open image: {e}")
            return None, None, f"❌ Failed to open image: {str(e)}", {}

        try:
            img_rgb = img.convert("RGB").resize(density_map.shape[::-1], resample=Image.BILINEAR)
            img_np = np.array(img_rgb, dtype=np.float32)
            img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8)
            if img_np.max() > 1.5:
                img_np = img_np / 255.0
        except Exception as e:
            print(f"❌ Error in image conversion/resizing: {e}")
            return None, None, f"❌ Error in image conversion/resizing: {str(e)}", {}

        
        density_normalized = density_map.copy()
        if density_normalized.max() > 0:
            density_normalized = (density_normalized - density_normalized.min()) / (density_normalized.max() - density_normalized.min())
        
        overlay_img = render_count_overlay(img_np, density_normalized, overlay_alpha)
        result_text = f"✅ Detected {round(count)} objects"
        if use_box_choice == "Yes" and box_array:
            result_text += f"\n📦 Using bounding box: {box_array}"
        

        print(f"✅ Counting done - Count: {count:.1f}")

        count_vis_cache = {"img_np": img_np, "density_normalized": density_normalized}
        return overlay_img, temp_density_file.name, result_text, count_vis_cache
        
        
    except Exception as e:
        print(f"❌ Counting error: {e}")
        import traceback
        traceback.print_exc()
        return None, None, f"❌ Counting failed: {str(e)}", {}


def find_tif_dir(root_dir):
    """Recursively find the first directory containing .tif files"""
    for dirpath, _, filenames in os.walk(root_dir):
        if '__MACOSX' in dirpath:
            continue
        if any(f.lower().endswith('.tif') for f in filenames):
            return dirpath
    return None

def is_valid_tiff(filepath):
    """Check if a file is a valid TIFF image"""
    try:
        with Image.open(filepath) as img:
            img.verify()
            return True
    except Exception as e:
        return False

def find_valid_tif_dir(root_dir):
    """Recursively find the first directory containing valid .tif files"""
    for dirpath, dirnames, filenames in os.walk(root_dir):
        if '__MACOSX' in dirpath:
            continue
        
        potential_tifs = [
            os.path.join(dirpath, f) 
            for f in filenames 
            if f.lower().endswith(('.tif', '.tiff')) and not f.startswith('._')
        ]
        
        if not potential_tifs:
            continue
        
        valid_tifs = [f for f in potential_tifs if is_valid_tiff(f)]
        
        if valid_tifs:
            print(f"✅ Found {len(valid_tifs)} valid TIFF files in: {dirpath}")
            return dirpath
    
    return None

def create_ctc_results_zip(output_dir):
    """
    Create a ZIP file with CTC format results
    
    Parameters:
    -----------
    output_dir : str
        Directory containing tracking results (res_track.txt, etc.)
    
    Returns:
    --------
    zip_path : str
        Path to created ZIP file
    """
    # Create temp directory for ZIP
    temp_zip_dir = tempfile.mkdtemp()
    zip_filename = f"tracking_results_{time.strftime('%Y%m%d_%H%M%S')}.zip"
    zip_path = os.path.join(temp_zip_dir, zip_filename)
    
    print(f"📦 Creating results ZIP: {zip_path}")
    
    # Create ZIP with all tracking results
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        # Add all files from output directory
        for root, dirs, files in os.walk(output_dir):
            for file in files:
                file_path = os.path.join(root, file)
                arcname = os.path.relpath(file_path, output_dir)
                zipf.write(file_path, arcname)
                print(f"  📄 Added: {arcname}")
        
        # Add a README with summary
        readme_content = f"""Tracking Results Summary
            ========================

            Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}

            Files:
            ------
            - res_track.txt: CTC format tracking data
            Format: track_id start_frame end_frame parent_id
            
            - Segmentation masks

            For more information on CTC format:
            http://celltrackingchallenge.net/
            """
        zipf.writestr("README.txt", readme_content)
    
    print(f"✅ ZIP created: {zip_path} ({os.path.getsize(zip_path) / 1024:.1f} KB)")
    return zip_path


def get_well_spaced_color(track_id, num_colors=256):
    """Generate well-spaced colors, using contrasting colors for adjacent IDs"""

    golden_ratio = 0.618033988749895
    hue = (track_id * golden_ratio) % 1.0

    import colorsys
    rgb = colorsys.hsv_to_rgb(hue, 0.9, 0.95)
    return np.array(rgb)


def extract_first_frame(tif_dir):
    """
    Extract the first frame from a directory of TIF files
    
    Returns:
    --------
    first_frame_path : str
        Path to the first TIF frame
    """
    tif_files = natsorted(glob(os.path.join(tif_dir, "*.tif")) + 
                      glob(os.path.join(tif_dir, "*.tiff")))
    valid_tif_files = [f for f in tif_files 
                      if not os.path.basename(f).startswith('._') and is_valid_tiff(f)]
    
    if valid_tif_files:
        return valid_tif_files[0]
    return None

def create_tracking_visualization(tif_dir, output_dir, valid_tif_files, overlay_alpha=0.3):
    """
    Create an animated GIF/video showing tracked objects with consistent colors
    
    Parameters:
    -----------
    tif_dir : str
        Directory containing input TIF frames
    output_dir : str
        Directory containing tracking results (masks)
    valid_tif_files : list
        List of valid TIF file paths
    
    Returns:
    --------
    video_path : str
        Path to generated visualization (GIF or first frame)
    """
    import numpy as np
    from matplotlib import colormaps
    from skimage import measure
    import tifffile
    
    # Look for tracking mask files in output directory
    # Common CTC formats: man_track*.tif, mask*.tif, or numbered masks
    mask_files = natsorted(glob(os.path.join(output_dir, "mask*.tif")) + 
                       glob(os.path.join(output_dir, "man_track*.tif")) +
                       glob(os.path.join(output_dir, "*.tif")))
    
    if not mask_files:
        print("⚠️  No mask files found in output directory")
        # Return first frame as fallback
        return valid_tif_files[0]
    
    print(f"📊 Found {len(mask_files)} mask files")

    
    frames = []
    alpha = float(np.clip(overlay_alpha, 0.0, 1.0))  # Transparency for overlay
    
    # Process each frame
    num_frames = min(len(valid_tif_files), len(mask_files))
    for i in range(num_frames):
        try:
            # Load original image using tifffile (handles ZSTD compression)
            try:
                img_np = tifffile.imread(valid_tif_files[i])

                # Normalize to [0, 1] range based on actual data type and values
                if img_np.dtype == np.uint8:
                    img_np = img_np.astype(np.float32) / 255.0
                elif img_np.dtype == np.uint16:
                    # Normalize uint16 to [0, 1] using actual min/max
                    img_min, img_max = img_np.min(), img_np.max()
                    if img_max > img_min:
                        img_np = (img_np.astype(np.float32) - img_min) / (img_max - img_min)
                    else:
                        img_np = img_np.astype(np.float32) / 65535.0
                else:
                    # For float or other types, normalize based on actual range
                    img_np = img_np.astype(np.float32)
                    img_min, img_max = img_np.min(), img_np.max()
                    if img_max > img_min:
                        img_np = (img_np - img_min) / (img_max - img_min)
                    else:
                        img_np = np.clip(img_np, 0, 1)

                # Convert to RGB if grayscale
                if img_np.ndim == 2:
                    img_np = np.stack([img_np]*3, axis=-1)
                img_np = img_np.astype(np.float32)
                if img_np.max() > 1.5:
                    img_np = img_np / 255.0
            except Exception as e:
                print(f"⚠️  Error loading image frame {i}: {e}")
                # Fallback to PIL
                img = Image.open(valid_tif_files[i]).convert("RGB")
                img_np = np.array(img, dtype=np.float32) / 255.0
            
            # Load tracking mask using tifffile (handles ZSTD compression)
            try:
                mask = tifffile.imread(mask_files[i])
            except Exception as e:
                print(f"⚠️  Error loading mask frame {i}: {e}")
                # Fallback to PIL
                mask = np.array(Image.open(mask_files[i]))
            
            # Resize mask to match image if needed
            if mask.shape[:2] != img_np.shape[:2]:
                from scipy.ndimage import zoom
                zoom_factors = [img_np.shape[0] / mask.shape[0], img_np.shape[1] / mask.shape[1]]
                mask = zoom(mask, zoom_factors, order=0).astype(mask.dtype)
            
            # Create overlay
            overlay = img_np.copy()
            
            # Get unique track IDs (excluding background 0)
            track_ids = np.unique(mask)
            track_ids = track_ids[track_ids != 0]
            
            # Color each tracked object
            for track_id in track_ids:
                # Create binary mask for this track
                binary_mask = (mask == track_id)
                
                # Get consistent color for this track ID
                # color = np.array(cmap(int(track_id) % 256)[:3])
                color = get_well_spaced_color(int(track_id))
                
                # Blend color onto image
                overlay[binary_mask] = (1 - alpha) * overlay[binary_mask] + alpha * color
                
                # Draw contours (optional, adds yellow boundaries)
                try:
                    contours = measure.find_contours(binary_mask.astype(np.uint8), 0.5)
                    for contour in contours:
                        contour = contour.astype(np.int32)
                        valid_y = np.clip(contour[:, 0], 0, overlay.shape[0] - 1)
                        valid_x = np.clip(contour[:, 1], 0, overlay.shape[1] - 1)
                        overlay[valid_y, valid_x] = [1.0, 1.0, 0.0]  # Yellow contour
                except:
                    pass  # Skip contours if they fail
            
            # Convert to uint8
            overlay_uint8 = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
            frames.append(Image.fromarray(overlay_uint8))
            
            if i % 10 == 0 or i == num_frames - 1:
                print(f"  📸 Processed frame {i+1}/{num_frames}")
        
        except Exception as e:
            print(f"⚠️  Error processing frame {i}: {e}")
            import traceback
            traceback.print_exc()
            continue
    
    if not frames:
        print("⚠️  No frames were processed successfully")
        return valid_tif_files[0]
    
    # Save as animated GIF
    try:
        temp_gif = tempfile.NamedTemporaryFile(delete=False, suffix=".gif")
        frames[0].save(
            temp_gif.name,
            save_all=True,
            append_images=frames[1:],
            duration=200,  # 200ms per frame = 5fps
            loop=0
        )
        temp_gif.close()  # Close the file handle
        print(f"✅ Created tracking visualization GIF: {temp_gif.name}")
        print(f"   Size: {os.path.getsize(temp_gif.name)} bytes, Frames: {len(frames)}")
        return temp_gif.name
    except Exception as e:
        print(f"⚠️  Failed to create GIF: {e}")
        import traceback
        traceback.print_exc()
        # Return first frame as static image fallback
        try:
            temp_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
            frames[0].save(temp_img.name)
            temp_img.close()
            return temp_img.name
        except:
            return valid_tif_files[0]

@spaces.GPU
def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay_alpha, prev_track_vis_cache):
    """
    Tracking handler - processes a ZIP of TIF frames, supports bounding box, returns visualization and results ZIP
    
    Parameters:
    -----------
    use_box_choice : str
        "Yes" or "No" - whether to use bounding box annotation for tracking
    first_frame_annot : tuple or None
        (image_path, bboxes) from BBoxAnnotator, only used if user annotated first frame
    zip_file_obj : File
        Uploaded ZIP file containing TIF sequence
    """
    if zip_file_obj is None:
        return None, "⚠️  Please upload a ZIP file containing video frames (.zip)", None, None, {}
    
    cleanup_tracking_cache(prev_track_vis_cache)
    temp_dir = None
    output_temp_dir = None
    
    try:
        # Parse bounding box if provided
        box_array = None
        if use_box_choice == "Yes" and first_frame_annot is not None:
            if isinstance(first_frame_annot, (list, tuple)) and len(first_frame_annot) > 1:
                bboxes = first_frame_annot[1]
                if bboxes:
                    box = parse_bboxes(bboxes)
                    if box:
                        box_array = box
                        print(f"📦 Using bounding boxes: {box_array}")
        
        # Extract input ZIP
        temp_dir = tempfile.mkdtemp()
        print(f"\n📦 Extracting to temporary directory: {temp_dir}")

        with zipfile.ZipFile(zip_file_obj.name, 'r') as zip_ref:
            extracted_count = 0
            skipped_count = 0
            
            for member in zip_ref.namelist():
                basename = os.path.basename(member)
                
                if ('__MACOSX' in member or 
                    basename.startswith('._') or 
                    basename.startswith('.DS_Store') or
                    member.endswith('/')):
                    skipped_count += 1
                    continue
                
                try:
                    zip_ref.extract(member, temp_dir)
                    extracted_count += 1
                    if basename.lower().endswith(('.tif', '.tiff')):
                        print(f"📄 Extracted TIFF: {basename}")
                except Exception as e:
                    print(f"⚠️  Failed to extract {member}: {e}")

            print(f"\n📊 Extracted: {extracted_count} files, Skipped: {skipped_count} files")

        # Find valid TIFF directory
        tif_dir = find_valid_tif_dir(temp_dir)
        
        if tif_dir is None:
            return None, "❌ Did not find valid TIF directory", None, None, {}
        
        # Validate TIFF files
        tif_files = natsorted(glob(os.path.join(tif_dir, "*.tif")) + 
                          glob(os.path.join(tif_dir, "*.tiff")))
        valid_tif_files = [f for f in tif_files 
                          if not os.path.basename(f).startswith('._') and is_valid_tiff(f)]
        
        if len(valid_tif_files) == 0:
            return None, "❌ Did not find valid TIF files", None, None, {}

        print(f"📈 Using {len(valid_tif_files)} TIF files")

        # Store paths for later visualization
        first_frame_path = valid_tif_files[0]

        # Create temporary output directory for CTC results
        output_temp_dir = tempfile.mkdtemp()
        print(f"💾 CTC-format results will be saved to: {output_temp_dir}")

        # Run tracking with optional bounding box
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        result = run_track(
            TRACK_MODEL,
            video_dir=tif_dir,
            box=box_array,  # Pass bounding box if specified
            device=device,
            output_dir=output_temp_dir
        )
        
        if 'error' in result:
            return None, f"❌ Tracking failed: {result['error']}", None, None, {}
        
        # Create visualization video of tracked objects
        print("\n🎬 Creating tracking visualization...")
        try:
            tracking_video = create_tracking_visualization(
                tif_dir, 
                output_temp_dir, 
                valid_tif_files,
                overlay_alpha=overlay_alpha
            )
        except Exception as e:
            print(f"⚠️  Failed to create visualization: {e}")
            import traceback
            traceback.print_exc()
            # Fallback to first frame if visualization fails
            try:
                tracking_video = Image.open(first_frame_path)
            except:
                tracking_video = None
        
        # Create downloadable ZIP with results
        try:
            results_zip = create_ctc_results_zip(output_temp_dir)
        except Exception as e:
            print(f"⚠️  Failed to create ZIP: {e}")
            results_zip = None
        
        bbox_info = ""
        if box_array:
            bbox_info = f"\n🔲 Using bounding box: [{box_array[0][0]}, {box_array[0][1]}, {box_array[0][2]}, {box_array[0][3]}]"

        result_text = (
            f"✅ Tracking completed!\n"
            f"\n"
            f"🖼️  Processed frames: {len(valid_tif_files)}{bbox_info}\n"
            f"\n"
            f"📥 Use the \"Download (.zip)\" button above to get the results, which include:\n"
            f"- res_track.txt (CTC-format tracking data)\n"
            f"- Other tracking-related files\n"
            f"- README.txt (Results description)"
        )

        if use_box_choice == "Yes" and box_array:
            result_text += f"\n📦 Using bounding box: {box_array}"

        print(f"\n✅ Tracking completed")

        track_vis_cache = {
            "tif_dir": tif_dir,
            "valid_tif_files": valid_tif_files,
            "output_dir": output_temp_dir,
            "input_temp_dir": temp_dir,
        }

        return results_zip, result_text, gr.update(visible=True), tracking_video, track_vis_cache

    except zipfile.BadZipFile:
        return None, "❌ Not a valid ZIP file", None, None, {}
    except Exception as e:
        import traceback
        traceback.print_exc()
        
        # Clean up on error
        for d in [temp_dir, output_temp_dir]:
            if d:
                try:
                    shutil.rmtree(d)
                except:
                    pass
        return None, f"❌ Tracking failed: {str(e)}", None, None, {}



# ===== Example Images =====
example_images_seg = [f for f in glob("example_imgs/seg/*")]
example_images_cnt = [f for f in glob("example_imgs/cnt/*")]
example_tracking_zips = [f for f in glob("example_imgs/tra/*.zip")]

_CHANNEL_NOTE = {1: "grayscale", 3: "RGB", 4: "RGBA"}


def describe_image_info(raw_path, frame=1):
    """Metadata line under the result: file, size, channels, bit depth, frames.

    Reads the *original* upload (not the standardized preview) so the numbers
    describe what the user actually provided. Metadata only - no pixel decode.
    """
    if not raw_path or not os.path.exists(raw_path):
        return gr.update(value="", visible=False)

    info = inspect_image(raw_path)
    bits = bit_depth(raw_path)

    fields = [f"<span class='ii-k'>File:</span> {html.escape(os.path.basename(raw_path))}"]
    if info.width and info.height:
        fields.append(f"<span class='ii-k'>Size:</span> {info.width} × {info.height} px")
    note = _CHANNEL_NOTE.get(info.channels)
    fields.append(f"<span class='ii-k'>Channels:</span> {info.channels}"
                  + (f" ({note})" if note else ""))
    if bits:
        fields.append(f"<span class='ii-k'>Bit depth:</span> {bits}-bit")
    if info.frames > 1:
        fields.append(f"<span class='ii-k'>Frames:</span> {int(frame)} / {info.frames}")
    if info.sub_frames > 1:
        fields.append(f"<span class='ii-k'>{info.sub_label.capitalize()}s:</span> {info.sub_frames}")

    body = "".join(f"<span class='ii-f'>{f}</span>" for f in fields)
    return gr.update(
        value=f"<div class='ii-wrap'><span class='ii-title'>ⓘ Image information</span>{body}</div>",
        visible=True,
    )


# ===== Gradio UI =====
CSS = """
/* ── Layout ──────────────────────────────────────────── */
.gradio-container {
    max-width: 1320px !important;
    margin: 0 auto !important;
    font-family: 'Inter', 'Segoe UI', system-ui, sans-serif !important;
    background: #fff !important;
}

/* ── Header markdown polish ───────────────────────────── */
.gradio-container .prose h1 {
    font-size: 1.4rem !important;
    font-weight: 700 !important;
    color: #1c1c1a !important;
    letter-spacing: -0.01em !important;
    margin-bottom: 2px !important;
}
.gradio-container .prose h3 {
    font-size: 1rem !important;
    font-weight: 600 !important;
    color: #bf1e2e !important;
    margin-top: 14px !important;
    margin-bottom: 4px !important;
}
.gradio-container .prose p {
    margin-top: 4px !important;
    margin-bottom: 6px !important;
    color: #5c5a52 !important;
    line-height: 1.7 !important;
}
.gradio-container .prose ul,
.gradio-container .prose ol {
    margin-top: 4px !important;
    margin-bottom: 6px !important;
}
.gradio-container .prose li {
    color: #5c5a52 !important;
    line-height: 1.7 !important;
}

/* ── Top-level header section: slim, no gradient card ─── */
.gradio-container > .gap > .prose:first-child {
    background: transparent !important;
    border: none !important;
    border-radius: 0 !important;
    padding: 18px 4px 10px !important;
    margin-bottom: 8px !important;
    box-shadow: none !important;
}

/* ── Mode selector: two-line task cards ──────────────── */
/* Gradio 5.x markup: .tabs > .tab-wrapper > .tab-container[role=tablist] > button */

/* Undo Gradio's fixed-height, clipped, underlined nav strip */
.tabs > .tab-wrapper {
    height: auto !important;
    padding-bottom: 0 !important;
    margin-bottom: 20px !important;
}
.tabs > .tab-wrapper > .tab-container {
    height: auto !important;
    overflow: visible !important;
    gap: 14px !important;
    align-items: stretch !important;
}
.tabs > .tab-wrapper > .tab-container::after { display: none !important; }  /* bottom hairline */
.tabs > .tab-wrapper > .overflow-menu { display: none !important; }         /* "…" menu */

/* The card itself */
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button {
    flex: 1 1 0 !important;
    display: grid !important;
    grid-template-columns: auto 1fr !important;
    grid-template-rows: auto auto !important;
    column-gap: 14px !important;
    row-gap: 2px !important;
    align-items: center !important;
    justify-items: start !important;
    text-align: left !important;
    white-space: normal !important;
    height: auto !important;
    padding: 16px 20px !important;
    border-radius: 10px !important;
    border: 1.5px solid #e7e5e1 !important;
    background: #fff !important;
    box-shadow: none !important;
    color: #1c1c1a !important;
    font-size: 15px !important;
    font-weight: 700 !important;
    line-height: 1.3 !important;
    transition: border-color 0.12s, box-shadow 0.12s, color 0.12s !important;
}

/* Left icon: spans both rows */
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button::before {
    grid-column: 1 !important;
    grid-row: 1 / span 2 !important;
    font-size: 26px !important;
    line-height: 1 !important;
    align-self: center !important;
}

/* Subtitle: row 2 of the text column.
   NOTE: also resets Gradio's .selected::after underline bar (position/size/bg/animation). */
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button::after {
    grid-column: 2 !important;
    grid-row: 2 !important;
    position: static !important;
    width: auto !important;
    height: auto !important;
    background: none !important;
    animation: none !important;
    font-size: 13px !important;
    font-weight: 400 !important;
    color: #7a7873 !important;
    line-height: 1.35 !important;
}

/* Per-tab icon + description */
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button:nth-of-type(1)::before { content: "🎨"; }
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button:nth-of-type(1)::after  { content: "Instance segmentation of microscopic objects"; }
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button:nth-of-type(2)::before { content: "🔢"; }
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button:nth-of-type(2)::after  { content: "Microscopy Object Counting Analysis"; }
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button:nth-of-type(3)::before { content: "🎬"; }
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button:nth-of-type(3)::after  { content: "Microscopy Object Video Tracking - Supports ZIP Upload"; }

/* States */
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button:hover {
    border-color: #e6bdb8 !important;
    background: #fdf3f2 !important;
}
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button.selected {
    color: #bf1e2e !important;
    border-color: #bf1e2e !important;
    box-shadow: 0 0 0 1px #bf1e2e !important;
    background: #fff !important;
}
.tabs > .tab-wrapper > .tab-container:not(.visually-hidden) > button.selected::after {
    color: #7a7873 !important;   /* description stays gray when selected */
}

/* ── Buttons ─────────────────────────────────────────── */
button.primary {
    background: #bf1e2e !important;
    border: none !important;
    border-radius: 8px !important;
    color: #fff !important;
    font-weight: 600 !important;
    font-size: 15px !important;
    box-shadow: none !important;
    transition: background 0.12s ease !important;
}
button.primary:hover {
    background: #a61a28 !important;
}
button.secondary {
    border-radius: 8px !important;
    font-weight: 500 !important;
    border: 1.5px solid #d8d5cd !important;
    color: #3a3a36 !important;
    transition: border-color 0.12s, color 0.12s, background 0.12s !important;
}
button.secondary:hover {
    border-color: #bf1e2e !important;
    color: #bf1e2e !important;
    background: #fdf3f2 !important;
}

/* ── Blocks and panels ───────────────────────────────── */
.gradio-container .block { border-radius: 8px !important; }
.gradio-container .gr-form,
.gradio-container .gr-box,
.gradio-container .gr-panel {
    border-radius: 8px !important;
    border-color: #e7e5e1 !important;
}

/* ── Labels ──────────────────────────────────────────── */
label { font-weight: 500 !important; color: #374151 !important; }

/* File/media title pills use .block-label on a <label>, so the rule above
   forces their text dark. Restore the theme color/weight so they match the
   other native labels (e.g. "Overlay Opacity"). */
.block-label,
.block-label span {
    color: var(--block-label-text-color) !important;
    font-weight: var(--block-label-text-weight) !important;
}

/* ── Collapsible "Image requirements" accordion ───────── */

/* The label is a <button class="label-wrap"> whose text sits in an inner <span>.
   Gradio styles that span directly:
       span.svelte-xxx { font-weight: var(--section-header-text-weight);  // = 400
                         font-size:   var(--section-header-text-size); }
   so font-size/weight set on the button is overridden and has no effect - the
   span itself must be targeted. Colour is not set on the span, so it inherits. */
.req-accordion .label-wrap {
    color: var(--block-label-text-color) !important;
}
.req-accordion .label-wrap span:not(.icon) {
    font-size: 0.9rem !important;
    font-weight: 700 !important;
    color: inherit !important;
}

/* ── Frame headings (labels above media frames) ──────── */
/* Neutralize the Gradio block wrapper so it doesn't clip/round the pill */
.frame-heading {
    /* The column has a 16px flex gap; pull the pill back down so it sits
       close to the frame below instead of floating far above it. */
    margin-bottom: -8px !important;
    padding: 0 !important;
    min-width: 0 !important;
    min-height: 0 !important;
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
    border-radius: 0 !important;
    overflow: visible !important;
}
/* Match the theme's native field labels (e.g. "Overlay Opacity",
   "Download Original Prediction") exactly: same colors, borderless, fully
   rounded — driven by theme variables so it adapts to light/dark mode. */
.frame-heading p {
    display: inline-block !important;
    width: fit-content !important;
    background: var(--block-label-background-fill) !important;
    color: var(--block-label-text-color) !important;
    border: none !important;
    border-radius: var(--block-label-radius) !important;
    font-weight: var(--block-label-text-weight) !important;
    font-size: var(--block-label-text-size) !important;
    line-height: 1.3 !important;
    padding: 5px 12px !important;
    margin: 0 !important;
}

/* ── Image output ────────────────────────────────────── */
.uniform-height {
    height: 480px !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    border-radius: 8px !important;
    background: #f0efeb !important;
}
.uniform-height img, .uniform-height canvas {
    max-height: 480px !important;
    object-fit: contain !important;
}

/* ── Input annotator: match the 480px result panel ───── */
/* width/height auto keeps the <img> box equal to the rendered picture (no
   letterbox gaps), so the annotator's ResizeObserver still maps bbox
   coordinates onto the right rect after scaling. */
.uniform-height-input img {
    max-height: 480px !important;
    width: auto !important;
    height: auto !important;
    max-width: 100% !important;
    object-fit: contain !important;
    margin: 0 auto !important;
}

/* ── Density map output ──────────────────────────────── */
#density_map_output { height: 480px !important; }
#density_map_output .image-container { height: 480px !important; }
#density_map_output img {
    height: 460px !important;
    width: auto !important;
    max-width: 95% !important;
    object-fit: contain !important;
}

/* ── "How to use" collapsible instructions (was a permanent wall of text) ── */
.howto-accordion {
    border: 1px solid #e7e5e1 !important;
    border-radius: 8px !important;
    margin-bottom: 14px !important;
}
.howto-accordion .label-wrap {
    color: #bf1e2e !important;
}
.howto-accordion .label-wrap span:not(.icon) {
    font-size: 0.9rem !important;
    font-weight: 700 !important;
    color: #bf1e2e !important;
}
.howto-accordion .prose p,
.howto-accordion .prose li {
    font-size: 0.87rem !important;
    color: #5c5a52 !important;
}
.howto-accordion .prose strong {
    color: #1c1c1a !important;
}

/* ── "How to use": top-right button + centered modal ─── */
#header_row {
    align-items: center !important;
    justify-content: space-between !important;
    flex-wrap: nowrap !important;
}
#header_title { flex: 1 1 auto !important; }
/* scoped to the app title so the generic .prose h1 rule stays put */
#header_title h1 {
    font-size: 1.85rem !important;
    letter-spacing: -0.02em !important;
}
#how_to_use_btn {
    flex: 0 0 auto !important;
    width: auto !important;
    min-width: 0 !important;
    align-self: center !important;
    background: #bf1e2e !important;
    color: #fff !important;
    font-weight: 600 !important;
    font-size: 13.5px !important;
    border: none !important;
    border-radius: 6px !important;
    padding: 9px 16px !important;
    box-shadow: none !important;
}
#how_to_use_btn:hover { background: #a61a28 !important; }
#how_to_use_modal {
    position: fixed !important;
    top: 8% !important;
    left: 50% !important;
    transform: translateX(-50%) !important;
    width: min(600px, 92vw) !important;
    max-height: 82vh !important;
    overflow-y: auto !important;
    background: #fff !important;
    border: 1px solid #e7e5e1 !important;
    border-radius: 12px !important;
    box-shadow: 0 24px 70px rgba(0,0,0,0.28) !important;
    z-index: 1000 !important;
    padding: 22px 26px !important;
}
#how_to_use_modal h3 { margin-top: 12px !important; }
#close_how_to_use_btn {
    width: auto !important;
    min-width: 0 !important;
    margin-top: 8px !important;
    border: 1px solid #d8d5cd !important;
    border-radius: 6px !important;
    color: #3a3a36 !important;
    background: #fff !important;
}

/* ── Footer note ─────────────────────────────────────── */
#footer_note {
    border-top: 1px solid #e7e5e1 !important;
    padding-top: 16px !important;
    margin-top: 24px !important;
}
#footer_note p { font-size: 12.5px !important; color: #5c5a52 !important; line-height: 1.8 !important; }
#footer_note strong { color: #1c1c1a !important; }

/* ── Requirements: collapsible panel (open by default) ── */
.req-panel {
    border: 1.5px solid #e6bdb8 !important;
    border-radius: 10px !important;
    padding: 14px 16px !important;
    margin-bottom: 14px !important;
    background: #fdf3f2 !important;
    box-shadow: 0 1px 4px rgba(28, 28, 26, 0.07) !important;
}
.req-panel > *, .req-panel .block, .req-panel .form { background: transparent !important; border: none !important; }
/* Accordion header: crimson, with a divider only while expanded */
.req-panel .label-wrap {
    margin: 0 !important;
    padding: 0 !important;
    color: #bf1e2e !important;
}
.req-panel .label-wrap.open {
    margin-bottom: 10px !important;
    padding-bottom: 8px !important;
    border-bottom: 1.5px solid #f3d9d5 !important;
}
.req-panel .label-wrap span:not(.icon) {
    font-size: 0.92rem !important;
    font-weight: 800 !important;
    color: #bf1e2e !important;
}
.req-panel .label-wrap .icon { color: #bf1e2e !important; }
.req-warn {
    font-size: 12px !important;
    color: #8a5a1c !important;
    background: #faf3e6 !important;
    border: 1px solid #f0dfc0 !important;
    border-radius: 5px !important;
    padding: 8px 10px !important;
    margin-bottom: 8px !important;
}
.req-row {
    display: flex !important;
    justify-content: space-between !important;
    gap: 16px !important;
    font-size: 12.5px !important;
    padding: 7px 0 !important;
    border-bottom: 1px solid #f0efeb !important;
}
.req-row:last-of-type { border-bottom: none !important; }
.req-label { color: #8b8981 !important; flex-shrink: 0 !important; }
.req-val { color: #1c1c1a !important; font-weight: 500 !important; text-align: right !important; }
.req-val a { color: #bf1e2e !important; }
.req-val code {
    font-size: 12px !important;
    background: #f5f3f0 !important;
    padding: 1px 5px !important;
    border-radius: 4px !important;
    color: #1c1c1a !important;
}
.req-note { font-size: 12px !important; color: #5c5a52 !important; margin-top: 10px !important; line-height: 1.5 !important; }

/* ── Result card ─────────────────────────────────────── */
/* The row carries the -8px hug to the frame below, so the pill keeps a normal
   box and centre-aligns with the download button instead of sitting lower. */
.result-header {
    align-items: center !important;
    gap: 12px !important;
    /* the download button makes this row taller than a bare pill, so it needs
       a deeper pull-down than .frame-heading's -8px to hug the card below */
    margin-bottom: -16px !important;
}
.result-header .frame-heading { margin-bottom: 0 !important; }
.result-download {
    border: 1px solid #d9d7d2 !important;
    background: #fff !important;
    color: #1c1c1a !important;
    font-weight: 600 !important;
    border-radius: 8px !important;
    box-shadow: none !important;
}
.result-download:hover {
    border-color: #bf1e2e !important;
    color: #bf1e2e !important;
    background: #fdf3f2 !important;
}

/* Opacity on one line. Gradio nests the label and the number box together in
   .head; display:contents lifts them into .wrap's flex row so the track can be
   ordered between them. */
.opacity-row .wrap {
    display: flex !important;
    flex-direction: row !important;
    align-items: center !important;
    gap: 14px !important;
}
.opacity-row .head { display: contents !important; }
.opacity-row .head label { order: 1 !important; flex: 0 0 auto !important; margin: 0 !important; }
/* the theme gives block labels a tinted pill; plain text here */
.opacity-row .head label,
.opacity-row .head label span {
    background: transparent !important;
    padding: 0 !important;
    color: #1c1c1a !important;
    font-weight: 600 !important;
    font-size: 14px !important;
    white-space: nowrap !important;
}
.opacity-row .slider_input_container { order: 2 !important; flex: 1 1 auto !important; }
.opacity-row .tab-like-container { order: 3 !important; flex: 0 0 auto !important; }
.opacity-row .min_value, .opacity-row .max_value { display: none !important; }

/* ── Section cards: give each block its own frame ────── */
.section-card {
    border: 1px solid #e7e5e1 !important;
    border-radius: 12px !important;
    background: #fff !important;
    padding: 16px !important;
    box-shadow: 0 1px 3px rgba(28, 28, 26, 0.04) !important;
}

/* No card here - it's one short control. But Gradio wraps lone inputs in a
   .form that paints its own fill, which shows as a grey band behind the row
   once the control is only one line tall. */
.bbox-slot .form,
.bbox-slot > .block {
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
}

/* Radio label + options on one line (Block = BlockTitle span + .wrap) */
.inline-radio {
    display: flex !important;
    flex-direction: row !important;
    align-items: center !important;
    flex-wrap: wrap !important;
    gap: 12px !important;
    padding: 0 !important;
    border: none !important;
    background: transparent !important;
    box-shadow: none !important;
}
.inline-radio > span { margin: 0 !important; flex: 0 0 auto !important; }
.inline-radio .wrap { flex: 0 0 auto !important; margin: 0 !important; }

/* ── Image information line under the result ─────────── */
/* gr.HTML puts elem_classes on BOTH the outer Block and the inner .prose div,
   so draw the box once on the inner one and strip the outer wrapper. */
.block.img-info {
    border: none !important;
    background: transparent !important;
    padding: 0 !important;
    box-shadow: none !important;
}
.img-info.prose {
    border: 1px solid #e7e5e1 !important;
    border-radius: 10px !important;
    background: #fff !important;
    padding: 12px 14px !important;
    margin: 4px 0 8px !important;
}
.ii-wrap {
    display: flex !important;
    flex-wrap: wrap !important;
    align-items: center !important;
    gap: 6px 20px !important;
    font-size: 12.5px !important;
    color: #1c1c1a !important;
    line-height: 1.5 !important;
}
.ii-title { color: #8b8981 !important; white-space: nowrap !important; }
.ii-f { white-space: nowrap !important; }
.ii-k { color: #5c5a52 !important; font-weight: 700 !important; margin-right: 3px !important; }

/* ── Feedback: one row (rating + comment + submit) ───── */
.feedback-title p { font-size: 14px !important; font-weight: 700 !important; color: #1c1c1a !important; margin: 4px 0 8px !important; }
.feedback-row {
    align-items: center !important;
    gap: 10px !important;
    /* wrap to a second line rather than crushing the comment box when the
       column gets narrow */
    flex-wrap: wrap !important;
    margin: 0 !important;
}
/* kill the per-child block padding/margins that were spreading the row out */
.feedback-row > *,
.feedback-row .block,
.feedback-row .form,
.feedback-row .wrap {
    margin: 0 !important;
    padding: 0 !important;
    border: none !important;
    background: transparent !important;
    min-width: 0 !important;
}
/* the comment box (2nd child) keeps a usable width; it takes the free space
   but never collapses — below its floor the row wraps instead */
.feedback-row > *:nth-child(2) {
    flex: 1 1 200px !important;
    min-width: 170px !important;
}
/* the reset above must not let the button shrink below its own label */
.feedback-row > *:last-child {
    flex: 0 0 auto !important;
    width: auto !important;
    min-width: max-content !important;
}
.feedback-row button {
    color: #bf1e2e !important;
    border: 1px solid #bf1e2e !important;
    background: #fff !important;
    font-weight: 700 !important;
    font-size: 14px !important;
    white-space: nowrap !important;
    padding: 9px 16px !important;
    border-radius: 8px !important;
    line-height: 1.2 !important;
    width: auto !important;
    min-width: max-content !important;
}
/* the input sits in a container=False textbox, so it has no frame of its own */
.feedback-row textarea,
.feedback-row input[type="text"] {
    border: 1px solid #ddd9d2 !important;
    border-radius: 8px !important;
    background: #fff !important;
    padding: 9px 12px !important;
    font-size: 14px !important;
    line-height: 1.3 !important;
    resize: none !important;
    overflow: hidden !important;
}

/* Star rating: single glyphs, reversed DOM + row-reverse so :checked fills left→right */
.star-rating, .star-rating * { overflow: visible !important; max-height: none !important; }
.star-rating { border: none !important; background: transparent !important; width: auto !important; min-width: max-content !important; flex: 0 0 auto !important; margin: 0 !important; }
.star-rating > * { width: auto !important; min-width: max-content !important; }
.star-rating .wrap {
    display: flex !important;
    flex-direction: row-reverse !important;
    justify-content: flex-end !important;
    gap: 2px !important;
    flex-wrap: nowrap !important;
}
.star-rating label {
    border: none !important;
    background: transparent !important;
    box-shadow: none !important;
    padding: 0 !important;
    margin: 0 !important;
    cursor: pointer !important;
    line-height: 1 !important;
}
.star-rating label input { position: absolute !important; opacity: 0 !important; width: 0 !important; height: 0 !important; pointer-events: none !important; margin: 0 !important; }
.star-rating label span { font-size: 0 !important; color: transparent !important; }
.star-rating label span::before {
    content: "★";
    font-size: 19px !important;
    color: #dcd8d0 !important;
    transition: color .12s ease;
}
/* hover preview: fill hovered + all visually-left (DOM-following) siblings */
.star-rating label:hover span::before,
.star-rating label:hover ~ label span::before { color: #ef8b93 !important; }
/* selection: fill checked + all visually-left (DOM-following) siblings */
.star-rating label:has(input:checked) span::before,
.star-rating label:has(input:checked) ~ label span::before { color: #bf1e2e !important; }
"""

HOWTO_TITLE_SEG = "### ⓘ How to use (Segmentation)"
HOWTO_TITLE_COUNT = "### ⓘ How to use (Counting)"
HOWTO_TITLE_TRACK = "### ⓘ How to use (Tracking)"

HOWTO_SEG = """
**Instructions:**
1. Select an example image from the Example Image Gallery or upload your own image (**Please see the "Image Upload Requirements" section before uploading**)
2. (Optional) Specify a target object with a bounding box and select "Yes", or click "Run Segmentation" directly
3. Click "Run Segmentation"
4. View the segmentation results (you can adjust the overlay opacity by sliding the opacity bar below the visualization), download the original predicted mask (.tif format); if needed, click "Clear Selection" to choose a new image

💡 Want to reuse an image? Upload it under "➕ Upload New Example Image to Gallery" and click "➕ Add to Gallery". The uploaded image will appear at the front of the gallery, and can be loaded by clicking its thumbnail. Gallery additions last for the current session only and are cleared when you reload the page.

🤘 Tell us about your experience by rating and submitting feedback, which would greatly help us improve the framework!
"""

HOWTO_COUNT = """
**Instructions:**
1. Select an example image from the Example Image Gallery or upload your own image (**Please see the "Image Upload Requirements" section before uploading**)
2. (Optional) Specify a target object with a bounding box and select "Yes", or click "Run Counting" directly
3. Click "Run Counting"
4. View the density map (you can adjust the density opacity by sliding the opacity bar below the visualization), download the original prediction (.npy format); if needed, click "Clear Selection" to choose a new image to run

💡 Want to reuse an image? Upload it under "➕ Upload New Example Image to Gallery" and click "➕ Add to Gallery". The uploaded image will appear at the front of the gallery, and can be loaded by clicking its thumbnail. Gallery additions last for the current session only and are cleared when you reload the page.

🤘 Tell us about your experience by rating and submitting feedback, which would greatly help us improve the framework!
"""

# The whole stylesheet is written for a light background, so dark mode renders
# badly. Gradio has no server-side switch: the frontend picks light/dark from
# the `__theme` query param, else the OS preference. Pin the param, and also
# strip the `dark` class directly so there's no flash before the reload.
FORCE_LIGHT_JS = """
function () {
    document.documentElement.classList.remove('dark');
    document.body.classList.remove('dark');
    const url = new URL(window.location);
    if (url.searchParams.get('__theme') !== 'light') {
        url.searchParams.set('__theme', 'light');
        window.location.replace(url.href);
    }
}
"""

HOWTO_TRACK = """
**Instructions:**
1. Select a video from the example library or Upload your own video ZIP file. The ZIP should contain a sequence of TIF images named in chronological order (e.g., t000.tif, t001.tif...)
2. (Optional) Specify a target object with a bounding box on the first frame and select "Yes", or click "Run Tracking" directly
3. Click "Run Tracking"
4. View the tracking results (you can adjust the overlay opacity by sliding the opacity bar below the visualization), download the CTC format results; if needed, click "Clear Selection" to choose a new ZIP file to run

💡 Want to reuse a video? Upload the ZIP under "➕ Upload New Example Image to Gallery" and click "➕ Add to Gallery". The uploaded video will appear at the front of the gallery, and can be loaded by clicking its thumbnail. Gallery additions last for the current session only and are cleared when you reload the page.

🤘 Tell us about your experience by rating and submitting feedback, which would greatly help us improve the framework!
"""

with gr.Blocks(
    title="Microscopy Analysis Suite",
    theme=gr.themes.Soft(
        primary_hue=gr.themes.Color(c50="#fdeceb", c100="#fbd6d3", c200="#f5a8a3", c300="#ef7a73",
                                     c400="#c94a44", c500="#bf1e2e", c600="#a9192a", c700="#8f1424",
                                     c800="#75101e", c900="#5c0c17", c950="#42080f"),
        secondary_hue=gr.themes.Color(c50="#fdeceb", c100="#fbd6d3", c200="#f5a8a3", c300="#ef7a73",
                                       c400="#c94a44", c500="#bf1e2e", c600="#a9192a", c700="#8f1424",
                                       c800="#75101e", c900="#5c0c17", c950="#42080f"),
        neutral_hue=gr.themes.colors.slate,
        # A weight must be listed here to be usable; the default is (400, 600).
        font=gr.themes.GoogleFont("Inter", weights=(400, 600, 700, 800)),
    ),
    css=CSS,
    js=FORCE_LIGHT_JS,
) as demo:
    with gr.Row(elem_id="header_row"):
        gr.Markdown("# 🔬 MicroscopyMatching: Microscopy Image Analysis Suite", elem_id="header_title")
        how_to_use_btn = gr.Button("ⓘ How to use", elem_id="how_to_use_btn", scale=0)
    with gr.Column(visible=False, elem_id="how_to_use_modal") as how_to_use_modal:
        how_to_use_title = gr.Markdown(HOWTO_TITLE_SEG, elem_id="how_to_use_title")
        howto_md = gr.Markdown(HOWTO_SEG)
        close_how_to_use_btn = gr.Button("Close", elem_id="close_how_to_use_btn", size="sm")
    how_to_use_btn.click(lambda: gr.update(visible=True), outputs=how_to_use_modal)
    close_how_to_use_btn.click(lambda: gr.update(visible=False), outputs=how_to_use_modal)
    

    # Must stay a callable: a plain value is deepcopied, so every session would
    # share one id.
    current_query_id = gr.State(lambda: str(uuid.uuid4()))
    user_uploaded_examples = gr.State(example_images_seg.copy())
    seg_vis_state = gr.State({})
    count_vis_state = gr.State({})
    track_vis_state = gr.State({})
    # The annotator only holds the flattened preview, so keep the original for
    # re-extracting other stack frames.
    seg_raw_state = gr.State(None)
    count_raw_state = gr.State(None)
    
    with gr.Tabs():
        # ===== Tab 1: Segmentation =====
        with gr.Tab("Segment") as seg_tab:
            # gr.Markdown("## Instance Segmentation of Microscopic Objects")
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("🖼️ Upload Image (Optional: Provide a Bounding Box)", elem_classes="frame-heading")
                    with gr.Column(elem_classes="section-card"):
                        annotator = BBoxAnnotator(
                            show_label=False,
                            categories=["cell"],
                            elem_classes="uniform-height-input",
                        )
                        seg_frame_slider = gr.Slider(
                            minimum=1, maximum=1, step=1, value=1,
                            label="📚 Stack frame (choose frame to use)",
                            visible=False,
                        )
                        seg_zplane_slider = gr.Slider(
                            minimum=1, maximum=1, step=1, value=1,
                            label="🔬 Z-plane (choose plane to use)",
                            visible=False,
                        )

                    # Example Images Gallery
                    gr.Markdown("📁 Example Image Gallery", elem_classes="frame-heading")
                    with gr.Column(elem_classes="section-card"):
                        example_gallery = gr.Gallery(
                            show_label=True,
                            columns=len(example_images_seg),
                            rows=1,
                            height=120,
                            object_fit="cover",
                            show_download_button=False
                        )

                    with gr.Column(elem_classes="bbox-slot"):
                        use_box_radio = gr.Radio(
                            choices=["Yes", "No"],
                            value="No",
                            label="🔲 Specify Bounding Box?",
                            elem_classes="inline-radio",
                        )
                    with gr.Row():
                        run_seg_btn = gr.Button("▶️ Run Segmentation", variant="primary", size="lg")
                        clear_btn = gr.Button("🔄 Clear Selection", variant="secondary")

                    with gr.Accordion("📋 Image Upload Requirements", open=True, elem_classes="req-panel"):
                        gr.HTML(
                            """
                            <div class="req-warn">⚠️ Do not upload images or data containing protected health information (PHI) or personally identifiable information (PII).</div>
                            <div class="req-row"><span class="req-label">Formats</span><span class="req-val">TIFF / OME-TIFF (8/16/32-bit, stacks, multi-channel), PNG, JPG, and other formats supported by <a href="https://github.com/cgohlke/tifffile" target="_blank">tifffile</a></span></div>
                            <div class="req-row"><span class="req-label">Max size</span><span class="req-val">4096 × 4096 px</span></div>
                            <div class="req-row"><span class="req-label">Channels</span><span class="req-val">RGB (3-channel) or grayscale; channels beyond the first 3 are ignored</span></div>
                            <div class="req-row"><span class="req-label">Stacks</span><span class="req-val">first frame loaded by default; frame slider appears after upload</span></div>
                            <div class="req-note">Alternatively, convert to an 8-bit RGB/grayscale PNG/JPG/TIFF first (e.g. with ImageJ/Fiji).</div>
                            """
                        )

                    # Upload Example Image
                    gr.Markdown("➕ Upload New Example Image to Gallery", elem_classes="frame-heading")
                    with gr.Column(elem_classes="section-card"):
                        image_uploader = gr.File(
                            show_label=False,
                            file_types=["image"] + list(TIFF_EXTENSIONS),
                            type="filepath"
                        )
                        add_to_gallery_btn = gr.Button("➕ Add to Gallery", variant="secondary", size="sm")
                        add_gallery_status = gr.Markdown(visible=False)


                with gr.Column(scale=1):
                    with gr.Row(elem_classes="result-header"):
                        gr.Markdown("📸 Segmentation Result", elem_classes="frame-heading")
                        download_mask_btn = gr.DownloadButton(
                            "⬇️ Download (.tif)",
                            size="sm",
                            scale=0,
                            elem_classes="result-download",
                        )
                    with gr.Column(elem_classes="section-card"):
                        seg_output = gr.Image(
                            type="pil",
                            show_label=False,
                            elem_classes="uniform-height"
                        )
                        seg_alpha_slider = gr.Slider(
                            minimum=0.0,
                            maximum=1.0,
                            step=0.05,
                            value=0.5,
                            label="Overlay opacity",
                            elem_classes="opacity-row",
                        )
                        seg_info_html = gr.HTML(visible=False, elem_classes="img-info")

                        # Feedback: rating + comment + submit on one row
                        gr.Markdown("How was this result?", elem_classes="feedback-title")
                        with gr.Row(elem_classes="feedback-row"):
                            score_slider = gr.Radio(choices=[("5", 5), ("4", 4), ("3", 3), ("2", 2), ("1", 1)], value=5, show_label=False, container=False, scale=0, elem_classes="star-rating")
                            feedback_box = gr.Textbox(placeholder="Tell us about your experience (optional)", lines=1, max_lines=1, show_label=False, container=False, scale=2)
                            submit_feedback_btn = gr.Button("Submit feedback", variant="secondary", scale=0)

                        feedback_status = gr.Textbox(
                            label="✅ Submission Status",
                            lines=1,
                            visible=False
                        )
            
            annotator.upload(
                fn=prepare_uploaded_image,
                inputs=annotator,
                outputs=[annotator, seg_raw_state, seg_frame_slider, seg_zplane_slider]
            )
            for _slider in (seg_frame_slider, seg_zplane_slider):
                _slider.release(
                    fn=select_frame,
                    inputs=[seg_raw_state, seg_frame_slider, seg_zplane_slider],
                    outputs=annotator
                )

            # Image information line. Driven off the raw-path state so it covers
            # every way an image arrives (upload, gallery, clear); the frame
            # slider refreshes the "Frames: n / total" counter.
            seg_raw_state.change(
                fn=describe_image_info,
                inputs=[seg_raw_state, seg_frame_slider],
                outputs=seg_info_html
            )
            seg_frame_slider.release(
                fn=describe_image_info,
                inputs=[seg_raw_state, seg_frame_slider],
                outputs=seg_info_html
            )

            # click event for segmentation
            run_seg_btn.click(
                fn=segment_with_choice,
                inputs=[use_box_radio, annotator, seg_alpha_slider],
                outputs=[seg_output, download_mask_btn, seg_vis_state]
            )
            seg_alpha_slider.input(
                fn=update_seg_overlay_alpha,
                inputs=[seg_alpha_slider, seg_vis_state],
                outputs=seg_output
            )

            # click event for clear button
            clear_btn.click(
                fn=lambda: (None, {}, None, gr.update(visible=False, value=1),
                            gr.update(visible=False, value=1)),
                inputs=None,
                outputs=[annotator, seg_vis_state, seg_raw_state, seg_frame_slider, seg_zplane_slider]
            )
            
            # init Gallery with example images
            demo.load(
                fn=lambda: example_images_seg.copy(),
                outputs=example_gallery
            )
            
            # Add the uploaded image to the gallery on button click, confirm with
            # a status message, and clear the uploader so the upload box returns.
            def add_to_gallery(img_path, current_imgs):
                if not img_path:
                    return (current_imgs, current_imgs,
                            gr.update(value="⚠️ Please upload an image first.", visible=True), None)
                rejected = check_image(img_path)
                if rejected:
                    return (current_imgs, current_imgs,
                            gr.update(value=f"❌ {rejected}", visible=True), None)
                if len(current_imgs) - len(example_images_seg) >= MAX_GALLERY_UPLOADS:
                    return (current_imgs, current_imgs,
                            gr.update(value=f"⚠️ Gallery upload limit reached "
                                            f"({MAX_GALLERY_UPLOADS} images). Reload the page "
                                            f"to start over.", visible=True), None)
                std_path = standardize_image(img_path)
                _remember_gallery_source(std_path, img_path)
                if std_path not in current_imgs:
                    current_imgs.insert(0, std_path)
                return (current_imgs, current_imgs,
                        gr.update(value="✅ Added to the Example Gallery below.", visible=True), None)

            add_to_gallery_btn.click(
                fn=add_to_gallery,
                inputs=[image_uploader, user_uploaded_examples],
                outputs=[user_uploaded_examples, example_gallery, add_gallery_status, image_uploader]
            )
            
            # Reuse the upload path so the toast and frame slider behave the same.
            def load_from_gallery(evt: gr.SelectData, all_imgs):
                if evt.index is not None and evt.index < len(all_imgs):
                    item = all_imgs[evt.index]
                    return prepare_uploaded_image(_GALLERY_RAW_SOURCE.get(item, item))
                return None, None, gr.update(visible=False, value=1), gr.update(visible=False, value=1)

            example_gallery.select(
                fn=load_from_gallery,
                inputs=user_uploaded_examples,
                outputs=[annotator, seg_raw_state, seg_frame_slider, seg_zplane_slider]
            )
            
            # click event for submitting feedback
            def submit_user_feedback(query_id, score, comment, annot_val):
                try:
                    img_path = annot_val[0] if annot_val and len(annot_val) > 0 else None
                    bboxes = annot_val[1] if annot_val and len(annot_val) > 1 else []
                    
                    # save_feedback(
                    #     query_id=query_id,
                    #     feedback_type=f"score_{int(score)}",
                    #     feedback_text=comment,
                    #     img_path=img_path,
                    #     bboxes=bboxes
                    # )

                    save_feedback_to_hf(
                        query_id=query_id,
                        feedback_type=f"score_{int(score)}",
                        feedback_text=comment,
                        img_path=img_path,
                        bboxes=bboxes
                    )
                    return "✅ Feedback submitted, thank you!", gr.update(visible=True)
                except Exception as e:
                    return f"❌ Submission failed: {str(e)}", gr.update(visible=True)

            submit_feedback_btn.click(
                fn=submit_user_feedback,
                inputs=[current_query_id, score_slider, feedback_box, annotator],
                outputs=[feedback_status, feedback_status]
            )
        
        # ===== Tab 2: Counting =====
        with gr.Tab("Count") as count_tab:
            # gr.Markdown("## Microscopy Object Counting Analysis")
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("🖼️ Upload Image (Optional: Provide a Bounding Box)", elem_classes="frame-heading")
                    with gr.Column(elem_classes="section-card"):
                        count_annotator = BBoxAnnotator(
                            show_label=False,
                            categories=["cell"],
                            elem_classes="uniform-height-input",
                        )
                        count_frame_slider = gr.Slider(
                            minimum=1, maximum=1, step=1, value=1,
                            label="📚 Stack frame (choose frame to use)",
                            visible=False,
                        )
                        count_zplane_slider = gr.Slider(
                            minimum=1, maximum=1, step=1, value=1,
                            label="🔬 Z-plane (choose plane to use)",
                            visible=False,
                        )

                    # Example gallery with "add" functionality
                    gr.Markdown("📁 Example Image Gallery", elem_classes="frame-heading")
                    with gr.Column(elem_classes="section-card"):
                        count_example_gallery = gr.Gallery(
                            show_label=False,
                            columns=len(example_images_cnt),
                            rows=1,
                            object_fit="cover",
                            height=120,
                            value=example_images_cnt.copy(),  # Initialize with examples
                            show_download_button=False
                        )

                    with gr.Column(elem_classes="bbox-slot"):
                        count_use_box_radio = gr.Radio(
                            choices=["Yes", "No"],
                            value="No",
                            label="🔲 Specify Bounding Box?",
                            elem_classes="inline-radio",
                        )

                    with gr.Row():
                        count_btn = gr.Button("▶️ Run Counting", variant="primary", size="lg")
                        clear_btn = gr.Button("🔄 Clear Selection", variant="secondary")

                    with gr.Accordion("📋 Image Upload Requirements", open=True, elem_classes="req-panel"):
                        gr.HTML(
                            """
                            <div class="req-warn">⚠️ Do not upload images or data containing protected health information (PHI) or personally identifiable information (PII).</div>
                            <div class="req-row"><span class="req-label">Formats</span><span class="req-val">TIFF / OME-TIFF (8/16/32-bit, stacks, multi-channel), PNG, JPG, and other formats supported by <a href="https://github.com/cgohlke/tifffile" target="_blank">tifffile</a></span></div>
                            <div class="req-row"><span class="req-label">Max size</span><span class="req-val">4096 × 4096 px</span></div>
                            <div class="req-row"><span class="req-label">Channels</span><span class="req-val">RGB (3-channel) or grayscale; channels beyond the first 3 are ignored</span></div>
                            <div class="req-row"><span class="req-label">Stacks</span><span class="req-val">first frame loaded by default; frame slider appears after upload</span></div>
                            <div class="req-note">Alternatively, convert to an 8-bit RGB/grayscale PNG/JPG/TIFF first (e.g. with ImageJ/Fiji).</div>
                            """
                        )
                    
                    # Add button to upload new examples
                    gr.Markdown("➕ Add Example Image to Gallery", elem_classes="frame-heading")
                    with gr.Column(elem_classes="section-card"):
                        count_image_uploader = gr.File(
                            show_label=False,
                            file_types=["image"] + list(TIFF_EXTENSIONS),
                            type="filepath"
                        )
                        add_to_count_gallery_btn = gr.Button("➕ Add to Gallery", variant="secondary", size="sm")
                        count_add_status = gr.Markdown(visible=False)

                
                with gr.Column(scale=1):
                    with gr.Row(elem_classes="result-header"):
                        gr.Markdown("📸 Density Map", elem_classes="frame-heading")
                        download_density_btn = gr.DownloadButton(
                            "⬇️ Download (.npy)",
                            size="sm",
                            scale=0,
                            elem_classes="result-download",
                        )
                    with gr.Column(elem_classes="section-card"):
                        count_output = gr.Image(
                            show_label=False,
                            type="filepath",
                            elem_id="density_map_output"

                        )
                        count_alpha_slider = gr.Slider(
                            minimum=0.0,
                            maximum=1.0,
                            step=0.05,
                            value=0.3,
                            label="Density opacity",
                            elem_classes="opacity-row",
                        )
                        count_info_html = gr.HTML(visible=False, elem_classes="img-info")
                        count_status = gr.Textbox(
                            label="📊 Statistics",
                            lines=2
                        )

                        # Feedback: rating + comment + submit on one row
                        gr.Markdown("How was this result?", elem_classes="feedback-title")
                        with gr.Row(elem_classes="feedback-row"):
                            score_slider = gr.Radio(choices=[("5", 5), ("4", 4), ("3", 3), ("2", 2), ("1", 1)], value=5, show_label=False, container=False, scale=0, elem_classes="star-rating")
                            feedback_box = gr.Textbox(placeholder="Tell us about your experience (optional)", lines=1, max_lines=1, show_label=False, container=False, scale=2)
                            submit_feedback_btn = gr.Button("Submit feedback", variant="secondary", scale=0)

                        feedback_status = gr.Textbox(
                            label="✅ Submission Status",
                            lines=1,
                            visible=False
                        )

            # State for managing gallery images
            count_user_examples = gr.State(example_images_cnt.copy())
            
            # Add the uploaded image to the gallery on button click, confirm with
            # a status message, and clear the uploader afterward.
            def add_to_count_gallery(new_img_file, current_imgs):
                """Add uploaded image to gallery"""
                if new_img_file is None:
                    return (current_imgs, current_imgs,
                            gr.update(value="⚠️ Please upload an image first.", visible=True), None)
                rejected = check_image(new_img_file)
                if rejected:
                    return (current_imgs, current_imgs,
                            gr.update(value=f"❌ {rejected}", visible=True), None)
                if len(current_imgs) - len(example_images_cnt) >= MAX_GALLERY_UPLOADS:
                    return (current_imgs, current_imgs,
                            gr.update(value=f"⚠️ Gallery upload limit reached "
                                            f"({MAX_GALLERY_UPLOADS} images). Reload the page "
                                            f"to start over.", visible=True), None)
                std_path = standardize_image(new_img_file)
                _remember_gallery_source(std_path, new_img_file)
                if std_path not in current_imgs:
                    current_imgs.insert(0, std_path)
                    print(f"✅ Added image to gallery: {std_path}")
                return (current_imgs, current_imgs,
                        gr.update(value="✅ Added to the Example Gallery above.", visible=True), None)

            # When user clicks "Add to Gallery"
            add_to_count_gallery_btn.click(
                fn=add_to_count_gallery,
                inputs=[count_image_uploader, count_user_examples],
                outputs=[count_user_examples, count_example_gallery, count_add_status, count_image_uploader]
            )
            
            count_annotator.upload(
                fn=prepare_uploaded_image,
                inputs=count_annotator,
                outputs=[count_annotator, count_raw_state, count_frame_slider, count_zplane_slider]
            )
            for _slider in (count_frame_slider, count_zplane_slider):
                _slider.release(
                    fn=select_frame,
                    inputs=[count_raw_state, count_frame_slider, count_zplane_slider],
                    outputs=count_annotator
                )

            # Image information line (same pattern as the Segment tab).
            count_raw_state.change(
                fn=describe_image_info,
                inputs=[count_raw_state, count_frame_slider],
                outputs=count_info_html
            )
            count_frame_slider.release(
                fn=describe_image_info,
                inputs=[count_raw_state, count_frame_slider],
                outputs=count_info_html
            )

            # When user selects from gallery, load into annotator
            def load_from_count_gallery(evt: gr.SelectData, all_imgs):
                """Load a gallery image, reusing the upload path."""
                if evt.index is not None and evt.index < len(all_imgs):
                    selected_img = all_imgs[evt.index]
                    print(f"📸 Loading image from gallery: {selected_img}")
                    return prepare_uploaded_image(_GALLERY_RAW_SOURCE.get(selected_img, selected_img))
                return None, None, gr.update(visible=False, value=1), gr.update(visible=False, value=1)

            count_example_gallery.select(
                fn=load_from_count_gallery,
                inputs=count_user_examples,
                outputs=[count_annotator, count_raw_state, count_frame_slider, count_zplane_slider]
            )

            # Run counting
            count_btn.click(
                fn=count_cells_handler,
                inputs=[count_use_box_radio, count_annotator, count_alpha_slider],
                outputs=[count_output, download_density_btn, count_status, count_vis_state]
            )
            count_alpha_slider.input(
                fn=update_count_overlay_alpha,
                inputs=[count_alpha_slider, count_vis_state],
                outputs=count_output
            )

            # Clear selection
            clear_btn.click(
                fn=lambda: (None, {}, None, gr.update(visible=False, value=1),
                            gr.update(visible=False, value=1)),
                inputs=None,
                outputs=[count_annotator, count_vis_state, count_raw_state, count_frame_slider, count_zplane_slider]
            )

            # Submit feedback
            def submit_user_feedback(query_id, score, comment, annot_val):
                try:
                    img_path = annot_val[0] if annot_val and len(annot_val) > 0 else None
                    bboxes = annot_val[1] if annot_val and len(annot_val) > 1 else []


                    save_feedback_to_hf(
                        query_id=query_id,
                        feedback_type=f"score_{int(score)}",
                        feedback_text=comment,
                        img_path=img_path,
                        bboxes=bboxes
                    )
                    return "✅ Feedback submitted successfully, thank you!", gr.update(visible=True)
                except Exception as e:
                    return f"❌ Submission failed: {str(e)}", gr.update(visible=True)

            submit_feedback_btn.click(
                fn=submit_user_feedback,
                inputs=[current_query_id, score_slider, feedback_box, annotator],
                outputs=[feedback_status, feedback_status]
            )
        
        # ===== Tab 3: Tracking =====
        with gr.Tab("Track") as track_tab:
            # gr.Markdown("## Microscopy Object Video Tracking - Supports ZIP Upload")
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("📦 Upload Image Sequence in ZIP File (Optional: Provide a Bounding Box)", elem_classes="frame-heading")
                    with gr.Column(elem_classes="section-card"):
                        track_zip_upload = gr.File(
                            show_label=False,
                            file_types=[".zip"]
                        )

                    # First frame annotation for bounding box (heading + annotator
                    # share one column so they show/hide together)
                    with gr.Column(visible=False) as track_annot_group:
                        gr.Markdown("🖼️ (Optional) First Frame Bounding Box Annotation", elem_classes="frame-heading")
                        with gr.Column(elem_classes="section-card"):
                            track_first_frame_annotator = BBoxAnnotator(
                                show_label=False,
                                categories=["cell"],
                                elem_classes="uniform-height-input",
                            )

                    # Example ZIP gallery
                    gr.Markdown("📁 Example Video Gallery (Click to Select)", elem_classes="frame-heading")
                    with gr.Column(elem_classes="section-card"):
                        track_example_gallery = gr.Gallery(
                            show_label=False,
                            columns=10,
                            rows=1,
                            height=120,
                            object_fit="contain",
                            show_download_button=False
                        )

                    with gr.Column(elem_classes="bbox-slot"):
                        track_use_box_radio = gr.Radio(
                            choices=["Yes", "No"],
                            value="No",
                            label="🔲 Specify Bounding Box?",
                            elem_classes="inline-radio",
                        )

                    with gr.Row():
                        track_btn = gr.Button("▶️ Run Tracking", variant="primary", size="lg")
                        clear_btn = gr.Button("🔄 Clear Selection", variant="secondary")

                    with gr.Accordion("📋 Upload Requirements", open=True, elem_classes="req-panel"):
                        gr.HTML(
                            """
                            <div class="req-warn">⚠️ Do not upload images or data containing protected health information (PHI) or personally identifiable information (PII).</div>
                            <div class="req-row"><span class="req-label">Format</span><span class="req-val">A <code>.zip</code> containing a sequence of <code>.tif</code> / <code>.tiff</code> frames</span></div>
                            <div class="req-row"><span class="req-label">Frame order</span><span class="req-val">Sorted naturally by filename; use consistent names such as frame_001.tif, frame_002.tif, …</span></div>
                            <div class="req-row"><span class="req-label">Image size</span><span class="req-val">Every frame must share the same width and height</span></div>
                            <div class="req-row"><span class="req-label">Channels</span><span class="req-val">Grayscale and RGB TIFF frames; beyond 3 channels only the first three are used</span></div>
                            <div class="req-row"><span class="req-label">Bit depth</span><span class="req-val">8-bit and 16-bit TIFFs</span></div>
                            <div class="req-note">You can use ImageJ/Fiji or a similar tool to export your video as a numbered TIFF image sequence, then compress the frames into a ZIP before uploading.</div>
                            """
                        )

                    # Add to gallery button
                    gr.Markdown("➕ Add ZIP to Example Gallery", elem_classes="frame-heading")
                    with gr.Column(elem_classes="section-card"):
                        track_gallery_upload = gr.File(
                            show_label=False,
                            file_types=[".zip"],
                            type="filepath"
                        )
                
                with gr.Column(scale=1):
                    with gr.Row(elem_classes="result-header"):
                        gr.Markdown("📸 Tracking Visualization", elem_classes="frame-heading")
                        # keep the wrapper so the existing show/hide wiring still
                        # targets a container, with the button inside it
                        with gr.Column(visible=False, scale=0, min_width=170) as track_dl_group:
                            track_download = gr.DownloadButton(
                                "⬇️ Download (.zip)",
                                size="sm",
                                elem_classes="result-download",
                            )
                    with gr.Column(elem_classes="section-card"):
                        track_first_frame_preview = gr.Image(
                            show_label=False,
                            type="filepath",
                            # height=400,
                            elem_classes="uniform-height",
                            interactive=False
                        )
                        track_alpha_slider = gr.Slider(
                            minimum=0.0,
                            maximum=1.0,
                            step=0.05,
                            value=0.3,
                            label="Overlay opacity",
                            elem_classes="opacity-row",
                        )

                        track_output = gr.Textbox(
                            label="📊 Tracking Information",
                            lines=9,
                            max_lines=20,
                            interactive=False
                        )

                        # Feedback: rating + comment + submit on one row
                        gr.Markdown("How was this result?", elem_classes="feedback-title")
                        with gr.Row(elem_classes="feedback-row"):
                            score_slider = gr.Radio(choices=[("5", 5), ("4", 4), ("3", 3), ("2", 2), ("1", 1)], value=5, show_label=False, container=False, scale=0, elem_classes="star-rating")
                            feedback_box = gr.Textbox(placeholder="Tell us about your experience (optional)", lines=1, max_lines=1, show_label=False, container=False, scale=2)
                            submit_feedback_btn = gr.Button("Submit feedback", variant="secondary", scale=0)

                        feedback_status = gr.Textbox(
                            label="✅ Submission Status",
                            lines=1,
                            visible=False
                        )
            
            # State for tracking examples
            track_user_examples = gr.State(example_tracking_zips.copy())
            
            # Function to get preview image from ZIP
            def get_zip_preview(zip_path):
                """Extract first frame from ZIP for gallery preview"""
                try:
                    temp_dir = tempfile.mkdtemp()
                    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                        for member in zip_ref.namelist():
                            basename = os.path.basename(member)
                            if ('__MACOSX' not in member and 
                                not basename.startswith('._') and
                                basename.lower().endswith(('.tif', '.tiff', '.png', '.jpg'))):
                                zip_ref.extract(member, temp_dir)
                                extracted_path = os.path.join(temp_dir, member)
                                
                                # Load and normalize for preview
                                import tifffile
                                import numpy as np
                                
                                img_np = tifffile.imread(extracted_path)
                                if img_np.dtype == np.uint16:
                                    img_min, img_max = img_np.min(), img_np.max()
                                    if img_max > img_min:
                                        img_np = ((img_np.astype(np.float32) - img_min) / (img_max - img_min) * 255).astype(np.uint8)
                                
                                if img_np.ndim == 2:
                                    img_np = np.stack([img_np]*3, axis=-1)
                                
                                # Save preview
                                preview_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
                                Image.fromarray(img_np).save(preview_path.name)
                                return preview_path.name
                except:
                    pass
                return None
            
            # Initialize gallery with previews
            def init_tracking_gallery():
                """Create preview images for ZIP examples"""
                previews = []
                for zip_path in example_tracking_zips:
                    if os.path.exists(zip_path):
                        preview = get_zip_preview(zip_path)
                        if preview:
                            previews.append(preview)
                return previews
            
            # Load gallery on startup
            demo.load(
                fn=init_tracking_gallery,
                outputs=track_example_gallery
            )
            
            # Add ZIP to gallery
            def add_zip_to_gallery(zip_path, current_zips):
                if not zip_path:
                    return current_zips, track_example_gallery
                try:
                    if zip_path not in current_zips:
                        current_zips.append(zip_path)
                        print(f"✅ Added ZIP to gallery: {zip_path}")
                    # Regenerate previews
                    previews = []
                    for zp in current_zips:
                        preview = get_zip_preview(zp)
                        if preview:
                            previews.append(preview)
                    return current_zips, previews
                except Exception as e:
                    print(f"⚠️ Error: {e}")
                    return current_zips, []
            
            track_gallery_upload.upload(
                fn=add_zip_to_gallery,
                inputs=[track_gallery_upload, track_user_examples],
                outputs=[track_user_examples, track_example_gallery]
            )
            
            # Select ZIP from gallery
            def load_zip_from_gallery(evt: gr.SelectData, all_zips):
                if evt.index is not None and evt.index < len(all_zips):
                    selected_zip = all_zips[evt.index]
                    print(f"📁 Selected ZIP from gallery: {selected_zip}")
                    return selected_zip
                return None
            
            track_example_gallery.select(
                fn=load_zip_from_gallery,
                inputs=track_user_examples,
                outputs=track_zip_upload
            )

            # Load first frame when ZIP is uploaded
            def load_first_frame_for_annotation(zip_file_obj):
                '''Load and normalize first frame from ZIP for annotation'''
                if zip_file_obj is None:
                    return None, gr.update(visible=False)
                
                import tifffile
                import numpy as np
                
                try:
                    temp_dir = tempfile.mkdtemp()
                    with zipfile.ZipFile(zip_file_obj.name, 'r') as zip_ref:
                        for member in zip_ref.namelist():
                            basename = os.path.basename(member)
                            if ('__MACOSX' not in member and 
                                not basename.startswith('._') and
                                basename.lower().endswith(('.tif', '.tiff'))):
                                zip_ref.extract(member, temp_dir)
                    
                    tif_dir = find_valid_tif_dir(temp_dir)
                    if tif_dir:
                        first_frame = extract_first_frame(tif_dir)
                        if first_frame:
                            # Load and normalize the first frame
                            try:
                                img_np = tifffile.imread(first_frame)
                                
                                # Normalize to [0, 255] uint8 range for display
                                if img_np.dtype == np.uint8:
                                    pass  # Already uint8
                                elif img_np.dtype == np.uint16:
                                    # Normalize uint16 using actual min/max
                                    img_min, img_max = img_np.min(), img_np.max()
                                    if img_max > img_min:
                                        img_np = ((img_np.astype(np.float32) - img_min) / (img_max - img_min) * 255).astype(np.uint8)
                                    else:
                                        img_np = (img_np.astype(np.float32) / 65535.0 * 255).astype(np.uint8)
                                else:
                                    # Float or other types
                                    img_np = img_np.astype(np.float32)
                                    img_min, img_max = img_np.min(), img_np.max()
                                    if img_max > img_min:
                                        img_np = ((img_np - img_min) / (img_max - img_min) * 255).astype(np.uint8)
                                    else:
                                        img_np = np.clip(img_np * 255, 0, 255).astype(np.uint8)
                                
                                # Convert to RGB if grayscale
                                if img_np.ndim == 2:
                                    img_np = np.stack([img_np]*3, axis=-1)
                                elif img_np.ndim == 3 and img_np.shape[2] > 3:
                                    img_np = img_np[:, :, :3]
                                
                                # Save normalized image to temp file
                                temp_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
                                Image.fromarray(img_np).save(temp_img.name)
                                
                                print(f"✅ Loaded and normalized first frame: {first_frame}")
                                print(f"   Original dtype: {tifffile.imread(first_frame).dtype}")
                                print(f"   Normalized to uint8 RGB for annotation")
                                
                                return temp_img.name, gr.update(visible=True)
                            except Exception as e:
                                print(f"⚠️  Error normalizing first frame: {e}")
                                import traceback
                                traceback.print_exc()
                                # Fallback to original file
                                return first_frame, gr.update(visible=True)
                except Exception as e:
                    print(f"⚠️  Error loading first frame: {e}")
                    import traceback
                    traceback.print_exc()
                return None, gr.update(visible=False)
            
            # Load first frame when ZIP is uploaded
            track_zip_upload.change(
                fn=load_first_frame_for_annotation,
                inputs=track_zip_upload,
                outputs=[track_first_frame_annotator, track_annot_group]
            )
            
            # Run tracking
            track_btn.click(
                fn=track_video_handler,
                inputs=[track_use_box_radio, track_first_frame_annotator, track_zip_upload, track_alpha_slider, track_vis_state],
                outputs=[track_download, track_output, track_dl_group, track_first_frame_preview, track_vis_state]
            )
            track_alpha_slider.change(
                fn=update_tracking_overlay_alpha,
                inputs=[track_alpha_slider, track_vis_state],
                outputs=track_first_frame_preview
            )

            # Clear selection
            clear_btn.click(
                fn=lambda: (None, {}),
                inputs=None,
                outputs=[track_first_frame_annotator, track_vis_state]
            )

            # Submit feedback
            def submit_user_feedback(query_id, score, comment, annot_val):
                try:
                    img_path = annot_val[0] if annot_val and len(annot_val) > 0 else None
                    bboxes = annot_val[1] if annot_val and len(annot_val) > 1 else []

                    save_feedback_to_hf(
                        query_id=query_id,
                        feedback_type=f"score_{int(score)}",
                        feedback_text=comment,
                        img_path=img_path,
                        bboxes=bboxes
                    )
                    return "✅ Feedback submitted successfully, thank you!", gr.update(visible=True)
                except Exception as e:
                    return f"❌ Submission failed: {str(e)}", gr.update(visible=True)

            submit_feedback_btn.click(
                fn=submit_user_feedback,
                inputs=[current_query_id, score_slider, feedback_box, annotator],
                outputs=[feedback_status, feedback_status]
            )

    seg_tab.select(lambda: (HOWTO_TITLE_SEG, HOWTO_SEG), outputs=[how_to_use_title, howto_md])
    count_tab.select(lambda: (HOWTO_TITLE_COUNT, HOWTO_COUNT), outputs=[how_to_use_title, howto_md])
    track_tab.select(lambda: (HOWTO_TITLE_TRACK, HOWTO_TRACK), outputs=[how_to_use_title, howto_md])

    gr.Markdown(
        """
        **Supporting three key tasks:**

        🎨 **Segmentation**: Instance segmentation of microscopic objects
         
        🔢 **Counting**: Counting microscopic objects based on density maps 
        
        🎬 **Tracking**: Tracking microscopic objects in video sequences

        📒 **Note:** This project is currently available with usage limits for research trial use and feedback collection. Response speed may vary depending on GPU allocation queues. We plan to release a free public version in the future. We are actively improving the toolkit and greatly appreciate your feedback!
        """,
        elem_id="footer_note"
    )


if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        ssr_mode=False,
        show_error=True,
    )