Blame for src/streamlit_app.py

# src/streamlit_app.py
"""
MyoSeg  —  Myotube & Nuclei Analyzer
========================================
Drop-in replacement for streamlit_app.py on Hugging Face Spaces.

Features:
  ✦ Animated count-up metrics (9 counters)
  ✦ Instance overlay — nucleus IDs (1,2,3…) + myotube IDs (M1,M2…)
  ✦ Separate nuclei outline + myotube outline tabs
  ✦ Watershed nuclei splitting for accurate counts
  ✦ Myotube surface area (total, mean, max µm²) + per-tube bar chart
  ✦ Active learning — upload corrected masks → saved to corrections/
  ✦ Low-confidence auto-flagging → image queued for retraining
  ✦ Retraining queue status panel
  ✦ All original sidebar controls preserved

  v10 changes (core refactor):
  ✦ All analysis code moved to myoseg_core.py (no Streamlit) so the same
    functions can back a chat assistant, an MCP server / API and scripts.
  ✦ U-Net output cached per image: changing thresholds, postprocessing or
    FI settings and pressing Rerun skips the network.
  ✦ Cytoplasm-hole FI, aspect filter and myotube split cropped to each
    object's bounding box — identical results, ~36x faster postprocessing.
  ✦ Review mode no longer recomputes the nuclei watershed twice.
  ✦ Image-quality checks + rule-based "suggested parameter changes" panels.
  ✦ Optional "Use split myotubes for counts" (off = v9 behaviour).
  ✦ A failing image is reported under "Images failed" instead of stopping the batch.
  ✦ Low-confidence images are queued once, not again on every rerun.

  v9 changes:
  ✦ FIXED: SVG viewer myotube ID count now matches live metrics count.
    Root cause: viewer showed all connected components (myo_lab), but
    metrics only counted those with ≥1 MyHC+ nucleus.  Now the viewer
    badge shows the biological myotube_count from compute_bio_metrics,
    and non-bio myotube regions are shown as faint outlines (not labelled).
  ✦ Outlines split into two separate tabs: "Nuclei outlines" and
    "Myotube outlines" per collaborator request.
  ✦ Privacy mode: sidebar toggle for "Private mode — do not use my data
    for training".  When enabled, images are NOT queued for retraining
    (no low_confidence queue, no corrections submission).
  ✦ Training contribution mode: explicit user-initiated action to submit
    current image + tuned parameters as a training contribution.  Only
    runs when user clicks "Submit for training" after finding good params.
  ✦ Parameter learning: when user submits, the current sidebar parameter
    set (thresholds, postprocessing knobs) is saved alongside the image
    so self_train.py can learn optimal parameters per image type.
  ✦ All v8 fixes preserved (no closing, shape filter, erode+dilate, etc).
"""

import io
import os
import json
import time
import zipfile
import hashlib
import re
from datetime import datetime
from pathlib import Path

import numpy as np
import pandas as pd
from PIL import Image

import streamlit as st
import streamlit.components.v1
import torch
import torch.nn as nn
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from huggingface_hub import hf_hub_download

import scipy.ndimage as ndi
from skimage.morphology import remove_small_objects, disk
from skimage import measure
from skimage.segmentation import watershed, find_boundaries
from skimage.feature import peak_local_max

import sys as _sys
_sys.path.insert(0, str(Path(__file__).resolve().parent))
import myoseg_core as core
from myoseg_core import (                      # analysis engine (no Streamlit inside)
    _safe_binary_dilation, sha256_file, get_channel, resize_u8_to_float01,
    postprocess_masks, label_cc, split_large_myotubes, label_nuclei_watershed,
    compute_surface_area, classify_nucleus_in_myotube, compute_bio_metrics,
    DoubleConv, UNet, PRESETS,
)



# ─────────────────────────────────────────────────────────────────────────────
# CONFIG  ← edit these two lines to match your HF model repo
# ─────────────────────────────────────────────────────────────────────────────
MODEL_REPO_ID  = "skarugu/myotube-unet"
MODEL_FILENAME = "model_final.pt"
APP_VERSION    = "v12-final-2026-10-21"

CONF_FLAG_THR   = 0.60    # legacy: max-probability confidence (almost never fires)
# Images whose uncertainty (0 = sure … 1 = coin-flip, see core.uncertainty_score) exceeds
# this are queued for labelling. Calibrate with benchmark.py on your own labelled data.
UNC_FLAG_THR    = float(os.environ.get("UNCERTAINTY_FLAG_THR", 0.30))
QUEUE_DIR       = Path("retrain_queue")
CORRECTIONS_DIR = Path("corrections")


# ─────────────────────────────────────────────────────────────────────────────
# Helpers  (identical to originals so nothing breaks)
# ─────────────────────────────────────────────────────────────────────────────

def png_bytes(arr_u8: np.ndarray) -> bytes:
    # compress_level=1: lossless (identical pixels), ~5x faster than the default 6,
    # files ~15% larger. PNG encoding was the slowest step of Review mode.
    buf = io.BytesIO()
    Image.fromarray(arr_u8).save(buf, format="PNG", compress_level=1)
    return buf.getvalue()


def hex_to_rgb(h: str):
    h = h.lstrip("#")
    return tuple(int(h[i:i+2], 16) for i in (0, 2, 4))


# ─────────────────────────────────────────────────────────────────────────────
# Overlay helpers
# ─────────────────────────────────────────────────────────────────────────────

def make_simple_overlay(rgb_u8, nuc_mask, myo_mask, nuc_color, myo_color, alpha):
    """Flat colour overlay — used for the ZIP export (fast, no matplotlib)."""
    base  = rgb_u8.astype(np.float32)
    H0, W0 = rgb_u8.shape[:2]
    nuc   = np.array(Image.fromarray((nuc_mask*255).astype(np.uint8))
                     .resize((W0, H0), Image.NEAREST)) > 0
    myo   = np.array(Image.fromarray((myo_mask*255).astype(np.uint8))
                     .resize((W0, H0), Image.NEAREST)) > 0
    out   = base.copy()
    for mask, color in [(myo, myo_color), (nuc, nuc_color)]:
        c = np.array(color, dtype=np.float32)
        out[mask] = (1 - alpha) * out[mask] + alpha * c
    return np.clip(out, 0, 255).astype(np.uint8)


def make_coloured_overlay(rgb_u8: np.ndarray,
                           nuc_lab: np.ndarray,
                           myo_lab: np.ndarray,
                           alpha: float = 0.45,
                           nuc_color: tuple = None,
                           myo_color: tuple = None) -> np.ndarray:
    """
    Colour the mask regions only — NO text baked in.
    Returns an RGB uint8 array at original image resolution.

    nuc_color / myo_color: RGB tuple e.g. (0, 255, 255).
      If None, uses per-instance colourmaps (cool / autumn).
      If provided, uses a flat solid colour for all instances of that type —
      this is what the sidebar colour pickers control.
    """
    orig_h, orig_w = rgb_u8.shape[:2]
    nuc_cmap = plt.colormaps["cool"]
    myo_cmap = plt.colormaps["autumn"]

    def _resize_lab(lab, h, w):
        return np.array(
            Image.fromarray(lab.astype(np.int32)).resize((w, h), Image.NEAREST)
        )

    nuc_disp = _resize_lab(nuc_lab, orig_h, orig_w)
    myo_disp = _resize_lab(myo_lab, orig_h, orig_w)
    n_nuc    = int(nuc_disp.max())
    n_myo    = int(myo_disp.max())

    base = rgb_u8.astype(np.float32).copy()
    if n_myo > 0:
        mask = myo_disp > 0
        if myo_color is not None:
            colour_layer = np.array(myo_color, dtype=np.float32)
            base[mask] = (1 - alpha) * base[mask] + alpha * colour_layer
        else:
            myo_norm = (myo_disp / max(n_myo, 1)).astype(np.float32)
            myo_rgba = (myo_cmap(myo_norm)[:, :, :3] * 255).astype(np.float32)
            base[mask] = (1 - alpha) * base[mask] + alpha * myo_rgba[mask]

    if n_nuc > 0:
        mask = nuc_disp > 0
        if nuc_color is not None:
            colour_layer = np.array(nuc_color, dtype=np.float32)
            base[mask] = (1 - alpha) * base[mask] + alpha * colour_layer
        else:
            nuc_norm = (nuc_disp / max(n_nuc, 1)).astype(np.float32)
            nuc_rgba = (nuc_cmap(nuc_norm)[:, :, :3] * 255).astype(np.float32)
            base[mask] = (1 - alpha) * base[mask] + alpha * nuc_rgba[mask]

    return np.clip(base, 0, 255).astype(np.uint8)


def make_outline_overlay(rgb_u8: np.ndarray,
                          nuc_lab: np.ndarray,
                          myo_lab: np.ndarray,
                          nuc_color: tuple = (0, 255, 255),
                          myo_color: tuple = (0, 255, 0),
                          line_width: int = 2) -> np.ndarray:
    """
    Draw contour outlines around each detected instance on the original image.
    Shows exactly what the model considers each myotube/nucleus boundary.
    """
    orig_h, orig_w = rgb_u8.shape[:2]

    def _resize_lab(lab, h, w):
        return np.array(
            Image.fromarray(lab.astype(np.int32)).resize((w, h), Image.NEAREST)
        )

    nuc_disp = _resize_lab(nuc_lab, orig_h, orig_w)
    myo_disp = _resize_lab(myo_lab, orig_h, orig_w)

    out = rgb_u8.copy()

    # Myotube outlines
    if myo_disp.max() > 0:
        myo_bounds = find_boundaries(myo_disp, mode='outer')
        if line_width > 1:
            myo_bounds = _safe_binary_dilation(myo_bounds, disk(line_width - 1))
        out[myo_bounds] = myo_color

    # Nuclei outlines
    if nuc_disp.max() > 0:
        nuc_bounds = find_boundaries(nuc_disp, mode='outer')
        if line_width > 1:
            nuc_bounds = _safe_binary_dilation(nuc_bounds, disk(max(line_width - 2, 0)))
        out[nuc_bounds] = nuc_color

    return out


def collect_label_positions(nuc_lab: np.ndarray,
                             myo_lab: np.ndarray,
                             img_w: int, img_h: int,
                             bio_myo_ids: set = None) -> dict:
    """
    Collect centroid positions for every nucleus and myotube,
    scaled to the original image pixel dimensions.

    bio_myo_ids: set of myotube label IDs that have ≥1 MyHC+ nucleus.
                 If provided, only these are labelled as "M1", "M2", …
                 (renumbered sequentially).  Non-bio regions get no label
                 and are stored separately for faint outline rendering.
    Returns:
        { "nuclei":     [ {"id": "1",  "x": 123.4, "y": 56.7}, ... ],
          "myotubes":   [ {"id": "M1", "x": 200.1, "y": 300.5, "orig_label": 5}, ... ],
          "myotubes_nonbio": [ {"id": "", "x": ..., "y": ..., "orig_label": 3}, ... ] }
    """
    sx = img_w / nuc_lab.shape[1]
    sy = img_h / nuc_lab.shape[0]

    nuclei = []
    for prop in measure.regionprops(nuc_lab):
        r, c = prop.centroid
        nuclei.append({"id": str(prop.label), "x": round(c * sx, 1), "y": round(r * sy, 1)})

    sx2 = img_w / myo_lab.shape[1]
    sy2 = img_h / myo_lab.shape[0]
    myotubes = []
    myotubes_nonbio = []

    if bio_myo_ids is not None and len(bio_myo_ids) > 0:
        # Renumber biological myotubes sequentially: M1, M2, M3…
        sorted_bio = sorted(bio_myo_ids)
        bio_remap = {orig: idx + 1 for idx, orig in enumerate(sorted_bio)}
        for prop in measure.regionprops(myo_lab):
            r, c = prop.centroid
            pos = {"x": round(c * sx2, 1), "y": round(r * sy2, 1), "orig_label": prop.label}
            if prop.label in bio_remap:
                pos["id"] = f"M{bio_remap[prop.label]}"
                myotubes.append(pos)
            else:
                pos["id"] = ""
                myotubes_nonbio.append(pos)
    else:
        # Fallback: label all connected components (backward compat)
        for prop in measure.regionprops(myo_lab):
            r, c = prop.centroid
            myotubes.append({"id": f"M{prop.label}", "x": round(c * sx2, 1),
                             "y": round(r * sy2, 1), "orig_label": prop.label})

    return {"nuclei": nuclei, "myotubes": myotubes, "myotubes_nonbio": myotubes_nonbio}


