app.py

try:
    import spaces  # noqa: E402  must precede any CUDA-initializing import (ZeroGPU)
except ImportError:  # local dev without the spaces runtime
    spaces = None

import json
import os
import shutil
import subprocess
import time

import cv2
import gradio as gr
import numpy as np
import pandas as pd
import torch

from PIL import Image as PILImage

# ---------------------------------------------------------------------------
# AnyTraverse API surface
# ---------------------------------------------------------------------------
ANYTRAVERSE_AVAILABLE = False
try:
    from anytraverse import build_pipeline_from_paper
    from anytraverse.utils.state import TraversalState

    ANYTRAVERSE_AVAILABLE = True
except Exception as _e:  # pragma: no cover - depends on environment
    print(f"[app] anytraverse not importable ({_e}); running in SIMULATION mode.")

    class TraversalState:
        """Stand-in so the dashboard is testable without the package."""
        OK = object()
        UNKNOWN_SCENE = object()
        UNKOWN_OBJ = object()  # source spelling (missing N) kept for parity


# Human-readable HOC labels requested by the user.
HOC_LABELS = {
    TraversalState.OK: "ok",
    TraversalState.UNKNOWN_SCENE: "unknown_scene",
    TraversalState.UNKOWN_OBJ: "unknown_object",
}


def get_vlm_device():
    if torch.cuda.is_available():
        return f"CUDA:0 ({torch.cuda.get_device_name(0)})"
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return "Apple MPS"
    return "CPU"


# ---------------------------------------------------------------------------
# Small helpers. NOTE: all maps / grids are kept in BGR internally and only
# converted to RGB at the very end, so the raw image is not color-swapped.
# ---------------------------------------------------------------------------
def to_numpy(val):
    if isinstance(val, torch.Tensor):
        return val.detach().cpu().numpy()
    if isinstance(val, (list, tuple)):
        return np.asarray(val[0])
    return np.asarray(val)


def to_float(val, default=0.0):
    if val is None:
        return default
    if isinstance(val, torch.Tensor):
        return float(val.detach().cpu().item())
    return float(val)


def colorize(arr, target_w, target_h, colormap=cv2.COLORMAP_INFERNO):
    """Normalize a 2D map and apply a color map, resized to target dims (BGR)."""
    arr = to_numpy(arr)
    if arr.ndim == 3:
        arr = arr.reshape(arr.shape[-2:])
    if arr.size == 0:
        arr = np.zeros((2, 2))
    lo, hi = float(arr.min()), float(arr.max())
    if hi - lo < 1e-9:
        norm = np.zeros(arr.shape, dtype=np.uint8)
    else:
        norm = ((arr - lo) / (hi - lo) * 255.0).astype(np.uint8)
    return cv2.resize(cv2.applyColorMap(norm, colormap), (int(target_w), int(target_h)))


