ui/workflow.py

# ui/workflow.py
from datetime import datetime
import streamlit as st
from PIL import Image, ImageOps
from streamlit_drawable_canvas import st_canvas
from core.drawing import (
    parse_boxes,  # from canvas JSON -> list[xyxy]
    boxes_to_canvas_json,  # (kept for compatibility; not used here)
    rects_only_json,
    get_rect_conf,
    inject_index_labels,
    seed_canvas_from_boxes,  # list[xyxy] -> canvas JSON
)
from core.state import set_defaults_from_preds, sync_samples_with_state
from core.exports import export_session  # not used here but retained
from core.detect_infer import DetConfig
import torch

# ---------------- helpers ----------------
def _auto_canvas_width():
    try:
        from streamlit_js_eval import get_page_info
        info = get_page_info() or {}
        w = int(info.get("clientWidth") or info.get("windowWidth") or 0)
        return 1200 if w >= 1800 else (1000 if w >= 1350 else 820)
    except Exception:
        return 820

def _det_cfg_from_state():
    c = st.session_state.get("det_cfg", {})
    return DetConfig(
        conf=float(c.get("conf", 0.25)),
        iou=float(c.get("iou", 0.50)),
        imgsz=int(c.get("imgsz", 640)),
        max_det=int(c.get("max_det", 300)),
        half=bool(c.get("half", True)),
    )

def _iou_xyxy(a, b, eps=1e-6):
    ax1, ay1, ax2, ay2 = map(float, a)
    bx1, by1, bx2, by2 = map(float, b)
    ix1, iy1 = max(ax1, bx1), max(ay1, by1)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
    inter = iw * ih
    aa = max(0.0, (ax2 - ax1)) * max(0.0, (ay2 - ay1))
    ba = max(0.0, (bx2 - bx1)) * max(0.0, (by2 - by1))
    union = aa + ba - inter + eps
    return inter / union

def _match_det_info_to_live_boxes(item, live_boxes, iou_thresh=0.90):
    """Align per-box detector info (conf/cls/name) to current live boxes."""
    base_boxes = item.get("boxes") or []
    info       = item.get("detector_info") or []
    out = []
    for b in live_boxes:
        best_iou, best_idx = 0.0, -1
        for j, bb in enumerate(base_boxes):
            i = _iou_xyxy(b, bb)
            if i > best_iou:
                best_iou, best_idx = i, j
        out.append(info[best_idx] if (best_iou >= iou_thresh and 0 <= best_idx < len(info)) else None)
    return out

def _reindex_det_info(old_boxes, old_info, new_boxes, iou_thresh=0.90):
    """Rebuild detector_info aligned to new_boxes by matching from (old_boxes, old_info)."""
    info_new = []
    for b in new_boxes:
        best_iou, best_idx = 0.0, -1
        for j, ob in enumerate(old_boxes):
            i = _iou_xyxy(b, ob)
            if i > best_iou:
                best_iou, best_idx = i, j
        info_new.append(old_info[best_idx] if (best_iou >= iou_thresh and 0 <= best_idx < len(old_info)) else None)
    return info_new

def _run_detection(item):
    det = st.session_state.get("detector")
    if det is None:
        return []
    cfg = _det_cfg_from_state()
    return det.predict_one(item["pil"], cfg) or []  # list of dicts {'bbox':[x1,y1,x2,y2], 'score','cls','name',...}

def _parse_canvas_into_boxes(item, scale, W, H):
    """Read current canvas JSON back into xyxy boxes on the original image grid."""
    if item.get("canvas_json"):
        item["boxes"] = parse_boxes(item["canvas_json"], scale, W, H)

def _classify_now(item, predict_fn):
    """Run classifier on current boxes, then sync to session samples."""
    if not item.get("boxes"):
        item["preds"] = []; item["user_labels"] = []; item["actions"] = []
        sync_samples_with_state(); return
    crops = [item["pil"].crop(b) for b in item["boxes"]]
    with torch.inference_mode():
        item["preds"] = predict_fn(crops, topk=3)
    # set_defaults fills current_label/confidence/margin/badges and seeds samples
    set_defaults_from_preds(item)
    item["actions"] = ["pending"] * len(item["boxes"])
    sync_samples_with_state()