def make_svg_viewer(img_b64: str,
                    img_w: int, img_h: int,
                    label_data: dict,
                    show_nuclei: bool = True,
                    show_myotubes: bool = True,
                    nuc_font_size: int = 11,
                    myo_font_size: int = 22,
                    viewer_height: int = 620) -> str:
    """
    Build a self-contained HTML string with:
      - A pan-and-zoom SVG viewer (mouse wheel + click-drag)
      - The coloured overlay PNG as the background
      - SVG <text> labels that stay pixel-perfect at any zoom level
      - A font-size slider that updates label sizes live
      - Toggle buttons for nuclei / myotubes labels
      - Count badges in the top-right corner

    Parameters
    ----------
    img_b64       : base64-encoded PNG of the coloured overlay (no text)
    img_w, img_h  : original pixel dimensions of the image
    label_data    : output of collect_label_positions()
    show_nuclei   : initial visibility of nucleus labels
    show_myotubes : initial visibility of myotube labels
    nuc_font_size : initial nucleus label font size (px)
    myo_font_size : initial myotube label font size (px)
    viewer_height : height of the viewer div in pixels
    """
    import json as _json
    labels_json = _json.dumps(label_data)
    n_nuc = len(label_data.get("nuclei", []))
    n_myo = len(label_data.get("myotubes", []))

    show_nuc_js = "true" if show_nuclei  else "false"
    show_myo_js = "true" if show_myotubes else "false"

    html = f"""
<style>
  .myo-viewer-wrap {{
    background: #0e0e1a;
    border: 1px solid #2a2a4e;
    border-radius: 10px;
    overflow: hidden;
    position: relative;
    user-select: none;
  }}
  .myo-toolbar {{
    display: flex;
    align-items: center;
    gap: 12px;
    padding: 8px 14px;
    background: #13132a;
    border-bottom: 1px solid #2a2a4e;
    flex-wrap: wrap;
  }}
  .myo-badge {{
    background: #1a1a3e;
    border: 1px solid #3a3a6e;
    border-radius: 6px;
    padding: 3px 10px;
    color: #e0e0e0;
    font-size: 13px;
    font-family: Arial, sans-serif;
    white-space: nowrap;
  }}
  .myo-badge span {{ font-weight: bold; }}
  .myo-btn {{
    padding: 4px 12px;
    border-radius: 6px;
    border: 1px solid #444;
    cursor: pointer;
    font-size: 12px;
    font-family: Arial, sans-serif;
    font-weight: bold;
    transition: opacity 0.15s;
  }}
  .myo-btn.nuc  {{ background: #003366; color: white; border-color: #4fc3f7; }}
  .myo-btn.myo  {{ background: #8B0000; color: white; border-color: #ff6666; }}
  .myo-btn.off  {{ opacity: 0.35; }}
  .myo-btn.reset {{ background: #1a1a2e; color: #90caf9; border-color: #3a3a6e; }}
  .myo-slider-wrap {{
    display: flex;
    align-items: center;
    gap: 6px;
    color: #aaa;
    font-size: 12px;
    font-family: Arial, sans-serif;
  }}
  .myo-slider-wrap input {{ width: 70px; accent-color: #4fc3f7; cursor: pointer; }}
  .myo-hint {{
    margin-left: auto;
    color: #555;
    font-size: 11px;
    font-family: Arial, sans-serif;
    white-space: nowrap;
  }}
  .myo-svg-wrap {{
    width: 100%;
    height: {viewer_height}px;
    overflow: hidden;
    cursor: grab;
    position: relative;
  }}
  .myo-svg-wrap:active {{ cursor: grabbing; }}
  svg.myo-svg {{
    width: 100%;
    height: 100%;
    display: block;
  }}
</style>

<div class="myo-viewer-wrap" id="myoViewer">
  <div class="myo-toolbar">
    <div class="myo-badge">🔵 Nuclei &nbsp;<span id="nucCount">{n_nuc}</span></div>
    <div class="myo-badge">🔴 Myotubes &nbsp;<span id="myoCount">{n_myo}</span></div>
    <button class="myo-btn nuc" id="btnNuc"  onclick="toggleLayer('nuc')">Nuclei IDs</button>
    <button class="myo-btn myo" id="btnMyo"  onclick="toggleLayer('myo')">Myotube IDs</button>
    <button class="myo-btn reset"             onclick="resetView()">⟳ Reset</button>
    <div class="myo-slider-wrap">
      Nucleus size:
      <input type="range" id="slNuc" min="4" max="40" value="{nuc_font_size}"
             oninput="setFontSize('nuc', this.value)" />
      <span id="lblNuc">{nuc_font_size}px</span>
    </div>
    <div class="myo-slider-wrap">
      Myotube size:
      <input type="range" id="slMyo" min="8" max="60" value="{myo_font_size}"
             oninput="setFontSize('myo', this.value)" />
      <span id="lblMyo">{myo_font_size}px</span>
    </div>
    <div class="myo-hint">Scroll to zoom &nbsp;·&nbsp; Drag to pan</div>
  </div>

  <div class="myo-svg-wrap" id="svgWrap">
    <svg class="myo-svg" id="mainSvg"
         viewBox="0 0 {img_w} {img_h}"
         preserveAspectRatio="xMidYMid meet">
      <defs>
        <filter id="dropshadow" x="-5%" y="-5%" width="110%" height="110%">
          <feDropShadow dx="0" dy="0" stdDeviation="1.5" flood-color="#000" flood-opacity="0.8"/>
        </filter>
      </defs>

      <!-- background image — the coloured overlay PNG -->
      <image href="data:image/png;base64,{img_b64}"
             x="0" y="0" width="{img_w}" height="{img_h}"
             preserveAspectRatio="xMidYMid meet"/>

      <!-- nuclei labels group -->
      <g id="gNuc" visibility="{'visible' if show_nuclei else 'hidden'}">
      </g>

      <!-- myotube labels group -->
      <g id="gMyo" visibility="{'visible' if show_myotubes else 'hidden'}">
      </g>
    </svg>
  </div>
</div>

<script>
(function() {{
  const labels  = {labels_json};
  const IMG_W   = {img_w};
  const IMG_H   = {img_h};

  let nucFontSize = {nuc_font_size};
  let myoFontSize = {myo_font_size};
  let showNuc = {show_nuc_js};
  let showMyo = {show_myo_js};

  // ── Build SVG label elements ─────────────────────────────────────────────
  const NS = "http://www.w3.org/2000/svg";

  function makeLabelGroup(items, fontSize, bgColor, borderColor, isMyo) {{
    const frag = document.createDocumentFragment();
    items.forEach(item => {{
      const g = document.createElementNS(NS, "g");
      g.setAttribute("class", isMyo ? "lbl-myo" : "lbl-nuc");

      // Background rect — sized after text is measured
      const rect = document.createElementNS(NS, "rect");
      rect.setAttribute("rx", isMyo ? "4" : "3");
      rect.setAttribute("ry", isMyo ? "4" : "3");
      rect.setAttribute("fill", bgColor);
      rect.setAttribute("stroke", borderColor);
      rect.setAttribute("stroke-width", isMyo ? "1.5" : "0");
      rect.setAttribute("opacity", isMyo ? "0.93" : "0.90");
      rect.setAttribute("filter", "url(#dropshadow)");

      // Text
      const txt = document.createElementNS(NS, "text");
      txt.textContent = item.id;
      txt.setAttribute("x", item.x);
      txt.setAttribute("y", item.y);
      txt.setAttribute("text-anchor", "middle");
      txt.setAttribute("dominant-baseline", "central");
      txt.setAttribute("fill", "white");
      txt.setAttribute("font-family", "Arial, sans-serif");
      txt.setAttribute("font-weight", "bold");
      txt.setAttribute("font-size", fontSize);
      txt.setAttribute("paint-order", "stroke");

      g.appendChild(rect);
      g.appendChild(txt);
      frag.appendChild(g);
    }});
    return frag;
  }}

  function positionRects() {{
    // After elements are in the DOM, size and position the backing rects
    document.querySelectorAll(".lbl-nuc, .lbl-myo").forEach(g => {{
      const txt  = g.querySelector("text");
      const rect = g.querySelector("rect");
      try {{
        const bb = txt.getBBox();
        const pad = parseFloat(txt.getAttribute("font-size")) * 0.22;
        rect.setAttribute("x",      bb.x - pad);
        rect.setAttribute("y",      bb.y - pad);
        rect.setAttribute("width",  bb.width  + pad * 2);
        rect.setAttribute("height", bb.height + pad * 2);
      }} catch(e) {{}}
    }});
  }}

  function rebuildLabels() {{
    const gNuc = document.getElementById("gNuc");
    const gMyo = document.getElementById("gMyo");
    gNuc.innerHTML = "";
    gMyo.innerHTML = "";
    gNuc.appendChild(makeLabelGroup(labels.nuclei,   nucFontSize, "#003366", "none",    false));
    gMyo.appendChild(makeLabelGroup(labels.myotubes, myoFontSize, "#8B0000", "#FF6666", true));
    // rAF so the browser has laid out the text before we measure it
    requestAnimationFrame(positionRects);
  }}

  // ── Font size controls ────────────────────────────────────────────────────
  window.setFontSize = function(which, val) {{
    val = parseInt(val);
    if (which === "nuc") {{
      nucFontSize = val;
      document.getElementById("lblNuc").textContent = val + "px";
      document.querySelectorAll(".lbl-nuc text").forEach(t => t.setAttribute("font-size", val));
    }} else {{
      myoFontSize = val;
      document.getElementById("lblMyo").textContent = val + "px";
      document.querySelectorAll(".lbl-myo text").forEach(t => t.setAttribute("font-size", val));
    }}
    requestAnimationFrame(positionRects);
  }};

  // ── Layer toggles ─────────────────────────────────────────────────────────
  window.toggleLayer = function(which) {{
    if (which === "nuc") {{
      showNuc = !showNuc;
      document.getElementById("gNuc").setAttribute("visibility", showNuc ? "visible" : "hidden");
      document.getElementById("btnNuc").classList.toggle("off", !showNuc);
    }} else {{
      showMyo = !showMyo;
      document.getElementById("gMyo").setAttribute("visibility", showMyo ? "visible" : "hidden");
      document.getElementById("btnMyo").classList.toggle("off", !showMyo);
    }}
  }};

  // ── Pan + Zoom (pure SVG viewBox manipulation) ────────────────────────────
  const wrap = document.getElementById("svgWrap");
  const svg  = document.getElementById("mainSvg");

  let vx = 0, vy = 0, vw = IMG_W, vh = IMG_H;   // current viewBox

  function setVB() {{
    svg.setAttribute("viewBox", `${{vx}} ${{vy}} ${{vw}} ${{vh}}`);
  }}

  // Scroll to zoom — zoom toward mouse cursor
  wrap.addEventListener("wheel", e => {{
    e.preventDefault();
    const rect  = wrap.getBoundingClientRect();
    const mx    = (e.clientX - rect.left) / rect.width;   // 0..1
    const my    = (e.clientY - rect.top)  / rect.height;
    const factor = e.deltaY < 0 ? 0.85 : 1.0 / 0.85;
    const nw    = Math.min(IMG_W, Math.max(IMG_W * 0.05, vw * factor));
    const nh    = Math.min(IMG_H, Math.max(IMG_H * 0.05, vh * factor));
    vx = vx + mx * (vw - nw);
    vy = vy + my * (vh - nh);
    vw = nw;
    vh = nh;
    // Clamp
    vx = Math.max(0, Math.min(IMG_W - vw, vx));
    vy = Math.max(0, Math.min(IMG_H - vh, vy));
    setVB();
  }}, {{ passive: false }});

  // Drag to pan
  let dragging = false, dragX0, dragY0, vx0, vy0;

  wrap.addEventListener("mousedown", e => {{
    dragging = true;
    dragX0 = e.clientX; dragY0 = e.clientY;
    vx0 = vx; vy0 = vy;
  }});
  window.addEventListener("mousemove", e => {{
    if (!dragging) return;
    const rect = wrap.getBoundingClientRect();
    const scaleX = vw / rect.width;
    const scaleY = vh / rect.height;
    vx = Math.max(0, Math.min(IMG_W - vw, vx0 - (e.clientX - dragX0) * scaleX));
    vy = Math.max(0, Math.min(IMG_H - vh, vy0 - (e.clientY - dragY0) * scaleY));
    setVB();
  }});
  window.addEventListener("mouseup",  () => {{ dragging = false; }});

  // Touch support
  let t0 = null, pinch0 = null;
  wrap.addEventListener("touchstart", e => {{
    if (e.touches.length === 1) {{
      t0 = e.touches[0]; vx0 = vx; vy0 = vy;
    }} else if (e.touches.length === 2) {{
      pinch0 = Math.hypot(
        e.touches[0].clientX - e.touches[1].clientX,
        e.touches[0].clientY - e.touches[1].clientY
      );
    }}
  }}, {{ passive: true }});
  wrap.addEventListener("touchmove", e => {{
    e.preventDefault();
    if (e.touches.length === 1 && t0) {{
      const rect = wrap.getBoundingClientRect();
      vx = Math.max(0, Math.min(IMG_W - vw, vx0 - (e.touches[0].clientX - t0.clientX) * vw / rect.width));
      vy = Math.max(0, Math.min(IMG_H - vh, vy0 - (e.touches[0].clientY - t0.clientY) * vh / rect.height));
      setVB();
    }} else if (e.touches.length === 2 && pinch0 !== null) {{
      const dist = Math.hypot(
        e.touches[0].clientX - e.touches[1].clientX,
        e.touches[0].clientY - e.touches[1].clientY
      );
      const factor = pinch0 / dist;
      const nw = Math.min(IMG_W, Math.max(IMG_W * 0.05, vw * factor));
      const nh = Math.min(IMG_H, Math.max(IMG_H * 0.05, vh * factor));
      vw = nw; vh = nh;
      vx = Math.max(0, Math.min(IMG_W - vw, vx));
      vy = Math.max(0, Math.min(IMG_H - vh, vy));
      pinch0 = dist;
      setVB();
    }}
  }}, {{ passive: false }});

  // ── Reset view ────────────────────────────────────────────────────────────
  window.resetView = function() {{
    vx = 0; vy = 0; vw = IMG_W; vh = IMG_H;
    setVB();
  }};

  // ── Init ──────────────────────────────────────────────────────────────────
  rebuildLabels();
}})();
</script>
"""
    return html