def add_caption(img_bgr, text):
    cv2.putText(img_bgr, str(text), (6, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                (255, 255, 255), 2)
    cv2.putText(img_bgr, str(text), (6, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
                (0, 0, 0), 1)
    return img_bgr


def _ffmpeg_bin():
    """Locate an ffmpeg binary: bundled (imageio-ffmpeg) first, else system."""
    try:
        import imageio_ffmpeg
        return imageio_ffmpeg.get_ffmpeg_exe()
    except Exception:
        pass
    return shutil.which("ffmpeg") or "ffmpeg"


def convert_to_h264(in_path, out_path):
    """FFmpeg wrapper producing a browser-playable H.264 video (no audio)."""
    if not in_path or not os.path.exists(in_path):
        return None
    try:
        subprocess.run(
            [_ffmpeg_bin(), "-y", "-i", in_path, "-vcodec", "libx264",
             "-pix_fmt", "yuv420p", "-preset", "fast", "-an", out_path],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True,
        )
        return out_path
    except Exception:
        return None


def bounds_box(rx_min, rx_max, ry_min, ry_max, w, h):
    return ((int(rx_min * w), int(ry_min * h)), (int(rx_max * w), int(ry_max * h)))


# ---------------------------------------------------------------------------
# Session: shared mutable state read by the streaming generator + button events
# ---------------------------------------------------------------------------
class AppSession:
    def __init__(self):
        self.pipeline = None
        self.cap = None
        self.writer = None
        self.video_path = None
        self.frame_idx = 0
        self.fps = 30
        self.vw = 0
        self.vh = 0
        self.is_paused = False
        self.resume_requested = False
        self.simulate_hoc_requested = False
        self.is_running = False
        self.last_grid = None
        self.last_attn = None
        self.traversal = TraversalState.OK
        self.preferences = {"road": 1.0, "grass": 0.5, "bush": -0.8, "rock": -0.6}
        self.uncert_thresh = 0.4
        self.sim_thresh = 0.8
        self.skip = 2
        self.telemetry = []
        self.raw_out = "raw_opencv_temp.mp4"
        self.h264_out = "anytraverse_h264_output.mp4"
        self.device = get_vlm_device()


session = AppSession()

TEL_COLUMNS = ["Frame", "ROI Trav", "ROI Unc", "Ref Sim", "State"]
EMPTY_DF = pd.DataFrame(columns=TEL_COLUMNS)
PLOT_COLUMNS = ["Frame", "ROI Trav", "ROI Unc", "Uncert Thresh"]


# ---------------------------------------------------------------------------
# Frame-state unpacking / simulation
# ---------------------------------------------------------------------------
def unpack_state(state_obj, bgr):
    """Map a real anytraverse AnyTraverseState to the dashboard's shado dict."""
    prompts = list(state_obj.traversability_preferences.keys())
    attn_maps = [(p, m) for p, m in zip(prompts, list(state_obj.attention_maps))]
    return {
        "raw_bgr": bgr,
        "roi_bbox": state_obj.roi_bbox,
        "trav": to_numpy(state_obj.traversability_map),
        "uncert": to_numpy(state_obj.uncertainty_map),
        "attn_maps": attn_maps,
        "roi_trav": to_float(state_obj.roi_traversability),
        "roi_uncert": to_float(state_obj.roi_uncertainty),
        "sim": to_float(state_obj.ref_scene_similarity),
        "state": state_obj.traversal_state,
    }


def _simulate_state(bgr, prefs, uncert_thresh, frame_idx):
    """Deterministic fake state so the UI is testable without the package."""
    h, w, _ = bgr.shape
    roi_u = float(np.clip(0.18 + 0.45 * np.sin(frame_idx / 7.0), 0.0, 1.0))
    trav = float(np.clip(0.65 + 0.35 * np.sin(frame_idx / 9.0), 0.05, 0.95))
    sim = float(np.clip(0.95 - frame_idx * 0.002, 0.2, 1.0))

    att = []
    phase = np.linspace(0, np.pi, w, dtype=np.float32)
    for k, p in enumerate(prefs.keys()):
        base = np.full((h, w), 0.5, dtype=np.float32)
        base[h // 2:, :] += (0.25 * np.sin(phase + k))[None, :]
        base[0:h // 2, :] = 0.9
        att.append((p, base))

    box = bounds_box(0.333, 0.667, 0.6, 0.95, w, h)
    if roi_u > uncert_thresh:
        st = TraversalState.UNKOWN_OBJ
    elif (frame_idx // 40) % 4 == 2:
        st = TraversalState.UNKNOWN_SCENE
    else:
        st = TraversalState.OK

    return {
        "raw_bgr": bgr, "roi_bbox": box,
        "trav": np.full((h, w), trav, dtype=np.float32),
        "uncert": np.full((h, w), roi_u, dtype=np.float32),
        "attn_maps": att, "roi_trav": trav, "roi_uncert": roi_u,
        "sim": sim, "state": st,
    }


def seed_state(bgr, box):  # pragma: no cover - kept unused (records placeholder)
    return None


# ---------------------------------------------------------------------------
# Rendering (no matplotlib anywhere)
# ---------------------------------------------------------------------------
def build_grid(p):
    """2x2 grid: [raw+ROI | traversability] / [uncertainty | ROI crop]."""
    h, w, _ = p["raw_bgr"].shape
    raw = p["raw_bgr"].copy()
    (x0, y0), (x1, y1) = p["roi_bbox"]
    cv2.rectangle(raw, (x0, y0), (x1, y1), (0, 255, 255), 2)  # BGR yellow

    trav_img = colorize(p["trav"], w, h)
    uncert_img = colorize(p["uncert"], w, h)

    xa, xb = max(x0, 0), min(x1, w)
    ya, yb = max(y0, 0), min(y1, h)
    roi_crop = raw[ya:yb + 1, xa:xb + 1]
    if roi_crop.size == 0:
        roi_crop = raw
    roi_crop = cv2.resize(roi_crop, (w, h))

    row1 = np.hstack([raw, trav_img])
    row2 = np.hstack([uncert_img, roi_crop])
    grid = np.vstack([row1, row2])
    return cv2.cvtColor(grid, cv2.COLOR_BGR2RGB)


def build_attn_strip(p):
    """Prompt attention maps as a grid: max 5 columns, wrapping to new rows."""
    raw = p["raw_bgr"]
    h, w, _ = raw.shape
    att = p["attn_maps"]
    if not att:
        return np.zeros((h, w, 3), dtype=np.uint8)
    cell_w = max(int(w // min(len(att), 5)), 80)
    blank = np.zeros((h, cell_w, 3), dtype=np.uint8)
    rows = []
    for i in range(0, len(att), 5):
        cells = [add_caption(colorize(m, cell_w, h).copy(), name)
                 for name, m in att[i:i + 5]]
        while len(cells) < 5:
            cells.append(blank.copy())
        rows.append(np.hstack(cells))
    strip = rows[0] if len(rows) == 1 else np.vstack(rows)
    return cv2.cvtColor(strip, cv2.COLOR_BGR2RGB)


def bars_html(trav, unc, thresh):
    """Two horizontal 0..1 gauge bars (pure HTML/CSS, no matplotlib)."""
    t = int(round(max(0.0, min(1.0, trav)) * 100))
    u = int(round(max(0.0, min(1.0, unc)) * 100))
    th = max(0.0, min(1.0, thresh)) * 100
    return (
        f"<div class='prog'><div style='display:flex;justify-content:space-between'>"
        f"<span style='font-weight:600'>ROI Traversability</span><span>{trav:.3f}</span></div>"
        f"<div style='position:relative;height:16px;background:#e9ecef;border-radius:8px;border:1px solid #ced4da'>"
        f"<div style='position:absolute;left:0;top:0;height:100%;width:{t}%;background:#2ca02c;border-radius:8px'></div></div></div>"
        f"<div class='bar' style='margin-top:10px'><div style='display:flex;justify-content:space-between'>"
        f"<span style='font-weight:600'>ROI Uncertainty</span><span>{unc:.3f}</span></div>"
        f"<div style='position:relative;height:16px;background:#e9ecef;border-radius:8px;border:1px solid #ced4da'>"
        f"<div style='position:absolute;left:0;top:0;height:100%;width:{u}%;background:#d62728;border-radius:8px'></div>"
        f"<div title='threshold' style='position:absolute;left:{th}%;top:-3px;bottom:-3px;width:2px;background:#343a40'></div>"
        f"</div></div>")


def lineplot_df():
    if not session.telemetry:
        return pd.DataFrame(columns=PLOT_COLUMNS)
    rows = []
    for t in session.telemetry:
        rows.append({"Frame": t["frame"], "ROI Trav": t["roi_trav"],
                     "ROI Unc": t["roi_uncert"], "Uncert Thresh": session.uncert_thresh})
    return pd.DataFrame(rows, columns=PLOT_COLUMNS)


def df_table():
    if not session.telemetry:
        return EMPTY_DF
    return pd.DataFrame(
        [{"Frame": t["frame"], "ROI Trav": t["roi_trav"], "ROI Unc": t["roi_uncert"],
          "Ref Sim": t["sim"], "State": t["state"]} for t in session.telemetry]
    )


# Order of the generator outputs (must mirror the `outputs` list).
def render(grid, status, op_visible, attn, m_frame, m_skip, m_state, m_trav,
           m_unc, m_sim, m_fps, m_lat, plot, bars, table, video=None):
    return (grid, status, gr.update(visible=op_visible), attn, str(m_frame),
            str(m_skip), str(m_state), f"{m_trav:.3f}", f"{m_unc:.3f}",
            f"{m_sim:.3f}", str(m_fps), f"{m_lat} ms", plot, bars, table, video)


def initial_render(msg):
    return (None, msg, gr.update(visible=False), None, "0", str(session.skip),
            "ok", "0.000", "0.000", "0.000", "0", "0 ms",
            lineplot_df(), bars_html(0.0, 0.0, session.uncert_thresh),
            EMPTY_DF, None)


# ---------------------------------------------------------------------------
# ZeroGPU worker(s). All CUDA work must live in a @spaces.GPU function; the
# decorator is a no-op when the `spaces` runtime is absent (local dev).
# The pipeline as a singleton is cached in `_VLM` so consecutive frames reuse
# the loaded model inside the GPU context.
# ---------------------------------------------------------------------------
_VLM = {"pipe": None}


def _pipeline_ready():
    return _VLM["pipe"] is not None


def _set_pipe():
    session.pipeline = _VLM["pipe"]


def _vlm_make_pipe(prefs, sim_thresh, uncert_thresh, rx, ry):
    return build_pipeline_from_paper(
        init_traversabilty_preferences=prefs,
        ref_scene_similarity_threshold=float(sim_thresh),
        roi_uncertainty_threshold=float(uncert_thresh),
        roi_x_bounds=(float(rx[0]), float(rx[1])),
        roi_y_bounds=(float(ry[0]), float(ry[1])),
    )


def _vlm_step_raw(frame_bgr, prefs, sim_thresh, uncert_thresh, rx, ry):
    if not _pipeline_ready():
        _VLM["pipe"] = _vlm_make_pipe(prefs, sim_thresh, uncert_thresh, rx, ry)
    st = _VLM["pipe"].step(image=PILImage.fromarray(cv2.cvtColor(frame_bgr,
                                                               cv2.COLOR_BGR2RGB)))
    return unpack_state(st, frame_bgr)


def _vlm_human_raw(text):
    if not _pipeline_ready():
        return None
    _VLM["pipe"].human_call(human_input=text)
    return dict(_VLM["pipe"].traversability_preferences)


def _vlm_register_raw():
    if not _pipeline_ready():
        return None
    _VLM["pipe"].register_scene()
    return None


def _gpu_decorate(fn):
    if spaces is not None:
        return spaces.GPU(duration=120)(fn)
    return fn


vlm_step = _gpu_decorate(_vlm_step_raw)
vlm_human = _gpu_decorate(_vlm_human_raw)
vlm_register = _gpu_decorate(_vlm_register_raw)


# ---------------------------------------------------------------------------
# Main streaming worker. Restarted on "Go / Reset" and on "Resume".
# ---------------------------------------------------------------------------
def run_evaluation(video_file, pref_json, sim_thresh, uncert_thresh,
                   rx_min, rx_max, ry_min, ry_max, frame_skip):
    if session.is_running and not session.is_paused and not session.resume_requested:
        yield initial_render("⏳ A live evaluation is already running.")
        return


    fresh = not session.resume_requested
    if fresh:
        session.telemetry = []
        session.uncert_thresh = float(uncert_thresh)
        session.sim_thresh = float(sim_thresh)
        session.skip = int(frame_skip) if frame_skip else 1
        session.is_paused = False
        session.simulate_hoc_requested = False
        session.is_running = True
        session.resume_requested = False

    if video_file:
        session.video_path = (video_file if isinstance(video_file, str)
                              else getattr(video_file, "name", str(video_file)))

    if fresh:
        try:
            session.preferences = json.loads(pref_json) or session.preferences
        except Exception:
            pass
        if ANYTRAVERSE_AVAILABLE:
            yield initial_render(
                "🔄 Building AnyTraverse pipeline (first run may download models)…")
        else:
            session.pipeline = None

    if not session.video_path:
        session.is_running = False
        yield initial_render("⚠ Please upload a video first.")
        return

    if session.cap is None or not session.cap.isOpened():
        session.cap = cv2.VideoCapture(session.video_path)
        session.fps = int(session.cap.get(cv2.CAP_PROP_FPS)) or 30
        session.vw = int(session.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        session.vh = int(session.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        if session.vw == 0 or session.vh == 0:
            session.is_running = False
            session.cap = None
            yield initial_render("❌ Could not read the uploaded video file.")
            return
        session.writer = cv2.VideoWriter(session.raw_out,
                                         cv2.VideoWriter_fourcc(*"mp4v"),
                                         session.fps, (session.vw * 2, session.vh * 2))
        session.frame_idx = 0

    skip = session.skip

    try:
        while session.cap.isOpened():
            # -------- operator pause handling -------------------------------------
            if session.is_paused:
                if session.resume_requested:
                    session.resume_requested = False
                    session.is_paused = False
                else:
                    last = session.telemetry[-1] if session.telemetry else None
                    yield render(
                        session.last_grid, f"🚨 **HALTED at frame {session.frame_idx}** "
                        f"— traversal state **`{HOC_LABELS.get(session.traversal,'?')}`**. "
                        "Enter a τ update (or `ok`) and press **Resume**.",
                        True, session.last_attn, session.frame_idx, skip,
                        HOC_LABELS.get(session.traversal, "?"),
                        last["roi_trav"] if last else 0.0,
                        last["roi_uncert"] if last else 0.0,
                        last["sim"] if last else 0.0, 0, 0,
                        lineplot_df(), bars_html(
                            last["roi_trav"] if last else 0.0,
                            last["roi_uncert"] if last else 0.0,
                            session.uncert_thresh), df_table(), video=None)
                    return

            # -------- simulated operator call (manual test trigger) ---------------
            if session.simulate_hoc_requested:
                session.simulate_hoc_requested = False
                session.is_paused = True
                session.traversal = TraversalState.UNKOWN_OBJ
                last = session.telemetry[-1] if session.telemetry else None
                yield render(
                    session.last_grid,
                    "🚨 **SIMULATED HUMAN-OPERATOR-CALL** — live loop paused. "
                    "Provide a τ update (or you can resume) and press **Resume**.",
                    True, session.last_attn, session.frame_idx,
                    skip, "unknown_object",
                    last["roi_trav"] if last else 0.0,
                    last["roi_uncert"] if last else 0.0,
                    last["sim"] if last else 0.0, 0, 0,
                    lineplot_df(), bars_html(
                        last["roi_trav"] if last else 0.0,
                        last["roi_uncert"] if last else 0.0,
                        session.uncert_thresh), df_table(), video=None)
                return

            # -------- read the next display frame ------------------------------
            t0 = time.time()
            ret, frame_bgr = session.cap.read()
            if not ret:
                break
            session.frame_idx += 1
            # Frame skipping: pass every k-th frame to anytraverse; the frames in
            # between are never shown in the UI nor written to the output video.
            if skip > 1 and (session.frame_idx - 1) % skip != 0:
                continue

            if ANYTRAVERSE_AVAILABLE:
                p = vlm_step(frame_bgr, session.preferences, session.sim_thresh,
                             session.uncert_thresh,
                             (float(rx_min), float(rx_max)),
                             (float(ry_min), float(ry_max)))
                _set_pipe()
            else:
                p = _simulate_state(frame_bgr, session.preferences,
                                    session.uncert_thresh, session.frame_idx)
            session.traversal = p["state"]

            fps = round(1.0 / max(time.time() - t0, 1e-3), 1)
            lat = round((time.time() - t0) * 1000, 1)
            lbl = HOC_LABELS.get(p["state"], "ok")

            # telemetry / chart row
            session.telemetry.append({
                "frame": session.frame_idx,
                "roi_trav": round(float(p["roi_trav"]), 4),
                "roi_uncert": round(float(p["roi_uncert"]), 4),
                "sim": round(float(p["sim"]), 4),
                "state": lbl,
            })

            grid = build_grid(p)
            attn = build_attn_strip(p)
            session.last_grid = grid
            session.last_attn = attn
            if session.writer is not None:
                session.writer.write(cv2.cvtColor(grid, cv2.COLOR_RGB2BGR))

            status = f"Frame {session.frame_idx} · state **`{lbl}`**"

            yield render(
                grid, status, False, attn, session.frame_idx, skip, lbl,
                p["roi_trav"], p["roi_uncert"], p["sim"], fps, lat,
                lineplot_df(), bars_html(p["roi_trav"], p["roi_uncert"],
                                        session.uncert_thresh), df_table(),
                video=None)

            if lbl != "ok":
                session.is_paused = True
                yield render(
                    grid, f"🚨 **HOC TRIGGERED at frame {session.frame_idx}** → "
                    f"**`{lbl}`**. Provide a τ update (or just `ok`) and **Resume**.",
                    True, attn, session.frame_idx, skip, lbl, p["roi_trav"],
                    p["roi_uncert"], p["sim"], fps, lat, lineplot_df(),
                    bars_html(p["roi_trav"], p["roi_uncert"],
                              session.uncert_thresh), df_table(), video=None)
                return

        # -------- normal completion -------------------------------------------------
        if session.writer is not None:
            session.writer.release()
            session.writer = None
        final_video = convert_to_h264(session.raw_out, session.h264_out)
        last = session.telemetry[-1] if session.telemetry else None
        yield render(
            session.last_grid, "🎉 **Evaluation complete.** Download the composed video below.",
            False, session.last_attn, session.frame_idx, skip,
            last["state"] if last else "ok",
            last["roi_trav"] if last else 0.0,
            last["roi_uncert"] if last else 0.0,
            last["sim"] if last else 0.0, 0, 0, lineplot_df(),
            bars_html(last["roi_trav"] if last else 0.0,
                      last["roi_uncert"] if last else 0.0, session.uncert_thresh),
            df_table(), video=final_video)
    finally:
        if not session.is_paused:
            if session.writer is not None:
                session.writer.release()
                session.writer = None
            if session.cap is not None:
                session.cap.release()
                session.cap = None
            session.is_running = False


# ---------------------------------------------------------------------------
# Operator intervention (Resume), Simulate-HOC, and live pipeline updates
# ---------------------------------------------------------------------------
def handle_operator_resume(operator_text):
    text = (operator_text or "").strip()
    if session.pipeline is not None:
        if text and text.lower() != "ok":
            prefs = vlm_human(text)
            if prefs:
                session.preferences = prefs
                msg = f"✅ Applied operator τ update `{text}` — resuming."
            else:
                msg = f"✅ (pipeline not built yet) resuming with `{text}`."
        else:
            vlm_register()
            msg = "✅ Scene registered (no τ change) — resuming."
    else:
        if text and text.lower() != "ok":
            try:
                for pw in text.split(";"):
                    if ":" in pw:
                        k, v = pw.split(":", 1)
                        session.preferences[k.strip()] = float(v)
            except Exception:
                pass
        msg = "✅ (Simulation) resuming."
    session.resume_requested = True
    session.is_paused = False
    session.simulate_hoc_requested = False
    return msg, gr.update(visible=False), json.dumps(session.preferences, indent=2)


def simulate_hoc():
    session.simulate_hoc_requested = True
    session.is_paused = False
    session.resume_requested = False
    return "⏸ Simulate-HOC requested — the live loop will pause on its next frame."


def live_pipeline_update(sim, unc, rxmin, rxmax, rymin, rymax):
    """Apply threshold / ROI changes to the running pipeline object on the fly."""
    pipe = _VLM["pipe"]
    if pipe is not None:
        try:
            pipe._threshold.ref_scene_similarity = float(sim)
            pipe._threshold.roi_uncertainty = float(unc)
            pipe._roi._x_bounds = (float(rxmin), float(rxmax))
            pipe._roi._y_bounds = (float(rymin), float(rymax))
        except Exception:
            return "⚠ live update failed"
    session.uncert_thresh = float(unc)
    session.sim_thresh = float(sim)
    return (f"Live cfg: sim={float(sim):.2f}, unc={float(unc):.2f}, "
            f"ROI x=({float(rxmin):.2f},{float(rxmax):.2f}) "
            f"y=({float(rymin):.2f},{float(rymax):.2f})")


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
MONO = [gr.themes.GoogleFont("IBM Plex Mono"), "DejaVu Sans Mono", "monospace"]

THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.blue,
    secondary_hue=gr.themes.colors.gray,
    neutral_hue=gr.themes.colors.gray,
    font=MONO,
    font_mono=MONO,
    radius_size=gr.themes.sizes.radius_sm,
    spacing_size=gr.themes.sizes.spacing_sm,
).set(
    body_background_fill="#f6f7f9",
    body_text_color="#1f2937",
    block_background_fill="#ffffff",
    block_border_color="#e2e8f0",
    block_title_background_fill="#f1f5f9",
    block_title_text_color="#475569",
    input_background_fill="#ffffff",
    input_border_color="#cbd5e1",
    button_primary_background_fill="#1f6feb",
    button_primary_background_fill_hover="#2f7bf5",
    button_primary_text_color="#ffffff",
    button_secondary_background_fill="#eef2f6",
    button_secondary_text_color="#334155",
)

CUSTOM_CSS = """
html, body { color-scheme: light; }
.gradio-container {
    max-width: 1320px !important;
    padding: 8px 16px 16px !important;
}
.prose h1, .prose h2, .prose h3, .prose p, .prose li, .prose code {
    font-family: 'IBM Plex Mono', 'DejaVu Sans Mono', monospace;
}
:root { --body-font: 'IBM Plex Mono', 'DejaVu Sans Mono', monospace; }
footer { display: none !important; }
#status-banner { border-left: 4px solid #1f6feb; padding-left: 12px; }
.blocks-wrap .wrap { row-gap: 6px !important; }
.block { margin-bottom: 6px !important; }
.back { background: #f6f7f9; }
"""

with gr.Blocks(title="AnyTraverse Studio") as demo:
    gr.Markdown("# 🚜 AnyTraverse Studio — Live Evaluation & HITL Dashboard")

    with gr.Row():
        with gr.Column(scale=3):
            # -------- main view --------
            live_view = gr.Image(label="Raw+ROI (TL) | Traversability (TR) | "
                                        "Uncertainty (BL) | ROI crop (BR)",
                                 height=300)
            attn_view = gr.Image(label="Attention maps (all prompts) — live only",
                                 height=160)
            status_banner = gr.Markdown(
                "### Status: ready — upload a video and press ▶️ Go.",
                elem_id="status-banner")

            with gr.Row():
                live_plot = gr.LinePlot(x="Frame", y=["ROI Trav", "ROI Unc"],
                                        title="Live ROI Metrics (ROI Trav & ROI "
                                              "Uncert vs threshold)",
                                        height=180, )
            metric_bars = gr.HTML(value=bars_html(0.0, 0.0, session.uncert_thresh),
                                  label="ROI Score Gauges")

        with gr.Column(scale=2):
            # -------- Controls --------
            video_in = gr.File(label="📹 Upload Off-Road Video (.mp4, .mov, .avi)",
                               file_count="single")
            pref_input = gr.Textbox(
                value="{}",
                placeholder='e.g. {"road": 1.0, "grass": 0.0, "bush": -0.8}',
                label="Traversability Preferences (τ, JSON) — leave {} for defaults")
            with gr.Row():
                sim_thresh = gr.Slider(0.05, 1.0, value=0.8, step=0.05,
                                       label="Ref Scene Sim. Threshold")
                uncert_thresh = gr.Slider(0.05, 1.0, value=0.4, step=0.05,
                                          label="ROI Uncertainty Threshold")
            frame_skip = gr.Slider(1, 10, value=2, step=1,
                                   label="Frame Skip (VLM inference interval)")
            gr.Markdown("#### ROI (normalized) — editable live")
            with gr.Row():
                rx_min = gr.Number(value=0.333, label="ROI X Min", step=0.01)
                rx_max = gr.Number(value=0.667, label="ROI X Max", step=0.01)
            with gr.Row():
                ry_min = gr.Number(value=0.600, label="ROI Y Min", step=0.01)
                ry_max = gr.Number(value=0.950, label="ROI Y Max", step=0.01)
            cfg_status = gr.Markdown("_Live-threshold / ROI edits apply to the "
                                     "running pipeline instantly._")
            with gr.Row():
                run_btn = gr.Button("▶️ Go / Reset", variant="primary")
                sim_btn = gr.Button("⏸ Simulate HOC", variant="secondary")

            with gr.Group(visible=False) as operator_box:
                gr.Markdown("### 🚨 HUMAN OPERATOR CALL")
                gr.Markdown(
                    "Enter τ updates as `prompt`: `weight; prompt: weight`, e.g. "
                    "`mud: -0.7; gravel: 0.6`. Type **ok** (or leave blank) to "
                    "resume without changing preferences (registers the scene).")
                operator_text = gr.Textbox(label="Operator τ update / ok",
                                           placeholder="mud: -0.7; gravel: 0.6")
                resume_btn = gr.Button("✅ Apply & Resume", variant="primary")

            gr.Markdown("#### Per-frame outputs")
            with gr.Row():
                m_frame = gr.Textbox(label="Frame", value="0", interactive=False)
                m_skip = gr.Textbox(label="Skip", value="2", interactive=False)
                m_state = gr.Textbox(label="State", value="ok", interactive=False)
            with gr.Row():
                m_trav = gr.Textbox(label="ROI Trav", value="0.000", interactive=False)
                m_unc = gr.Textbox(label="ROI Unc", value="0.000", interactive=False)
                m_sim = gr.Textbox(label="Ref Sim", value="0.000", interactive=False)
            with gr.Row():
                m_fps = gr.Textbox(label="FPS", value="0", interactive=False)
                m_lat = gr.Textbox(label="Latency", value="0 ms", interactive=False)
                m_dev = gr.Textbox(label="Device", value=session.device, interactive=False)

            with gr.Row():
                log_table = gr.DataFrame(headers=TEL_COLUMNS, interactive=False,
                                         label="Telemetry")
                download_out = gr.DownloadButton(label="⬇ Download composed video (.mp4)",
                                                 value=None, variant="primary")

    # ---------------- Event wiring ----------------
    inputs = [video_in, pref_input, sim_thresh, uncert_thresh, rx_min, rx_max,
              ry_min, ry_max, frame_skip]
    outputs = [live_view, status_banner, operator_box, attn_view, m_frame, m_skip,
               m_state, m_trav, m_unc, m_sim, m_fps, m_lat, live_plot, metric_bars,
               log_table, download_out]

    cfg_inputs = [sim_thresh, uncert_thresh, rx_min, rx_max, ry_min, ry_max]
    for ctl in (sim_thresh, uncert_thresh, rx_min, rx_max, ry_min, ry_max):
        ctl.change(live_pipeline_update, inputs=cfg_inputs, outputs=[cfg_status])

    run_btn.click(run_evaluation, inputs=inputs, outputs=outputs)
    sim_btn.click(simulate_hoc, outputs=[status_banner])
    resume_btn.click(
        handle_operator_resume, inputs=[operator_text],
        outputs=[status_banner, operator_box, pref_input],
    ).then(run_evaluation, inputs=inputs, outputs=outputs)


if __name__ == "__main__":
    on_spaces = bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE"))
    demo.queue().launch(
        share=not on_spaces, theme=THEME, css=CUSTOM_CSS,
        allowed_paths=["."],
        server_name="0.0.0.0",
    )