def _bump_on_image_switch(ak: str, item: dict):
    """Force a canvas remount when the active image changes."""
    if st.session_state.get("_last_active_key") != ak:
        item["canvas_rev"] = (item.get("canvas_rev", 0) + 1)
        st.session_state["_last_active_key"] = ak

# ---------------- main UI ----------------
def render_workflow_tab(predict_fn, class_names, class_to_id, model_meta, export_root):
    CANVAS_W = _auto_canvas_width()
    if "uploader_rev" not in st.session_state:
        st.session_state.uploader_rev = 0
    
    # =============== (1) Load images ===============
    st.subheader("1) Load image(s)")
    ups = st.file_uploader(
        "Upload one or more images",
        type=["jpg", "jpeg", "png", "bmp"],
        accept_multiple_files=True,
        key=f"uploader_{st.session_state.uploader_rev}"  # resettable
    )
    for up in (ups or []):
        key = f"{up.name}-{up.size}"
        if key not in st.session_state.images:
            pil = Image.open(up).convert("RGB")
            st.session_state.images[key] = {
                "key": key, "name": up.name, "pil": pil,
                "boxes": [], "preds": [], "user_labels": [], "actions": [],
                "detector_info": [],  # <— conf/cls/name aligned with boxes
                "canvas_json": {"objects": []},
                "canvas_rev": 0,
            }
            if st.session_state.active_key is None:
                st.session_state.active_key = key

    keys = list(st.session_state.images.keys())
    if not keys:
        st.info("Upload images or load from S3/Drive in the sidebar.")
        return

    names = [st.session_state.images[k]["name"] for k in keys]
    idx   = keys.index(st.session_state.active_key) if st.session_state.active_key in keys else 0
    chosen = st.selectbox("Active image", names, index=idx)
    ak = keys[names.index(chosen)]
    st.session_state.active_key = ak

    # =============== (2) Canvas with Assist+ autorun ===============
    item = st.session_state.images[ak]
    base = item["pil"]
    disp_w = CANVAS_W
    scale  = float(disp_w) / float(base.width)
    disp_h = int(round(base.height * scale))
    _bump_on_image_switch(ak, item)

    # Assist+: auto-detect → seed canvas → auto-classify (only if no boxes yet)
    if st.session_state.assist_plus and not item.get("boxes") and not item.get("_skip_autodetect_once"):
        det_out = _run_detection(item)
        if det_out:
            boxes_xyxy = [tuple(int(v) for v in o["bbox"]) for o in det_out]
            det_info = []
            for o in det_out:
                conf  = o.get("score", o.get("conf", o.get("confidence")))
                clsid = o.get("cls", o.get("class"))
                name  = o.get("name")
                det_info.append({
                    "conf": float(conf) if conf is not None else None,
                    "cls":  int(clsid) if clsid is not None else None,
                    "name": name if name is not None else None,
                })
            item["boxes"] = boxes_xyxy
            item["detector_info"] = det_info
            item["canvas_json"] = seed_canvas_from_boxes(
                boxes_xyxy, scale, det_scores=[d["conf"] for d in det_info]
            )
            item["canvas_mode"] = "edit"
            _parse_canvas_into_boxes(item, scale, base.width, base.height)
            _classify_now(item, predict_fn)
            st.rerun()

    # Build display background
    bg_pil = ImageOps.exif_transpose(base).resize((disp_w, disp_h), Image.BILINEAR).convert("RGB").copy()

    # 2) Draw / edit rectangles
    st.subheader("2) Draw / edit rectangles")

    # Per-image canvas mode (so you can move/resize in Edit mode)
    if "canvas_mode" not in item:
        item["canvas_mode"] = "edit" if item.get("boxes") else "draw"

    mode_idx = 1 if item["canvas_mode"] == "edit" else 0
    canvas_mode_choice = st.radio(
        "Canvas mode",
        ["Draw", "Edit"],
        index=mode_idx,
        horizontal=True,
        key=f"canvas_mode_{ak}_{item.get('canvas_rev', 0)}"
    )
    item["canvas_mode"] = "edit" if canvas_mode_choice == "Edit" else "draw"
    drawing_mode = "transform" if item["canvas_mode"] == "edit" else "rect"

    # Hints
    if item["canvas_mode"] == "edit":
        st.caption("Tip: Click a box to move/resize. Use the live list below to remove a box.")
    else:
        st.caption("Tip: Drag on the image to draw new boxes.")

    # ---------- initial drawing ----------
    base_json = item.get("canvas_json") or {"objects": []}
    show_labels = st.checkbox("Show box labels on canvas", value=True, key=f"show_labels_{ak}")
    display_json = inject_index_labels(base_json, show_conf=False) if show_labels else base_json

    # ---------- Canvas ----------
    canvas = st_canvas(
        fill_color="rgba(0,0,0,0)",
        stroke_width=3,
        stroke_color="#FF9900",
        background_image=bg_pil,
        update_streamlit=True,
        height=disp_h,
        width=disp_w,
        drawing_mode=drawing_mode,
        display_toolbar=True,
        initial_drawing=display_json,
        key=f"canvas_{ak}_{item.get('canvas_rev', 0)}_{disp_w}x{disp_h}"
    )

    # ---------- Live canvas snapshot ----------
    live_json = canvas.json_data or base_json
    objs = list(live_json.get("objects") or [])

    # Rect indices & objects (ignore overlay textboxes)
    rect_indices = [k for k, o in enumerate(objs) if o.get("type") == "rect"]
    rect_objs    = [objs[k] for k in rect_indices]

    # Live boxes in original image coords
    live_boxes = parse_boxes(live_json, scale, base.width, base.height)

    # Detector summary (use detector_info matched to live boxes)
    infos = _match_det_info_to_live_boxes(item, live_boxes)
    confs = [d["conf"] for d in infos if d and d.get("conf") is not None]
    if confs:
        n = len(confs)
        cmin, cmax = min(confs), max(confs)
        cmean = sum(confs)/n
        th = float(st.session_state.get("det_cfg", {}).get("conf", 0.25))
        below = sum(1 for c in confs if c < th)
        st.markdown(
            f"**Active image (detector)** — boxes: {n} · mean **{cmean:.3f}** · "
            f"min **{cmin:.3f}** · max **{cmax:.3f}** · below {th:.2f}: **{below}/{n}**"
        )

    # Commit canvas into boxes (Clear / Remove)
    col_clear, col_remove = st.columns([1, 1])

    if col_clear.button("Clear boxes (THIS image)", key=f"clear_{ak}"):
        item["canvas_json"] = {"objects": []}
        item["boxes"] = []; item["preds"] = []; item["user_labels"] = []; item["actions"] = []
        item["detector_info"] = []  # also clear detector stats
        item["canvas_mode"] = "draw"
        item["canvas_rev"]  = (item.get("canvas_rev", 0) + 1)
        item["_skip_autodetect_once"] = True
        st.session_state.all_samples = [s for s in st.session_state.all_samples if s.image_key != ak]
        st.rerun()

    if col_remove.button("Remove image", key=f"remove_{ak}"):
        remaining_keys = [k for k in st.session_state.images.keys() if k != ak]
        next_key = remaining_keys[0] if remaining_keys else None
        del st.session_state.images[ak]
        st.session_state.all_samples = [s for s in st.session_state.all_samples if s.image_key != ak]
        st.session_state.uploader_rev += 1
        st.session_state.active_key = next_key
        if next_key:
            nxt = st.session_state.images[next_key]
            nxt["_skip_autodetect_once"] = True
            nxt["canvas_mode"] = "edit" if nxt.get("boxes") else "draw"
            nxt["canvas_rev"]  = nxt.get("canvas_rev", 0) + 1
        st.rerun()

    # =============== (3) Detect / Classify controls ===============
    st.subheader("3) Detect & Classify")

    assist_on = bool(st.session_state.get("assist_plus"))
    has_boxes = bool(item.get("boxes"))

    canvas_objects = (canvas.json_data or {}).get("objects") or []
    has_canvas_rects = any(obj.get("type") == "rect" for obj in canvas_objects)
    can_classify = has_boxes or has_canvas_rects

    # ---------- DETECT ----------
    if assist_on:
        st.caption("Detect: find objects and write boxes onto the canvas.")
        det_mode = st.radio(
            "Detection mode",
            ["Replace current boxes", "Append to existing"],
            index=0,
            horizontal=True,
            key=f"det_mode_{ak}",
            help="Replace overwrites all boxes; Append merges new detections."
        )
        det_label = "Detect objects (active image)" if not has_boxes else "Run detection (active image)"
        if st.button(
            det_label,
            help="Run detection with the mode selected above.",
            key=f"detect_{ak}"
        ):
            det_out = _run_detection(item)
            new_boxes = [tuple(int(v) for v in o["bbox"]) for o in det_out]
            new_info  = []
            for o in det_out:
                conf  = o.get("score", o.get("conf", o.get("confidence")))
                clsid = o.get("cls", o.get("class"))
                name  = o.get("name")
                new_info.append({
                    "conf": float(conf) if conf is not None else None,
                    "cls":  int(clsid) if clsid is not None else None,
                    "name": name if name is not None else None,
                })
            if new_boxes:
                old_boxes = list(item.get("boxes") or [])
                old_info  = list(item.get("detector_info") or [])
                if det_mode.startswith("Replace"):
                    merged_boxes = new_boxes
                    merged_info  = new_info
                else:
                    merged_boxes = old_boxes + new_boxes
                    merged_info  = old_info + new_info

                item["boxes"] = merged_boxes
                item["detector_info"] = merged_info
                item["canvas_json"] = seed_canvas_from_boxes(
                    merged_boxes, scale, det_scores=[d["conf"] for d in merged_info]
                )
                item["canvas_mode"] = "edit"
                item["canvas_rev"] = (item.get("canvas_rev", 0) + 1)

                # auto-run classification so users don't need a second click after detection
                _parse_canvas_into_boxes(item, scale, base.width, base.height)
                _classify_now(item, predict_fn)
                st.rerun()
            else:
                st.info("No objects detected. Try adjusting thresholds or draw boxes.")

    # ---------- CLASSIFY ----------
    st.caption("Classify: run the classifier on the current canvas boxes.")
    has_preds = bool(item.get("preds"))
    boxes_changed = live_boxes != (item.get("boxes") or [])
    classify_label = "Run classification" if not has_preds else ("Update after edits" if boxes_changed else "Re-run classification")
    classify_help = "Add or detect a box first." if not can_classify else "Run the classifier on the current canvas boxes."

    if st.button(
        classify_label,
        disabled=not can_classify,
        help=classify_help,
        key=f"classify_{ak}"
    ):
        item["canvas_json"] = canvas.json_data or item.get("canvas_json") or {}
        _parse_canvas_into_boxes(item, scale, base.width, base.height)
        _classify_now(item, predict_fn)
        st.rerun()

    # ---------- Crops (single list for remove + relabel + stats) ----------
    st.subheader("Crops")
    st.caption(f"Crops in this image: **{len(live_boxes)}**")

    samples = [s for s in st.session_state.all_samples if s.image_key == ak]
    samples_by_idx = {s.crop_idx: s for s in samples}

    if not live_boxes:
        st.info("Draw boxes or run detection to see crops.")
    else:
        for i, b in enumerate(live_boxes):
            info = infos[i] if i < len(infos) else None
            det_conf = (info or {}).get("conf", None)
            det_label = None
            if info:
                if info.get("name") is not None:
                    det_label = str(info["name"])
                elif info.get("cls") is not None:
                    names = st.session_state.get("detector_names") or []
                    cid = info["cls"]
                    if isinstance(names, (list, tuple)) and isinstance(cid, int) and 0 <= cid < len(names):
                        det_label = str(names[cid])
                    else:
                        det_label = str(cid)

            sample = samples_by_idx.get(i)
            col_img, col_meta, col_actions = st.columns([1.2, 2.6, 1.3])

            with col_img:
                try:
                    crop_image = base.crop(b)
                    st.image(crop_image, caption=f"Crop {i+1}", width=180)
                except Exception:
                    st.caption(f"Crop {i+1}")

            with col_meta:
                st.write(f"Box #{i+1}  [x0={b[0]}, y0={b[1]}, x1={b[2]}, y1={b[3]}]")
                if det_conf is None and not det_label:
                    st.caption("det: —")
                else:
                    th = float(st.session_state.get("det_cfg", {}).get("conf", 0.25))
                    status = " ⚠️" if (det_conf is not None and det_conf < th) else ""
                    parts = []
                    if det_conf is not None: parts.append(f"{det_conf:.2f}")
                    if det_label: parts.append(det_label)
                    st.caption("det: " + " · ".join(parts) + status)

                if sample:
                    clf_top_name, clf_top_prob = (sample.top3[0] if sample.top3 else (sample.current_label, sample.confidence))
                    if det_conf is not None and clf_top_prob is not None:
                        labels_match = bool(det_label and clf_top_name and str(det_label) == str(clf_top_name))
                        det_flag = ":red[low det]" if det_conf < th else ":green[det ok]"
                        lbl_flag = ":green[label match]" if labels_match else (":orange[label diff]" if det_label else ":gray[det unlabeled]")
                        st.markdown(
                            f"**Det vs Clf:** det **{det_conf:.2f}**{f' ({det_label})' if det_label else ''} "
                            f"vs clf **{clf_top_prob:.3f}** ({clf_top_name}) · {det_flag} · {lbl_flag}"
                        )
                    top3_str = ", ".join([f"{n} ({p:.3f})" for n, p in sample.top3[:3]])
                    st.write(f"**Current Label:** {sample.current_label}")
                    st.write(f"**Top-3:** {top3_str}")
                    st.write(f"**Confidence:** {sample.confidence:.3f}")
                    st.write(f"**Margin:** {sample.margin:.3f}")
                    badge_str = " ".join([f":red[{b}]" for b in sample.badges]) if sample.badges else ":green[OK]"
                    st.write(f"**Badges:** {badge_str}")
                    st.write(f"**Action Status:** {(':green[' if sample.action=='accepted' else ':red[') + sample.action.upper() + ']'}")
                else:
                    st.caption("Not classified yet.")

            with col_actions:
                if sample:
                    try:
                        idx = class_names.index(sample.current_label) if sample.current_label in class_names else 0
                    except Exception:
                        idx = 0
                    new_label = st.selectbox(
                        f"Relabel Crop {sample.crop_idx+1}",
                        class_names,
                        index=idx,
                        key=f"wf_relabel_{ak}_{i}",
                    )
                    if new_label != sample.current_label:
                        sample.user_label = new_label
                        sample.action = "relabel"
                        from core.state import update_image_state_from_samples
                        update_image_state_from_samples()
                        st.rerun()
                else:
                    st.caption("Classify before relabeling.")

                if st.button("Remove", key=f"rm_crop_{ak}_{i}"):
                    true_idx = rect_indices[i] if i < len(rect_indices) else None
                    if true_idx is not None and 0 <= true_idx < len(objs):
                        old_boxes = item.get("boxes") or []
                        old_info  = item.get("detector_info") or []

                        objs.pop(true_idx)
                        if true_idx < len(objs) and objs[true_idx].get("type") == "textbox":
                            objs.pop(true_idx)

                        live_json["objects"] = objs
                        item["canvas_json"] = live_json
                        _parse_canvas_into_boxes(item, scale, base.width, base.height)

                        new_boxes = item.get("boxes") or []
                        item["detector_info"] = _reindex_det_info(old_boxes, old_info, new_boxes)

                        item["canvas_mode"] = "edit"
                        item["canvas_rev"]  = (item.get("canvas_rev", 0) + 1)
                        item["preds"] = []; item["user_labels"] = []; item["actions"] = []
                        sync_samples_with_state()
                        st.rerun()

            st.divider()