# ─────────────────────────────────────────────────────────────────────────────
# Animated counter
# ─────────────────────────────────────────────────────────────────────────────

def animated_metric(placeholder, label: str, final_val,
                    color: str = "#4fc3f7", steps: int = 20, delay: float = 0.025):
    is_float = isinstance(final_val, float)
    for i in range(1, steps + 1):
        v = final_val * i / steps
        display = f"{v:.1f}" if is_float else str(int(v))
        placeholder.markdown(
            f"""
            <div style='text-align:center;padding:12px 6px;border-radius:12px;
                        background:#1a1a2e;border:1px solid #2a2a4e;margin:4px 0;'>
                <div style='font-size:2rem;font-weight:800;color:{color};
                            line-height:1.1;'>{display}</div>
                <div style='font-size:0.75rem;color:#9e9e9e;margin-top:4px;'>{label}</div>
            </div>
            """,
            unsafe_allow_html=True,
        )
        time.sleep(delay)


# ─────────────────────────────────────────────────────────────────────────────
# Active-learning queue helpers
# ─────────────────────────────────────────────────────────────────────────────

def _ensure_dirs():
    QUEUE_DIR.mkdir(parents=True, exist_ok=True)
    CORRECTIONS_DIR.mkdir(parents=True, exist_ok=True)


def add_to_queue(image_array: np.ndarray, reason: str = "batch",
                 nuc_mask=None, myo_mask=None, metadata: dict = None):
    _ensure_dirs()
    ts   = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
    meta = {**(metadata or {}), "reason": reason, "timestamp": ts}

    if nuc_mask is not None and myo_mask is not None:
        folder = CORRECTIONS_DIR / ts
        folder.mkdir(parents=True, exist_ok=True)
        Image.fromarray(image_array).save(folder / "image.png")
        Image.fromarray((nuc_mask > 0).astype(np.uint8) * 255).save(folder / "nuclei_mask.png")
        Image.fromarray((myo_mask > 0).astype(np.uint8) * 255).save(folder / "myotube_mask.png")
        (folder / "meta.json").write_text(json.dumps({**meta, "has_masks": True}, indent=2))
    else:
        Image.fromarray(image_array).save(QUEUE_DIR / f"{ts}.png")
        (QUEUE_DIR / f"{ts}.json").write_text(json.dumps({**meta, "has_masks": False}, indent=2))


def safe_zip_stem(name: str) -> str:
    keep = []
    for ch in Path(str(name)).stem:
        keep.append(ch if ch.isalnum() or ch in ("-", "_") else "_")
    stem = "".join(keep).strip("_")
    return stem or "image"


def normalize_sample_id(filename: str) -> str:
    """Create a stable, human-readable ID from any uploaded filename."""
    stem = Path(str(filename)).stem
    stem = stem.replace("\u00a0", " ")
    stem = re.sub(r"\s+", " ", stem).strip()
    stem = re.sub(r"[\\/:*?\"<>|]+", "_", stem)
    stem = re.sub(r"_+", "_", stem).strip(" _.-")
    return stem or "image"


def unique_sample_id(filename: str, used_ids: set[str]) -> str:
    base = normalize_sample_id(filename)
    sample_id = base
    n = 2
    while sample_id in used_ids:
        sample_id = f"{base}_{n}"
        n += 1
    used_ids.add(sample_id)
    return sample_id


def sample_match_key(name: str) -> str:
    return normalize_sample_id(name).casefold()


def artifact_lookup_keys(display_name: str, artifact: dict = None) -> set[str]:
    keys = {sample_match_key(display_name), sample_match_key(safe_zip_stem(display_name))}
    if artifact:
        for value in (
            artifact.get("original_filename"),
            artifact.get("original_stem"),
            artifact.get("sample_id"),
        ):
            if value:
                keys.add(sample_match_key(value))
                keys.add(sample_match_key(safe_zip_stem(value)))
    return {k for k in keys if k}


def build_mask_batch_zip(selected_names, artifacts: dict) -> bytes:
    buf = io.BytesIO()
    manifest = []
    with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        for name in selected_names:
            if name not in artifacts:
                continue
            art = artifacts[name]
            folder = safe_zip_stem(name)
            rgb = art.get("rgb_u8")
            nuc = art.get("nuc_pp_arr")
            myo = art.get("myo_pp_arr")
            if rgb is None or nuc is None or myo is None:
                continue
            zf.writestr(f"{folder}/image.png", png_bytes(rgb.astype(np.uint8)))
            zf.writestr(f"{folder}/nuclei_mask.png", png_bytes((nuc > 0).astype(np.uint8) * 255))
            zf.writestr(f"{folder}/myotube_mask.png", png_bytes((myo > 0).astype(np.uint8) * 255))
            meta = {
                "source_image": name,
                "folder": folder,
                "expected_files": ["image.png", "nuclei_mask.png", "myotube_mask.png"],
            }
            zf.writestr(f"{folder}/meta.json", json.dumps(meta, indent=2))
            manifest.append(meta)
        zf.writestr("manifest.json", json.dumps(manifest, indent=2))
    return buf.getvalue()


def _zip_png_mask(zf: zipfile.ZipFile, member: str) -> np.ndarray:
    with zf.open(member) as f:
        return (np.array(Image.open(f).convert("L")) > 0).astype(np.uint8)


def _pick_mask_member(files: dict, base_name: str):
    for ext in (".png", ".tif", ".tiff"):
        member = files.get(f"{base_name}{ext}")
        if member:
            return member
    return None


def import_correction_zip(zip_file, artifacts: dict) -> tuple[int, list[str]]:
    submitted = 0
    errors = []
    artifact_by_key = {}
    for name, art in artifacts.items():
        for key in artifact_lookup_keys(name, art):
            artifact_by_key[key] = name

    try:
        data = zip_file.getvalue()
        zf = zipfile.ZipFile(io.BytesIO(data))
    except Exception as exc:
        return 0, [f"Could not read ZIP: {exc}"]

    members = [m for m in zf.namelist() if not m.endswith("/")]
    used_sources = set()

    # Preferred format: folders produced by build_mask_batch_zip.
    folders = sorted({str(Path(m).parent).replace("\\", "/") for m in members if "/" in m})
    for folder in folders:
        if folder in ("", "."):
            continue
        files = {Path(m).name.lower(): m for m in members if str(Path(m).parent).replace("\\", "/") == folder}
        nuc_member = _pick_mask_member(files, "nuclei_mask")
        myo_member = _pick_mask_member(files, "myotube_mask")
        if not (nuc_member and myo_member):
            continue
        source = artifact_by_key.get(sample_match_key(Path(folder).name))
        meta_member = files.get("meta.json")
        if meta_member:
            try:
                with zf.open(meta_member) as f:
                    meta = json.load(f)
                source = meta.get("source_image", source)
            except Exception:
                pass
        if source not in artifacts:
            errors.append(f"Skipped {folder}: source image is not in current analysis.")
            continue
        try:
            nuc_arr = _zip_png_mask(zf, nuc_member)
            myo_arr = _zip_png_mask(zf, myo_member)
            art = artifacts[source]
            orig_rgb = art["rgb_u8"]
            add_to_queue(orig_rgb, nuc_mask=nuc_arr, myo_mask=myo_arr,
                         reason="user_correction_batch",
                         metadata={"source_image": source,
                                   "sample_id": art.get("sample_id", source),
                                   "original_filename": art.get("original_filename", source),
                                   "batch_folder": folder,
                                   "timestamp": datetime.now().isoformat()})
            submitted += 1
            used_sources.add(source)
        except Exception as exc:
            errors.append(f"Skipped {folder}: {exc}")

    # Also accept flat files like sample_nuclei_mask.png + sample_myotube_mask.png.
    flat = {Path(m).name: m for m in members if "/" not in m}
    stems = {}
    for filename, member in flat.items():
        lower = filename.lower()
        for suffix, key in (
            ("_nuclei_mask.png", "nuc"),
            ("_nuclei_mask.tif", "nuc"),
            ("_nuclei_mask.tiff", "nuc"),
            ("_myotube_mask.png", "myo"),
            ("_myotube_mask.tif", "myo"),
            ("_myotube_mask.tiff", "myo"),
        ):
            if lower.endswith(suffix):
                stem = filename[: -len(suffix)]
                stems.setdefault(stem, {})[key] = member
    for stem, pair in stems.items():
        if not {"nuc", "myo"}.issubset(pair):
            continue
        source = artifact_by_key.get(sample_match_key(stem))
        if source not in artifacts or source in used_sources:
            continue
        try:
            nuc_arr = _zip_png_mask(zf, pair["nuc"])
            myo_arr = _zip_png_mask(zf, pair["myo"])
            art = artifacts[source]
            add_to_queue(art["rgb_u8"], nuc_mask=nuc_arr, myo_mask=myo_arr,
                         reason="user_correction_batch",
                         metadata={"source_image": source,
                                   "sample_id": art.get("sample_id", source),
                                   "original_filename": art.get("original_filename", source),
                                   "batch_stem": stem,
                                   "timestamp": datetime.now().isoformat()})
            submitted += 1
        except Exception as exc:
            errors.append(f"Skipped {stem}: {exc}")

    return submitted, errors


# ─────────────────────────────────────────────────────────────────────────────
# Model  (architecture identical to training script)
# ─────────────────────────────────────────────────────────────────────────────


def correction_stem(upload_name: str) -> str:
    stem = normalize_sample_id(upload_name)
    lower = stem.lower()
    for suffix in (
        "_nuclei_mask", "-nuclei-mask", "_nuclei", "-nuclei", "_nuc", "-nuc",
        "_myotube_mask", "-myotube-mask", "_myotubes_mask", "-myotubes-mask",
        "_myotube", "-myotube", "_myotubes", "-myotubes", "_myo", "-myo",
    ):
        if lower.endswith(suffix):
            return normalize_sample_id(stem[: -len(suffix)])
    return normalize_sample_id(stem)


def uploaded_mask_to_array(uploaded_file) -> np.ndarray:
    return (np.array(Image.open(uploaded_file).convert("L")) > 0).astype(np.uint8)


def submit_uploaded_correction_pairs(nuc_files, myo_files, artifacts: dict, selected_name: str):
    nuc_files = nuc_files or []
    myo_files = myo_files or []
    if len(nuc_files) != len(myo_files):
        return 0, [f"Upload the same number of nuclei and myotube masks ({len(nuc_files)} nuclei, {len(myo_files)} myotube)."]

    artifact_by_key = {}
    for name, art in artifacts.items():
        for key in artifact_lookup_keys(name, art):
            artifact_by_key[key] = name

    nuc_by_stem = {sample_match_key(correction_stem(f.name)): f for f in nuc_files}
    myo_by_stem = {sample_match_key(correction_stem(f.name)): f for f in myo_files}
    common_stems = sorted(set(nuc_by_stem) & set(myo_by_stem))

    if len(nuc_files) == len(myo_files) == 1 and not common_stems:
        selected_key = sample_match_key(selected_name)
        common_stems = [selected_key]
        nuc_by_stem[selected_key] = nuc_files[0]
        myo_by_stem[selected_key] = myo_files[0]

    submitted = 0
    errors = []
    missing_nuc = sorted(set(myo_by_stem) - set(nuc_by_stem))
    missing_myo = sorted(set(nuc_by_stem) - set(myo_by_stem))
    for stem in missing_nuc:
        errors.append(f"Skipped {stem}: no matching nuclei mask.")
    for stem in missing_myo:
        errors.append(f"Skipped {stem}: no matching myotube mask.")

    for stem in common_stems:
        source = artifact_by_key.get(stem)
        if source is None and len(common_stems) == 1:
            source = selected_name
        if source not in artifacts:
            errors.append(f"Skipped {stem}: no matching analyzed image is loaded.")
            continue

        try:
            nuc_arr = uploaded_mask_to_array(nuc_by_stem[stem])
            myo_arr = uploaded_mask_to_array(myo_by_stem[stem])
            art = artifacts[source]
            add_to_queue(
                art["rgb_u8"],
                nuc_mask=nuc_arr,
                myo_mask=myo_arr,
                reason="user_correction_batch" if len(common_stems) > 1 else "user_correction",
                metadata={
                    "source_image": source,
                    "sample_id": art.get("sample_id", source),
                    "original_filename": art.get("original_filename", source),
                    "nuclei_file": nuc_by_stem[stem].name,
                    "myotube_file": myo_by_stem[stem].name,
                    "timestamp": datetime.now().isoformat(),
                },
            )
            submitted += 1
        except Exception as exc:
            errors.append(f"Skipped {stem}: {exc}")

    return submitted, errors

@st.cache_resource
def load_model(device: str):
    # Where the weights come from, in order:
    #   1. MYOSEG_MODEL_PATH (env var / Space variable)
    #   2. model_final.pt committed in the Space repo itself (next to src/ or inside it)
    #   3. the Hugging Face model repo MODEL_REPO_ID (needs HF_TOKEN if that repo is private)
    app_dir = Path(__file__).resolve().parent
    candidates = [os.environ.get("MYOSEG_MODEL_PATH"),
                  str(app_dir.parent / MODEL_FILENAME), str(app_dir / MODEL_FILENAME)]
    local = next((c for c in candidates if c and os.path.exists(c)), None)
    source = local
    if local is None:
        try:
            local = hf_hub_download(repo_id=MODEL_REPO_ID, filename=MODEL_FILENAME,
                                    force_download=False, token=os.environ.get("HF_TOKEN"))
            source = f"hf://{MODEL_REPO_ID}/{MODEL_FILENAME}"
        except Exception as exc:
            st.error(
                f"Could not load the model. No `{MODEL_FILENAME}` was found in the Space, and "
                f"downloading it from `{MODEL_REPO_ID}` failed ({type(exc).__name__}). "
                "Fix: add model_final.pt to the Space repo, or make the model repo public, "
                "or add a valid HF_TOKEN secret."
            )
            st.stop()
    file_sha = sha256_file(local)
    info = {
        "sha256": file_sha,
        "modified": time.ctime(os.path.getmtime(local)),
        "size_mb": round(os.path.getsize(local) / 1e6, 2),
        "source": source,
    }
    ckpt  = torch.load(local, map_location=device, weights_only=False)
    state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
    model = UNet(in_ch=2, out_ch=2, base=32)
    model.load_state_dict(state)
    model.to(device).eval()
    return model, info


def show_model_details(info: dict):
    """Drawn on every run (outside the cache) so it is always visible."""
    with st.sidebar.expander("Model details", expanded=False):
        st.caption(f"App version: `{APP_VERSION}`")
        st.caption(f"Loaded from: `{info.get('source', MODEL_REPO_ID)}`")
        st.caption(f"Size: {info['size_mb']:.2f} MB")
        st.caption(f"Modified: {info['modified']}")
        st.code(info["sha256"], language=None)
        st.caption("Architecture: `unet_v1`")


def _metric_float(row, key, default=0.0):
    try:
        value = row.get(key, default)
        if pd.isna(value):
            return float(default)
        return float(value)
    except Exception:
        return float(default)


def validation_flags_for_row(row) -> str:
    flags = []
    total_nuclei = _metric_float(row, "total_nuclei")
    myotubes = _metric_float(row, "myotube_count")
    myhc_pos = _metric_float(row, "myHC_positive_nuclei")
    fusion_index = _metric_float(row, "fusion_index")
    total_area = _metric_float(row, "total_area_um2")
    mean_area = _metric_float(row, "mean_area_um2")
    max_area = _metric_float(row, "max_area_um2")

    if total_nuclei <= 0:
        flags.append("empty nuclei count")
    if myotubes <= 0:
        flags.append("empty myotube count")
    if myhc_pos > total_nuclei and total_nuclei > 0:
        flags.append("MyHC+ nuclei exceeds total nuclei")
    if myotubes <= 2 and total_area > 10000:
        flags.append("possible merged myotube region")
    if max_area > 0 and mean_area > 0 and max_area >= 8 * mean_area:
        flags.append("one myotube area dominates field")
    if fusion_index > 60:
        flags.append("very high fusion index")
    if fusion_index < 1 and myotubes > 5:
        flags.append("very low fusion index")

    return "; ".join(flags)


def build_validation_report_zip(df: pd.DataFrame, settings: dict, errors=None, masks=None) -> bytes:
    errors = errors or []
    masks = masks or {}
    display_cols = [c for c in df.columns if not c.startswith("_")]
    report_df = df[display_cols].copy()
    report_df["review_flags"] = (report_df.apply(validation_flags_for_row, axis=1)
                                 if len(report_df) else pd.Series(dtype=str))
    for _c in ("image", "source_filename"):
        if _c not in report_df.columns:
            report_df[_c] = pd.Series(dtype=str)
    numeric_means = report_df.mean(numeric_only=True) if len(report_df) else {}

    summary = {
        "generated_at": datetime.now().isoformat(timespec="seconds"),
        "app_version": APP_VERSION,
        "model_repo": MODEL_REPO_ID,
        "model_file": MODEL_FILENAME,
        "images_analyzed": int(len(report_df)),
        "images_failed": int(len(errors)),
        "analysis_mode": settings.get("run_mode", ""),
        "preset": settings.get("analysis_preset", ""),
        "mean_total_nuclei": round(_metric_float(numeric_means, "total_nuclei"), 3) if len(report_df) else 0,
        "mean_myotube_count": round(_metric_float(numeric_means, "myotube_count"), 3) if len(report_df) else 0,
        "mean_fusion_index": round(_metric_float(numeric_means, "fusion_index"), 3) if len(report_df) else 0,
    }

    out = io.BytesIO()
    with zipfile.ZipFile(out, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("metrics.csv", report_df.to_csv(index=False).encode("utf-8"))
        zf.writestr(
            "validation_flags.csv",
            report_df[["image", "source_filename", "review_flags"]].to_csv(index=False).encode("utf-8"),
        )
        zf.writestr("run_summary.csv", pd.DataFrame([summary]).to_csv(index=False).encode("utf-8"))
        zf.writestr("settings.json", json.dumps(settings, indent=2, default=str).encode("utf-8"))
        if errors:
            zf.writestr("analysis_errors.csv", pd.DataFrame(errors).to_csv(index=False).encode("utf-8"))
        if masks:
            for name, pair in masks.items():
                if pair.get("nuclei") is not None:
                    zf.writestr(f"masks/{name}_nuclei_pp.png", pair["nuclei"])
                if pair.get("myotube") is not None:
                    zf.writestr(f"masks/{name}_myotube_pp.png", pair["myotube"])
        zf.writestr(
            "README.txt",
            (
                "MyoSeg validation report\n"
                "metrics.csv contains one row per analyzed image.\n"
                "validation_flags.csv marks images that should be checked manually.\n"
                "settings.json records the model, thresholds, postprocessing, and run mode.\n"
                "Masks are included only when Validation report mode is selected.\n"
            ).encode("utf-8"),
        )
    return out.getvalue()


# ─────────────────────────────────────────────────────────────────────────────
# PAGE CONFIG  +  CSS
# ─────────────────────────────────────────────────────────────────────────────

st.set_page_config(page_title="MyoSeg — Myotube & Nuclei Analyzer",
                   layout="wide", page_icon="🔬")

app_theme = st.sidebar.radio("Theme", ["Dark", "Light"], index=0, horizontal=True)
if app_theme == "Light":
    _theme = {
        "app_bg": "#f7f8fb",
        "sidebar_bg": "#ffffff",
        "text": "#172033",
        "muted": "#4b5568",
        "heading": "#1f5f99",
        "panel": "#ffffff",
        "panel_border": "#d6dbe8",
        "input_bg": "#ffffff",
        "flag_bg": "#fff1f1",
        "flag_border": "#d64545",
    }
else:
    _theme = {
        "app_bg": "#0e0e1a",
        "sidebar_bg": "#282832",
        "text": "#e0e0e0",
        "muted": "#b8bccb",
        "heading": "#90caf9",
        "panel": "#17172a",
        "panel_border": "#34344a",
        "input_bg": "#262631",
        "flag_bg": "#3e1a1a",
        "flag_border": "#ef5350",
    }

st.markdown(f"""
<style>
.stApp {{
    background:{_theme["app_bg"]};
    color:{_theme["text"]};
}}
.stApp [data-testid="stSidebar"] {{
    background:{_theme["sidebar_bg"]};
}}
.stApp [data-testid="stMarkdownContainer"],
.stApp label,
.stApp p,
.stApp span,
.stApp div {{
    color:{_theme["text"]};
}}
.stApp small,
.stApp [data-testid="stCaptionContainer"],
.stApp [data-testid="stHelp"] {{
    color:{_theme["muted"]};
}}
.block-container {{ max-width:1200px; padding-top:1.25rem; }}
h1,h2,h3,h4 {{ color:{_theme["heading"]}; }}
.stApp [data-testid="stMetric"],
.stApp [data-testid="stDataFrame"],
.stApp [data-testid="stExpander"],
.stApp [data-testid="stFileUploader"] section {{
    background:{_theme["panel"]};
    border-color:{_theme["panel_border"]};
}}
.stApp input,
.stApp textarea,
.stApp select,
.stApp [data-baseweb="select"] > div,
.stApp [data-baseweb="input"] > div {{
    background:{_theme["input_bg"]};
    color:{_theme["text"]};
}}
.flag-box {{
    background:{_theme["flag_bg"]}; border-left:4px solid {_theme["flag_border"]};
    padding:10px 16px; border-radius:8px; margin:8px 0;
}}
</style>
""", unsafe_allow_html=True)

st.title("🔬 MyoSeg — Myotube & Nuclei Analyzer")
device = "cuda" if torch.cuda.is_available() else "cpu"

# ─────────────────────────────────────────────────────────────────────────────
# SIDEBAR
# ─────────────────────────────────────────────────────────────────────────────
with st.sidebar:
    st.caption(f"Device: **{device}**")

    st.header("Input mapping")
    src1 = st.selectbox("Model channel 1 (MyHC / myotubes)",
                        ["Red", "Green", "Blue", "Grayscale"], index=0)
    inv1 = st.checkbox("Invert channel 1", value=False)
    src2 = st.selectbox("Model channel 2 (DAPI / nuclei)",
                        ["Red", "Green", "Blue", "Grayscale"], index=2)
    inv2 = st.checkbox("Invert channel 2", value=False)

    st.header("Preprocessing")
    image_size = st.select_slider("Model input size",
                                  options=[256, 384, 512, 640, 768, 1024], value=512)
    tta_label = st.selectbox(
        "Test-time augmentation",
        ["Off (fastest)", "4 flips", "8 flips + rotations"],
        index=0,
        help=(
            "Runs the U-Net on flipped/rotated copies of the image and averages them. "
            "Smoother, more stable masks and a per-pixel uncertainty estimate, at "
            "~4x / ~8x the inference time. Off reproduces v9 results exactly."
        ),
    )
    tta_mode = {"Off (fastest)": "none", "4 flips": "flip4", "8 flips + rotations": "d8"}[tta_label]

    st.header("Workflow")
    run_mode = st.radio(
        "Run mode",
        ["Batch speed", "Validation report", "Review"],
        index=0,
        help=(
            "Batch speed calculates metrics only. Validation report also saves masks. "
            "Review builds the full image viewer and overlays."
        ),
    )
    batch_speed = run_mode != "Review"
    validation_mode = run_mode == "Validation report"

    st.header("Analysis preset")
    analysis_preset = st.radio(
        "Preset",
        ["Tuned to manual counts", "Training defaults", "Standard", "Dense culture rescue", "Sparse cleanup"],
        index=0,
        help=(
            "Tuned to manual counts: settings chosen by cross-validated search against "
            "55 manually counted wells (lowest fusion-index / MyHC+ / myotube error). "
            "Training defaults reproduces earlier (v9) results. "
            "Dense culture rescue lowers the myotube threshold and uses more "
            "aggressive bridge-breaking/splitting to reduce myotube undercounting."
        ),
    )
    defaults = {**core.DEFAULT_PARAMS, **PRESETS[analysis_preset]}   # single source of truth: myoseg_core
    preset_key = analysis_preset.lower().replace(" ", "_")

    st.header("Thresholds")
    thr_nuc = st.slider(
        "Nuclei threshold", 0.05, 0.95, defaults["thr_nuc"], 0.01,
        key=f"thr_nuc_{preset_key}"
    )
    thr_myo = st.slider(
        "Myotube threshold", 0.05, 0.95, defaults["thr_myo"], 0.01,
        key=f"thr_myo_{preset_key}"
    )

    st.header("Fusion Index method")
    fi_method = st.radio(
        "FI classification method",
        ["Cytoplasm-hole (accurate, Lair 2025)", "Pixel-overlap (legacy)"],
        index=0 if defaults["fi_method"] == "cytoplasm-hole" else 1,
        key=f"fi_method_{preset_key}",
        help=(
            "Cytoplasm-hole: checks for a MyHC signal dip beneath each nucleus — "
            "eliminates false positives from nuclei sitting above/below myotubes in Z. "
            "Pixel-overlap: legacy method that overestimates FI (Lair et al. 2025)."
        )
    )
    use_hole_method = fi_method.startswith("Cytoplasm")
    hole_ratio_thr = st.slider(
        "Hole ratio threshold", 0.50, 0.99, float(defaults["hole_ratio_thr"]), 0.01,
        key=f"hole_ratio_thr_{preset_key}",
        help=(
            "A nucleus is counted as fused if its MyHC signal is less than "
            "this fraction of the surrounding cytoplasm ring signal. "
            "Lower = stricter (fewer nuclei counted as fused). "
            "0.85 is the value validated by Lair et al. 2025."
        ),
        disabled=not use_hole_method,
    )
    ring_width_px = st.number_input(
        "Cytoplasm ring width (px)", 2, 20, 6, 1,
        help="Width of the ring around each nucleus used to measure local MyHC intensity.",
        disabled=not use_hole_method,
    )

    st.header("Postprocessing")
    min_nuc_area = st.number_input(
        "Min nucleus area (px)", 0, 10000, defaults["min_nuc_area"], 1,
        key=f"min_nuc_area_{preset_key}"
    )
    min_myo_area = st.number_input(
        "Min myotube area (px)", 0, 200000, defaults["min_myo_area"], 10,
        key=f"min_myo_area_{preset_key}"
    )
    nuc_close_radius = st.number_input(
        "Nuclei close radius", 0, 50, defaults["nuc_close_radius"], 1,
        key=f"nuc_close_radius_{preset_key}"
    )
    myo_open_radius = st.number_input(
        "Myotube open radius", 0, 50, defaults["myo_open_radius"], 1,
        key=f"myo_open_radius_{preset_key}",
        help="Opening removes small noise without merging separate myotubes. "
             "Replaces the old closing radius which was merging adjacent myotubes.")

    st.header("Myotube separation")
    st.caption(
        "These controls break apart touching/bridged myotubes that would "
        "otherwise be counted as a single object."
    )
    myo_erode_radius = st.number_input(
        "Myotube erode radius (px)", 0, 15, defaults["myo_erode_radius"], 1,
        key=f"myo_erode_radius_{preset_key}",
        help=(
            "Erode + re-dilate breaks thin pixel bridges between adjacent "
            "myotubes while preserving their overall size. "
            "Start at 2 px; increase to 3–4 px for very dense cultures. "
            "Set 0 to disable."
        )
    )
    min_myo_aspect_ratio = st.number_input(
        "Min myotube aspect ratio", 0.0, 10.0, defaults["min_myo_aspect_ratio"], 0.1,
        key=f"min_myo_aspect_ratio_{preset_key}",
        help=(
            "Rejects round blobs (debris/artifacts) that are not real myotubes. "
            "Myotubes are elongated (aspect ratio > 3). Round objects have ~1. "
            "Set to 1.5–2.0 to filter false positives in sparse cultures. "
            "Set to 0 to disable (default)."
        )
    )
    myo_max_area_px = st.number_input(
        "Max myotube area before split (px)", 0, 500000, defaults["myo_max_area_px"], 500,
        key=f"myo_max_area_px_{preset_key}",
        help=(
            "Any connected myotube region larger than this is split using "
            "nucleus-seeded watershed. Set to 0 to disable. "
            "Increase for cultures with legitimately large single myotubes."
        )
    )
    myo_split_min_seeds = st.number_input(
        "Min nuclei seeds to split", 2, 20, defaults["myo_split_min_seeds"], 1,
        key=f"myo_split_min_seeds_{preset_key}",
        help=(
            "Minimum nucleus centroids required before splitting a large region. "
            "Set to 2 to split merged pairs; increase if single large myotubes "
            "are being incorrectly split."
        )
    )

    apply_split_to_metrics = st.checkbox(
        "Use split myotubes for counts", value=bool(defaults["apply_split_to_metrics"]),
        key=f"apply_split_{preset_key}",
        help=(
            "Off (default, matches v9 and the validation data): the large-myotube split "
            "above only changes the viewer/overlays; counts use connected components. "
            "On: the split is applied before counting, so myotube count, nuclei per "
            "myotube and fusion index reflect the separated myotubes."
        ),
    )
    myhc_mask_label = st.radio(
        "Mask used to decide MyHC⁺ nuclei",
        ["Separated (counting) mask", "Full mask (no bridge-breaking)"],
        index=0 if defaults["myhc_mask"] == "separated" else 1,
        key=f"myhc_mask_{preset_key}",
        help=(
            "Bridge-breaking erosion trims thin myotube parts, so nuclei there can be "
            "lost as MyHC+. 'Full mask' counts myotubes on the separated mask but decides "
            "MyHC+ on the untrimmed mask."
        ),
    )
    myhc_mask = "separated" if myhc_mask_label.startswith("Separated") else "full"

    st.header("Watershed (nuclei splitting)")
    nuc_ws_min_dist = st.number_input("Min watershed distance", 1, 30, int(defaults["nuc_ws_min_distance"]), 1,
                                      key=f"nuc_ws_min_dist_{preset_key}")
    nuc_ws_min_area = st.number_input("Min watershed area (px)", 1, 500, int(defaults["nuc_ws_min_area"]), 1,
                                      key=f"nuc_ws_min_area_{preset_key}")

    st.header("Overlay")
    nuc_hex   = st.color_picker("Nuclei colour",  "#00FFFF")
    myo_hex   = st.color_picker("Myotube colour", "#FF0000")
    alpha     = st.slider("Overlay alpha", 0.0, 1.0, 0.45, 0.01)
    nuc_rgb   = hex_to_rgb(nuc_hex)
    myo_rgb   = hex_to_rgb(myo_hex)
    label_nuc = st.checkbox("Show nucleus IDs on overlay",  value=False)
    label_myo = st.checkbox("Show myotube IDs on overlay",  value=False)
    if batch_speed:
        label_nuc = False
        label_myo = False

    st.header("Surface area")
    px_um = st.number_input("Pixel size (µm) — set for real µm²",
                             value=1.0, min_value=0.01, step=0.01)

    st.header("Active learning")
    enable_al = st.toggle("Enable correction upload", value=False)

    st.header("Privacy & Training")
    private_mode = st.toggle(
        "🔒 Private mode",
        value=True,
        help=(
            "When enabled, your images are processed locally only. "
            "They are NOT added to the retraining queue, NOT saved to "
            "corrections/, and NOT used for model improvement in any way. "
            "Use this for unpublished data or sensitive research images."
        )
    )
    if private_mode:
        st.info(
            "🔒 **Private mode ON** — your images will not be used for "
            "training or stored beyond this session."
        )

    st.header("Metric definitions")
    with st.expander("Fusion Index"):
        st.write("100 × (nuclei in myotubes with ≥2 nuclei) / total nuclei")
    with st.expander("MyHC-positive nucleus"):
        st.write("Counted if ≥10% of nucleus pixels overlap a myotube.")
    with st.expander("Surface area"):
        st.write("Pixel count × px_um². Set pixel size for real µm² values.")


# ─────────────────────────────────────────────────────────────────────────────
# FILE UPLOADER
# ─────────────────────────────────────────────────────────────────────────────
_uploader_label = (
    "Upload 1+ images (png / jpg / tif). Any filename is accepted. 🔒 Private mode is ON — images will not be stored."
    if private_mode else
    "Upload 1+ images (png / jpg / tif). Any filename is accepted. Images may be used for model improvement."
)
if "image_upload_reset_id" not in st.session_state:
    st.session_state.image_upload_reset_id = 0

uploads = st.file_uploader(
    _uploader_label,
    type=["png", "jpg", "jpeg", "tif", "tiff"],
    accept_multiple_files=True,
    key=f"image_uploads_{st.session_state.image_upload_reset_id}",
)

for key in ("df", "artifacts", "zip_bytes", "bio_metrics", "validation_report_bytes", "analysis_errors", "run_mode",
            "qc_results", "run_info"):
    if key not in st.session_state:
        st.session_state[key] = None

if uploads:
    c_uploads, _ = st.columns([1, 4])
    with c_uploads:
        if st.button(
            "Delete all uploaded images",
            help="Clears the current batch and removes the generated results from this session.",
            key="clear_all_uploaded_images",
        ):
            st.session_state.image_upload_reset_id += 1
            for key in ("df", "artifacts", "zip_bytes", "bio_metrics", "validation_report_bytes", "analysis_errors", "run_mode",
                        "qc_results", "run_info"):
                st.session_state[key] = None
            st.rerun()

model, model_info = load_model(device=device)   # loaded before the upload check so
model_sha = model_info["sha256"]                  # "Model details" is always visible
show_model_details(model_info)

if not uploads:
    st.info("👆 Upload one or more fluorescence images to get started.")
    st.stop()


# Cache of U-Net outputs. Changing thresholds / postprocessing / FI settings and
# pressing Rerun reuses these instead of running the network again.
if "prob_cache" not in st.session_state:
    st.session_state.prob_cache = core.ProbCache(max_items=64)
prob_cache = st.session_state.prob_cache

# The same parameter dict the core / chat assistant / API use.
core_params = core.make_params(
    src_myo=src1, inv_myo=bool(inv1), src_nuc=src2, inv_nuc=bool(inv2),
    image_size=int(image_size), tta=tta_mode,
    fi_method="cytoplasm-hole" if use_hole_method else "pixel-overlap",
    hole_ratio_thr=float(hole_ratio_thr), ring_width=int(ring_width_px),
    nuc_ws_min_distance=int(nuc_ws_min_dist), nuc_ws_min_area=int(nuc_ws_min_area),
    px_um=float(px_um),
    thr_nuc=float(thr_nuc), thr_myo=float(thr_myo),
    min_nuc_area=int(min_nuc_area), min_myo_area=int(min_myo_area),
    nuc_close_radius=int(nuc_close_radius), myo_open_radius=int(myo_open_radius),
    myo_erode_radius=int(myo_erode_radius), min_myo_aspect_ratio=float(min_myo_aspect_ratio),
    myo_max_area_px=int(myo_max_area_px), myo_split_min_seeds=int(myo_split_min_seeds),
    apply_split_to_metrics=bool(apply_split_to_metrics), myhc_mask=myhc_mask,
)

analysis_settings = {
    "app_version": APP_VERSION,
    "app_theme": app_theme,
    "run_mode": run_mode,
    "analysis_preset": analysis_preset,
    "model_repo": MODEL_REPO_ID,
    "model_filename": MODEL_FILENAME,
    "model_sha256": model_sha,
    "device": device,
    "image_size": int(image_size),
    "test_time_augmentation": tta_mode,
    "uncertainty_flag_threshold": UNC_FLAG_THR,
    "channel_1": src1,
    "channel_1_inverted": bool(inv1),
    "channel_2": src2,
    "channel_2_inverted": bool(inv2),
    "threshold_nuclei": float(thr_nuc),
    "threshold_myotube": float(thr_myo),
    "fi_method": "cytoplasm-hole" if use_hole_method else "pixel-overlap",
    "hole_ratio_threshold": float(hole_ratio_thr),
    "cytoplasm_ring_width_px": int(ring_width_px),
    "min_nuclei_area_px": int(min_nuc_area),
    "min_myotube_area_px": int(min_myo_area),
    "nuclei_close_radius": int(nuc_close_radius),
    "myotube_open_radius": int(myo_open_radius),
    "myotube_erode_radius": int(myo_erode_radius),
    "min_myotube_aspect_ratio": float(min_myo_aspect_ratio),
    "max_myotube_area_before_split_px": int(myo_max_area_px),
    "min_nuclei_seeds_to_split": int(myo_split_min_seeds),
    "apply_split_to_metrics": bool(apply_split_to_metrics),
    "myhc_mask": myhc_mask,
    "core_version": core.CORE_VERSION,
    "nuclei_watershed_min_distance": int(nuc_ws_min_dist),
    "nuclei_watershed_min_area_px": int(nuc_ws_min_area),
    "pixel_size_um": float(px_um),
    "private_mode": bool(private_mode),
}

# ─────────────────────────────────────────────────────────────────────────────
# RUN ANALYSIS
# ─────────────────────────────────────────────────────────────────────────────
with st.form("run_form"):
    run = st.form_submit_button("▶  Run / Rerun analysis", type="primary")

if run:
    results          = []
    artifacts        = {}
    all_bio_metrics  = {}
    low_conf_flags   = []
    analysis_errors  = []
    report_masks     = {}
    qc_results       = {}
    cache_hits       = 0
    zip_buf          = io.BytesIO()
    build_visual_outputs = run_mode == "Review"
    save_report_masks = run_mode == "Validation report"

    run_t0 = time.perf_counter()
    with st.spinner("Analysing images…"):
        with zipfile.ZipFile(zip_buf, "w", compression=zipfile.ZIP_STORED) as zf:
            prog = st.progress(0.0)
            used_sample_ids = set()

            for i, up in enumerate(uploads):
                original_filename = up.name
                original_stem = Path(original_filename).stem
                name = unique_sample_id(original_filename, used_sample_ids)
                try:
                    rgb_u8 = np.array(
                        Image.open(io.BytesIO(up.getvalue())).convert("RGB"),
                        dtype=np.uint8
                    )

                    # Inference (skipped when this image + channel mapping + size + model
                    # were already run) → threshold → postprocess → metrics.
                    res = core.analyze_image(
                        model, rgb_u8, core_params, device=device,
                        cache=prob_cache, model_sha=model_sha,
                        with_labels=build_visual_outputs,
                    )
                    probs = res["probs"]
                    cache_hits += int(res["cache_hit"])
                    qc_results[name] = core.quality_check(rgb_u8, core_params, probs)

                    # Uncertainty check (active learning) — only queue if NOT in private
                    # mode, and only on fresh inference so reruns don't queue twice.
                    conf = res["confidence"]
                    unc = res["uncertainty"]["uncertainty"]
                    reason = ("high_uncertainty" if unc > UNC_FLAG_THR
                              else "low_confidence" if conf < CONF_FLAG_THR else None)
                    if reason:
                        low_conf_flags.append((name, unc))
                        if not private_mode and not res["cache_hit"]:
                            add_to_queue(rgb_u8, reason=reason,
                                         metadata={
                                             "confidence": conf,
                                             "uncertainty": res["uncertainty"],
                                             "tta": tta_mode,
                                             "filename": original_filename,
                                             "sample_id": name,
                                         })

                    nuc_raw, myo_raw = res["nuc_raw"], res["myo_raw"]
                    nuc_pp, myo_pp = res["nuc_pp"], res["myo_pp"]
                    bio = res["metrics"]
                    per_areas = bio.pop("_per_myotube_areas", [])
                    bio_myo_ids = bio.pop("_bio_myo_ids", set())
                    bio.pop("_total_cc_count", 0)
                    bio["image"] = name
                    bio["source_filename"] = original_filename
                    bio["model_uncertainty"] = unc
                    results.append(bio)
                    all_bio_metrics[name] = {**bio, "_per_myotube_areas": per_areas}

                    nuc_pp_bytes = png_bytes((nuc_pp * 255).astype(np.uint8))
                    myo_pp_bytes = png_bytes((myo_pp * 255).astype(np.uint8))
                    nuc_raw_bytes = png_bytes((nuc_raw * 255).astype(np.uint8))
                    myo_raw_bytes = png_bytes((myo_raw * 255).astype(np.uint8))

                    artifacts[name] = {
                        "sample_id": name,
                        "original_filename": original_filename,
                        "original_stem": original_stem,
                    }

                    if save_report_masks:
                        report_masks[name] = {"nuclei": nuc_pp_bytes, "myotube": myo_pp_bytes}

                    if build_visual_outputs:
                        artifacts[name].update({
                            "rgb_u8": rgb_u8,
                            "nuc_pp_arr": nuc_pp,
                            "myo_pp_arr": myo_pp,
                            "nuc_pp": nuc_pp_bytes,
                            "myo_pp": myo_pp_bytes,
                            "nuc_raw_bytes": nuc_raw_bytes,
                            "myo_raw_bytes": myo_raw_bytes,
                        })
                        simple_ov = make_simple_overlay(
                            rgb_u8, nuc_pp, myo_pp, nuc_rgb, myo_rgb, float(alpha)
                        )
                        # Labels come from the core (watershed nuclei + split myotubes),
                        # so they are no longer computed a second time here.
                        nuc_lab = res["nuc_lab"]
                        myo_lab = res["myo_lab"]

                        inst_px = make_coloured_overlay(rgb_u8, nuc_lab, myo_lab, alpha=float(alpha))
                        nuc_only_px = make_coloured_overlay(rgb_u8, nuc_lab, np.zeros_like(myo_lab), alpha=float(alpha))
                        myo_only_px = make_coloured_overlay(rgb_u8, np.zeros_like(nuc_lab), myo_lab, alpha=float(alpha))

                        orig_h_img, orig_w_img = rgb_u8.shape[:2]
                        label_positions = collect_label_positions(
                            nuc_lab, myo_lab, orig_w_img, orig_h_img,
                            bio_myo_ids=bio_myo_ids,
                        )
                        artifacts[name].update({
                            "nuc_lab": nuc_lab,
                            "myo_lab": myo_lab,
                            "label_positions": label_positions,
                            "img_w": orig_w_img,
                            "img_h": orig_h_img,
                        })

                        outline_ov = make_outline_overlay(rgb_u8, nuc_lab, myo_lab,
                                                           nuc_color=nuc_rgb, myo_color=(0, 255, 0),
                                                           line_width=2)
                        zf.writestr(f"{name}/overlay_combined.png",  png_bytes(simple_ov))
                        zf.writestr(f"{name}/overlay_instance.png",  png_bytes(inst_px))
                        zf.writestr(f"{name}/overlay_nuclei.png",    png_bytes(nuc_only_px))
                        zf.writestr(f"{name}/overlay_myotubes.png",  png_bytes(myo_only_px))
                        zf.writestr(f"{name}/overlay_outlines.png",  png_bytes(outline_ov))
                        zf.writestr(f"{name}/nuclei_pp.png",         nuc_pp_bytes)
                        zf.writestr(f"{name}/myotube_pp.png",        myo_pp_bytes)
                        zf.writestr(f"{name}/nuclei_raw.png",        nuc_raw_bytes)
                        zf.writestr(f"{name}/myotube_raw.png",       myo_raw_bytes)
                except Exception as exc:
                    # One unreadable/odd image no longer stops the whole batch.
                    analysis_errors.append({"image": name, "source_filename": original_filename,
                                            "error": f"{type(exc).__name__}: {exc}"})

                prog.progress((i + 1) / len(uploads))

            df = pd.DataFrame(results)
            if len(df):
                df = df.sort_values("image")
            zf.writestr("metrics.csv", df.to_csv(index=False).encode("utf-8"))
            zf.writestr("settings.json", json.dumps(analysis_settings, indent=2, default=str).encode("utf-8"))

    st.session_state.df          = df
    st.session_state.artifacts   = artifacts
    st.session_state.zip_bytes   = zip_buf.getvalue()
    st.session_state.bio_metrics = all_bio_metrics
    st.session_state.validation_report_bytes = build_validation_report_zip(
        df, analysis_settings, errors=analysis_errors, masks=report_masks
    )
    st.session_state.analysis_errors = analysis_errors
    st.session_state.run_mode = run_mode
    st.session_state.qc_results = qc_results
    st.session_state.run_info = {
        "seconds": round(time.perf_counter() - run_t0, 1),
        "cache_hits": cache_hits,
        "n": len(uploads),
    }

    if low_conf_flags:
        names_str = ", ".join(f"{n} (uncertainty={c:.2f})" for n, c in low_conf_flags)
        _queued = ("queued for labelling / retraining" if not private_mode
                   else "not queued (Private mode is on)")
        st.markdown(
            f"<div class='flag-box'>⚠️ <b>Uncertain images — {_queued}:</b> "
            f"{names_str}</div>",
            unsafe_allow_html=True,
        )

if st.session_state.df is None:
    st.info("Click **▶ Run / Rerun analysis** to generate results.")
    st.stop()

# ─────────────────────────────────────────────────────────────────────────────
# RESULTS TABLE  +  DOWNLOADS
# ─────────────────────────────────────────────────────────────────────────────
st.subheader("Run summary")
summary_cols = st.columns(4)
summary_cols[0].metric("Images analyzed", len(st.session_state.df))
summary_cols[1].metric("Images failed", len(st.session_state.analysis_errors or []))
summary_cols[2].metric("Mode", st.session_state.run_mode or run_mode)
summary_cols[3].metric("Preset", analysis_preset)
_ri = st.session_state.run_info
if _ri:
    st.caption(
        f"Finished in {_ri['seconds']} s. "
        + (f"Reused saved model output for {_ri['cache_hits']} of {_ri['n']} image(s) — "
           "only thresholds/postprocessing were recomputed."
           if _ri["cache_hits"] else
           "Tip: after changing thresholds or postprocessing, Rerun skips the neural network and is much faster.")
    )

# Image-quality checks (plain image statistics, no extra model)
_qc = st.session_state.qc_results or {}
if _qc:
    _qc_flags = {k: list(v["warnings"]) for k, v in _qc.items()}
    for k, extra in core.batch_quality_outliers(_qc).items():
        _qc_flags.setdefault(k, []).extend(extra)
    _qc_flags = {k: v for k, v in _qc_flags.items() if v}
    with st.expander(f"🩺 Image quality checks — {len(_qc_flags)} image(s) with warnings",
                     expanded=bool(_qc_flags)):
        if not _qc_flags:
            st.write("No quality problems detected (saturation, empty channels, channel mapping, blur, model uncertainty).")
        for k, warns in _qc_flags.items():
            st.markdown(f"**{k}**")
            for w in warns:
                st.markdown(f"- {w}")

# Rule-based diagnosis of suspicious results → suggested settings to try
_diag = {k: core.diagnose(v, core_params) for k, v in (st.session_state.bio_metrics or {}).items()}
_diag = {k: v for k, v in _diag.items() if v}
if _diag:
    with st.expander(f"💡 Suggested parameter changes — {len(_diag)} image(s)"):
        st.caption("Suggestions are starting points; nothing is changed automatically. "
                   "Adjust the sidebar and Rerun — the model output is reused, so it is quick.")
        for k, items in _diag.items():
            st.markdown(f"**{k}**")
            for d in items:
                sugg = ", ".join(f"`{p}` → {v}" for p, v in d["suggest"].items()) or "inspect manually"
                st.markdown(f"- *{d['issue']}* — {d['explanation']}  \n  Try: {sugg}")
if st.session_state.analysis_errors:
    with st.expander("Analysis errors"):
        st.dataframe(pd.DataFrame(st.session_state.analysis_errors), use_container_width=True)

st.subheader("📋 Results")
display_cols = [c for c in st.session_state.df.columns if not c.startswith("_")]
st.dataframe(st.session_state.df[display_cols], use_container_width=True, height=320)

c1, c2, c3, c4 = st.columns(4)
with c1:
    st.download_button("⬇️ Download metrics.csv",
                       st.session_state.df[display_cols].to_csv(index=False).encode(),
                       file_name="metrics.csv", mime="text/csv")
with c2:
    report_bytes = st.session_state.validation_report_bytes or build_validation_report_zip(
        st.session_state.df, analysis_settings, errors=st.session_state.analysis_errors
    )
    st.download_button("Download validation report",
                       report_bytes,
                       file_name=f"myoseg_validation_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip",
                       mime="application/zip")
with c3:
    st.download_button("⬇️ Download results.zip",
                       st.session_state.zip_bytes,
                       file_name="results.zip", mime="application/zip")
with c4:
    # Rebuild ZIP with CURRENT colour / alpha settings — no model rerun needed
    if (st.session_state.run_mode or run_mode) != "Review":
        st.caption("Colour rebuild is available in Review mode.")
    elif st.button("🎨 Rebuild ZIP with current colours", help=(
        "Regenerates the overlay images in the ZIP using the current "
        "colour picker and alpha values from the sidebar."
    )):
        import base64 as _b64_zip
        new_zip_buf = io.BytesIO()
        with zipfile.ZipFile(new_zip_buf, "w", compression=zipfile.ZIP_STORED) as zf:
            for img_name, art in st.session_state.artifacts.items():
                _r  = art["rgb_u8"]
                _nl = art["nuc_lab"]
                _ml = art["myo_lab"]
                _zn = np.zeros_like(_nl)
                _zm = np.zeros_like(_ml)
                ov_comb = make_coloured_overlay(_r, _nl, _ml,
                                                alpha=float(alpha),
                                                nuc_color=nuc_rgb, myo_color=myo_rgb)
                ov_nuc  = make_coloured_overlay(_r, _nl, _zm,
                                                alpha=float(alpha),
                                                nuc_color=nuc_rgb, myo_color=None)
                ov_myo  = make_coloured_overlay(_r, _zn, _ml,
                                                alpha=float(alpha),
                                                nuc_color=None, myo_color=myo_rgb)
                simple  = make_simple_overlay(_r,
                                              (_nl > 0).astype(np.uint8),
                                              (_ml > 0).astype(np.uint8),
                                              nuc_rgb, myo_rgb, float(alpha))
                outline = make_outline_overlay(_r, _nl, _ml,
                                               nuc_color=nuc_rgb, myo_color=(0, 255, 0),
                                               line_width=2)
                zf.writestr(f"{img_name}/overlay_combined.png",  png_bytes(simple))
                zf.writestr(f"{img_name}/overlay_instance.png",  png_bytes(ov_comb))
                zf.writestr(f"{img_name}/overlay_nuclei.png",    png_bytes(ov_nuc))
                zf.writestr(f"{img_name}/overlay_myotubes.png",  png_bytes(ov_myo))
                zf.writestr(f"{img_name}/overlay_outlines.png",  png_bytes(outline))
                zf.writestr(f"{img_name}/nuclei_pp.png",         png_bytes((art["nuc_pp_arr"] > 0).astype(np.uint8) * 255))
                zf.writestr(f"{img_name}/myotube_pp.png",        png_bytes((art["myo_pp_arr"] > 0).astype(np.uint8) * 255))
                zf.writestr(f"{img_name}/nuclei_raw.png",        art["nuc_raw_bytes"])
                zf.writestr(f"{img_name}/myotube_raw.png",       art["myo_raw_bytes"])
            df_cols = [c for c in st.session_state.df.columns if not c.startswith("_")]
            zf.writestr("metrics.csv", st.session_state.df[df_cols].to_csv(index=False).encode())
        st.session_state.zip_bytes = new_zip_buf.getvalue()
        st.success("ZIP rebuilt with current colours. Click Download results.zip above to save.")

st.divider()

# ─────────────────────────────────────────────────────────────────────────────
if (st.session_state.run_mode or run_mode) != "Review":
    st.info("Batch run complete. Switch Run mode to Review and rerun if you want overlays, per-image tabs, and live preview.")
    st.stop()

# PER-IMAGE PREVIEW  +  ANIMATED METRICS
# ─────────────────────────────────────────────────────────────────────────────
st.subheader("🖼️ Image preview & live metrics")
names = list(st.session_state.artifacts.keys())
pick  = st.selectbox("Select image", names)

col_img, col_metrics = st.columns([3, 2], gap="large")

with col_img:
    tabs = st.tabs([
        "🔵 Combined",
        "📐 Nuclei outlines",
        "📐 Myotube outlines",
        "🟣 Nuclei only",
        "🟠 Myotubes only",
        "📷 Original",
        "⬜ Nuclei mask",
        "⬜ Myotube mask",
    ])
    art     = st.session_state.artifacts[pick]
    bio_cur = st.session_state.bio_metrics.get(pick, {})
    lpos    = art["label_positions"]
    iw      = art["img_w"]
    ih      = art["img_h"]

    # Build coloured overlays RIGHT NOW using the current sidebar colour / alpha.
    # This means changing colour picker or alpha slider instantly updates the
    # viewer — no rerun needed for display changes.
    import base64 as _b64_disp
    def _b64png_disp(arr):
        return _b64_disp.b64encode(png_bytes(arr)).decode()

    _rgb  = art["rgb_u8"]
    _nl   = art["nuc_lab"]
    _ml   = art["myo_lab"]
    _zero_nuc = np.zeros_like(_nl)
    _zero_myo = np.zeros_like(_ml)

    inst_b64     = _b64png_disp(make_coloured_overlay(_rgb, _nl, _ml,
                                  alpha=float(alpha),
                                  nuc_color=nuc_rgb, myo_color=myo_rgb))
    nuc_only_b64 = _b64png_disp(make_coloured_overlay(_rgb, _nl, _zero_myo,
                                  alpha=float(alpha),
                                  nuc_color=nuc_rgb, myo_color=None))
    myo_only_b64 = _b64png_disp(make_coloured_overlay(_rgb, _zero_nuc, _ml,
                                  alpha=float(alpha),
                                  nuc_color=None, myo_color=myo_rgb))

    with tabs[0]:
        html_combined = make_svg_viewer(
            inst_b64, iw, ih, lpos,
            show_nuclei=True, show_myotubes=True,
        )
        st.components.v1.html(html_combined, height=680, scrolling=False)

    with tabs[1]:
        # Nuclei-only outlines
        nuc_outline_img = make_outline_overlay(
            _rgb, _nl, np.zeros_like(_ml),
            nuc_color=nuc_rgb, myo_color=(0, 255, 0),
            line_width=2,
        )
        nuc_outline_b64 = _b64png_disp(nuc_outline_img)
        nuc_outline_lpos = {"nuclei": lpos["nuclei"], "myotubes": [], "myotubes_nonbio": []}
        html_nuc_outline = make_svg_viewer(
            nuc_outline_b64, iw, ih, nuc_outline_lpos,
            show_nuclei=True, show_myotubes=False,
        )
        st.components.v1.html(html_nuc_outline, height=680, scrolling=False)

    with tabs[2]:
        # Myotube-only outlines
        myo_outline_img = make_outline_overlay(
            _rgb, np.zeros_like(_nl), _ml,
            nuc_color=nuc_rgb, myo_color=(0, 255, 0),
            line_width=2,
        )
        myo_outline_b64 = _b64png_disp(myo_outline_img)
        myo_outline_lpos = {"nuclei": [], "myotubes": lpos["myotubes"],
                            "myotubes_nonbio": lpos.get("myotubes_nonbio", [])}
        html_myo_outline = make_svg_viewer(
            myo_outline_b64, iw, ih, myo_outline_lpos,
            show_nuclei=False, show_myotubes=True,
        )
        st.components.v1.html(html_myo_outline, height=680, scrolling=False)

    with tabs[3]:
        nuc_only_lpos = {"nuclei": lpos["nuclei"], "myotubes": [], "myotubes_nonbio": []}
        html_nuc = make_svg_viewer(
            nuc_only_b64, iw, ih, nuc_only_lpos,
            show_nuclei=True, show_myotubes=False,
        )
        st.components.v1.html(html_nuc, height=680, scrolling=False)

    with tabs[4]:
        myo_only_lpos = {"nuclei": [], "myotubes": lpos["myotubes"],
                         "myotubes_nonbio": lpos.get("myotubes_nonbio", [])}
        html_myo = make_svg_viewer(
            myo_only_b64, iw, ih, myo_only_lpos,
            show_nuclei=False, show_myotubes=True,
        )
        st.components.v1.html(html_myo, height=680, scrolling=False)

    with tabs[5]:
        st.image(art["rgb_u8"], use_container_width=True)
    with tabs[6]:
        nuc_mask_view = (art.get("nuc_pp_arr", np.zeros((_rgb.shape[0], _rgb.shape[1]), dtype=np.uint8)) > 0).astype(np.uint8) * 255
        if int(nuc_mask_view.max()) == 0:
            st.warning("The nuclei mask is empty for this image.")
        st.image(nuc_mask_view, clamp=True, use_container_width=True)
    with tabs[7]:
        myo_mask_view = (art.get("myo_pp_arr", np.zeros((_rgb.shape[0], _rgb.shape[1]), dtype=np.uint8)) > 0).astype(np.uint8) * 255
        if int(myo_mask_view.max()) == 0:
            st.warning("The myotube mask is empty for this image.")
        st.image(myo_mask_view, clamp=True, use_container_width=True)

with col_metrics:
    st.markdown("#### 📊 Live metrics")
    bio       = st.session_state.bio_metrics.get(pick, {})
    per_areas = bio.get("_per_myotube_areas", [])

    r1c1, r1c2, r1c3 = st.columns(3)
    r2c1, r2c2, r2c3 = st.columns(3)
    r3c1, r3c2, r3c3 = st.columns(3)

    placeholders = {
        "total_nuclei"            : r1c1.empty(),
        "myotube_count"           : r1c2.empty(),
        "myHC_positive_nuclei"    : r1c3.empty(),
        "myHC_positive_percentage": r2c1.empty(),
        "fusion_index"            : r2c2.empty(),
        "avg_nuclei_per_myotube"  : r2c3.empty(),
        "total_area_um2"          : r3c1.empty(),
        "mean_area_um2"           : r3c2.empty(),
        "max_area_um2"            : r3c3.empty(),
    }

    specs = [
        ("total_nuclei",             "Total nuclei",      "#4fc3f7", False),
        ("myotube_count",            "Myotubes",          "#ff8a65", False),
        ("myHC_positive_nuclei",     "MyHC⁺ nuclei",      "#a5d6a7", False),
        ("myHC_positive_percentage", "MyHC⁺ %",           "#ce93d8", True),
        ("fusion_index",             "Fusion index %",     "#80cbc4", True),
        ("avg_nuclei_per_myotube",   "Avg nuc/myotube",   "#80deea", True),
        ("total_area_um2",           f"Total area (µm²)",  "#fff176", True),
        ("mean_area_um2",            f"Mean area (µm²)",   "#ffcc80", True),
        ("max_area_um2",             f"Max area (µm²)",    "#ef9a9a", True),
    ]

    for key, label, color, is_float in specs:
        val = bio.get(key, 0)
        animated_metric(placeholders[key], label,
                        float(val) if is_float else int(val),
                        color=color)

    if per_areas:
        st.markdown("#### 📐 Per-myotube area")
        area_df = pd.DataFrame({
            "Myotube"      : [f"M{i+1}" for i in range(len(per_areas))],
            f"Area (µm²)"  : per_areas,
        }).set_index("Myotube")
        st.bar_chart(area_df, height=220)

st.divider()

# ─────────────────────────────────────────────────────────────────────────────
# TRAINING CONTRIBUTION — User-initiated parameter + image submission
# ─────────────────────────────────────────────────────────────────────────────
if not private_mode and names:
    st.subheader("📤 Submit image for training")
    st.caption(
        "Once you've tuned the sidebar parameters to get the best results for "
        "this image, you can download the current image/mask batch or submit "
        "the parameter snapshot.  FIJI-corrected masks submitted below are "
        "used as supervised correction data."
    )

    st.download_button(
        "Download all current masks as ZIP",
        data=build_mask_batch_zip(names, st.session_state.artifacts),
        file_name=f"myoseg_current_masks_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip",
        mime="application/zip",
        key="download_current_masks_batch",
        help="Downloads every analyzed image with the current postprocessed nuclei and myotube masks.",
    )

    train_pick = st.selectbox("Image to submit", names, key="train_pick")

    if st.button("📤 Submit for training", type="primary"):
        _ensure_dirs()
        ts = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        folder = CORRECTIONS_DIR / f"params_{ts}"
        folder.mkdir(parents=True, exist_ok=True)

        # Save the original image
        train_art = st.session_state.artifacts[train_pick]
        Image.fromarray(train_art["rgb_u8"]).save(folder / "image.png")

        # Save the postprocessed masks (from current parameter settings)
        nuc_pp_arr = train_art.get("nuc_pp_arr")
        myo_pp_arr = train_art.get("myo_pp_arr")
        if nuc_pp_arr is not None:
            Image.fromarray((nuc_pp_arr > 0).astype(np.uint8) * 255).save(folder / "nuclei_mask.png")
        if myo_pp_arr is not None:
            Image.fromarray((myo_pp_arr > 0).astype(np.uint8) * 255).save(folder / "myotube_mask.png")

        # Save the current parameter set — self_train can learn from these
        param_snapshot = {
            "reason": "user_optimized_params",
            "has_masks": nuc_pp_arr is not None and myo_pp_arr is not None,
            "timestamp": ts,
            "source_image": train_pick,
            "sample_id": train_art.get("sample_id", train_pick),
            "original_filename": train_art.get("original_filename", train_pick),
            "parameters": {
                "thr_nuc": float(thr_nuc),
                "thr_myo": float(thr_myo),
                "min_nuc_area": int(min_nuc_area),
                "min_myo_area": int(min_myo_area),
                "nuc_close_radius": int(nuc_close_radius),
                "myo_open_radius": int(myo_open_radius),
                "myo_erode_radius": int(myo_erode_radius),
                "min_myo_aspect_ratio": float(min_myo_aspect_ratio),
                "myo_max_area_px": int(myo_max_area_px),
                "myo_split_min_seeds": int(myo_split_min_seeds),
                "image_size": int(image_size),
            },
            "metrics": st.session_state.bio_metrics.get(train_pick, {}),
        }
        (folder / "meta.json").write_text(json.dumps(param_snapshot, indent=2, default=str))

        st.success(
            f"✅ **{train_pick}** submitted for training with your optimized parameters. "
            "The safe retraining script stores these settings for parameter review; only FIJI-corrected masks are used as ground-truth labels by default."
        )

# ─────────────────────────────────────────────────────────────────────────────
# ACTIVE LEARNING — CORRECTION UPLOAD
# ─────────────────────────────────────────────────────────────────────────────
if enable_al and not private_mode:
    st.subheader("🧠 Submit corrected labels (Active Learning)")
    flash_msg = st.session_state.pop("correction_flash", None)
    if flash_msg:
        st.success(flash_msg)
    st.caption(
        "Download predicted masks, correct them in ImageJ/FIJI, napari, or any "
        "image editor (white = foreground, black = background), then upload "
        "one corrected pair or many corrected pairs below."
    )

    # Active learning: correcting the images the model is least sure about
    # improves it fastest, so list them first.
    _unc = {n: st.session_state.bio_metrics.get(n, {}).get("model_uncertainty", 0) for n in names}
    names_by_unc = sorted(names, key=lambda n: -float(_unc.get(n) or 0))
    st.caption("Images are sorted by model uncertainty — correct the top ones first: "
               + ", ".join(f"{n} ({float(_unc[n] or 0):.2f})" for n in names_by_unc[:5]))
    al_pick = st.selectbox("Correct masks for image", names_by_unc, key="al_pick")
    if "correction_upload_reset_id" not in st.session_state:
        st.session_state.correction_upload_reset_id = 0

    # Download buttons for current masks — lets users edit and re-upload
    if al_pick in st.session_state.artifacts:
        al_art = st.session_state.artifacts[al_pick]
        dl1, dl2 = st.columns(2)
        with dl1:
            nuc_pp = al_art.get("nuc_pp_arr")
            if nuc_pp is not None:
                st.download_button(
                    "⬇️ Download nuclei mask",
                    data=png_bytes((nuc_pp * 255).astype(np.uint8)),
                    file_name=f"{al_pick}_nuclei_mask.png",
                    mime="image/png",
                    key="dl_nuc_mask",
                )
            else:
                st.download_button(
                    "⬇️ Download nuclei mask",
                    data=al_art["nuc_pp"],
                    file_name=f"{al_pick}_nuclei_mask.png",
                    mime="image/png",
                    key="dl_nuc_mask",
                )
        with dl2:
            myo_pp = al_art.get("myo_pp_arr")
            if myo_pp is not None:
                st.download_button(
                    "⬇️ Download myotube mask",
                    data=png_bytes((myo_pp * 255).astype(np.uint8)),
                    file_name=f"{al_pick}_myotube_mask.png",
                    mime="image/png",
                    key="dl_myo_mask",
                )
            else:
                st.download_button(
                    "⬇️ Download myotube mask",
                    data=al_art["myo_pp"],
                    file_name=f"{al_pick}_myotube_mask.png",
                    mime="image/png",
                    key="dl_myo_mask",
                )

    clear_col, _ = st.columns([1, 3])
    with clear_col:
        if st.button(
            "Remove uploaded masks",
            help=(
                "Clears masks currently sitting in the upload boxes. "
                "This does not delete corrections that were already submitted."
            ),
            key="clear_pending_correction_uploads",
        ):
            st.session_state.correction_upload_reset_id += 1
            st.rerun()

    upload_key_suffix = f"{al_pick}_{st.session_state.correction_upload_reset_id}"
    acol1, acol2 = st.columns(2)
    with acol1:
        corr_nuc = st.file_uploader(
            "Upload corrected NUCLEI mask(s)",
            type=["png", "tif", "tiff"],
            accept_multiple_files=True,
            key=f"nuc_corr_{upload_key_suffix}",
            help="Upload one nuclei mask for the selected image, or select many nuclei masks for batch correction.",
        )
    with acol2:
        corr_myo = st.file_uploader(
            "Upload corrected MYOTUBE mask(s)",
            type=["png", "tif", "tiff"],
            accept_multiple_files=True,
            key=f"myo_corr_{upload_key_suffix}",
            help="Upload the matching myotube mask(s). Batch files are paired by filename stem.",
        )

    if st.button("✅ Submit corrections", type="primary"):
        if not corr_nuc or not corr_myo:
            st.error("Please upload BOTH nuclei and myotube mask files.")
        else:
            n_submitted, correction_errors = submit_uploaded_correction_pairs(
                corr_nuc, corr_myo, st.session_state.artifacts, al_pick
            )
            if n_submitted == 1:
                st.session_state.correction_flash = "Saved 1 corrected image pair to corrections/."
            elif n_submitted > 1:
                st.session_state.correction_flash = f"Saved {n_submitted} corrected image pairs to corrections/."
            if correction_errors:
                st.warning("Some files were skipped:\n" + "\n".join(f"- {e}" for e in correction_errors[:12]))
            if not n_submitted and not correction_errors:
                st.error("No matching correction pairs were found.")
            if n_submitted:
                st.session_state.correction_upload_reset_id += 1
                st.rerun()

elif enable_al and private_mode:
    st.info(
        "🔒 Active learning and training submissions are disabled in Private mode. "
        "Toggle off Private mode in the sidebar to enable."
    )

st.divider()

# ─────────────────────────────────────────────────────────────────────────────
# RETRAINING QUEUE STATUS
# ─────────────────────────────────────────────────────────────────────────────
with st.expander("🔧 Self-training queue status"):
    _ensure_dirs()
    q_items = list(QUEUE_DIR.glob("*.json"))
    c_items = list(CORRECTIONS_DIR.glob("*/meta.json"))

    sq1, sq2 = st.columns(2)
    sq1.metric("Images in retraining queue", len(q_items))
    sq2.metric("Corrected label pairs",      len(c_items))

    if q_items:
        reasons = {}
        for p in q_items:
            try:
                r = json.loads(p.read_text()).get("reason", "unknown")
                reasons[r] = reasons.get(r, 0) + 1
            except Exception:
                pass
        st.write("Queue breakdown:", reasons)

    manifest = Path("manifest.json")
    if manifest.exists():
        try:
            history = json.loads(manifest.read_text())
            if history:
                st.markdown("**Last 5 retraining runs:**")
                hist_df = pd.DataFrame(history[-5:])
                st.dataframe(hist_df, use_container_width=True)
        except Exception:
            pass

    if st.button("🔄 Trigger retraining now"):
        import subprocess
        subprocess.Popen(["python", "self_train.py", "--manual"])
        st.info("Retraining started in the background. Check terminal / logs for progress.")