"""Flora-TempSeg demo: fine-grained video temporal segmentation on Qwen3.5-9B.

Reproduces the released model's native inference path (nvidia/Flora-TempSeg):
64 ordered temporal slots @ 4 Hz per 16-second window (12 s emit core, 2 s guards),
20 Hz local ticks, trained pool prompt with the canonical window envelope,
thinking disabled, greedy decoding, <=1024 new tokens, FP32 Gated-DeltaNet
parameters and FP32 LoRA adapter on a BF16 base.
"""
import spaces  # ZeroGPU: must be imported before torch
import json
import math
import re
import time
from bisect import bisect_left
from dataclasses import dataclass
from fractions import Fraction
from pathlib import Path

import av
import gradio as gr
import torch
from PIL import Image

# ---------------------------------------------------------------------------
# Model identity (pinned revisions from the release manifest models.json)
# ---------------------------------------------------------------------------
TEMPSEG_REPO = "nvidia/Flora-TempSeg"
TEMPSEG_REVISION = "114a1076f70eb53f4a58ac045505db3087e79347"
BASE_REPO = "Qwen/Qwen3.5-9B"
BASE_REVISION = "c202236235762e1c871ad0ccb60c8ee5ba337b9a"

TICK_US = 50_000                     # 20 Hz local ticks
SLOT_SPACING_US = 250_000            # 4 Hz slots
SLOT_COUNT = 64                       # 16 s view
SLOT_FPS = 4
LEFT_GUARD_US = RIGHT_GUARD_US = 2_000_000
CORE_US = 12_000_000
VIEW_US = LEFT_GUARD_US + CORE_US + RIGHT_GUARD_US
MAX_NEW_TOKENS = 1024
CONTEXT_BUDGET = 32_768
MAX_SEGMENTS = 32

MAX_DURATION_S = 120
MAX_UPLOAD_BYTES = 256 * 1024 * 1024
MAX_DIMENSION = 4096
MAX_FRAME_COUNT = 7200

# Trained pool prompts (verbatim from the release TEMPSEG_PROMPTS.json).
POOL_PROMPTS = {
    "human": (
        "<video>\n"
        "The video contains 64 ordered temporal slots sampled at 4 Hz for one 16-second local view.\n"
        "The emit core is 12 seconds with a 2-second guard on each side.\n"
        "Edge padding repeats the nearest real frame; ignore slots outside the valid tick envelope.\n"
        "Times use 20 Hz local ticks in the half-open range [0,320).\n"
        "Canonical window envelope:\n"
        '{"valid_start_tick":0,"valid_end_tick":320,"core_start_tick":40,"core_end_tick":280}\n'
        "Emit only subtask intervals with positive half-open intersection with [core_start_tick,core_end_tick).\n"
        "A quantized boundary alias may instead end exactly at core_start_tick or start exactly at core_end_tick.\n"
        "lo=true requires s==valid_start_tick; ro=true requires e==valid_end_tick.\n"
        "Return exactly one JSON object and no prose or Markdown.\n"
        "Use only the top-level key segments; each row's keys must appear exactly as s,e,lo,ro.\n"
        "Preserve concurrent or duplicate intervals as separate rows and sort rows by (s,e,lo,ro).\n"
        "An empty list is allowed only when no interval has emit-core support.\n"
    ),
    "yam": (
        "<video>\n"
        "The video contains 64 ordered temporal slots sampled at 4 Hz for one 16-second local view.\n"
        "The emit core is 12 seconds with a 2-second guard on each side.\n"
        "Segment meaningful local operations. Acquisition and delivery are distinct when observable; retain distinct directional handovers. Ordinary contacts, adjustments and low-motion pauses within the same ongoing goal do not create new events. Keep clear failure/retry or execution-quality transitions as separate intervals even when their intended operation is unchanged. Represent genuine pure rest as an interval, not a gap.\n"
        "Edge padding repeats the nearest real sampled frame; ignore slots outside the valid tick envelope.\n"
        "Times use 20 Hz local ticks in the half-open range [0,320).\n"
        "Canonical window envelope:\n"
        '{"core_end_tick":280,"core_start_tick":40,"valid_end_tick":320,"valid_start_tick":0}\n'
        "Emit only intervals with positive half-open intersection with [core_start_tick,core_end_tick). Clip intervals at valid video limits; lo/ro indicate continuation beyond the visible limit. lo=true requires s==valid_start_tick; ro=true requires e==valid_end_tick.\n"
        "Return exactly one JSON object and no prose or Markdown, with only top-level key segments. Each row has keys exactly in order s,e,lo,ro. Sort rows by (s,e,lo,ro)."
    ),
}
POOL_SHA256 = {
    "human": "21521d5e3b854488f022f402cbb3c31137e2a1a8a114837b16e353c631bc350a",
    "yam": "a50644d790524e1d4c80a5ad154e940d9bb512b0816eb9d2703cf5721a2f1943",
}

# ---------------------------------------------------------------------------
# Media probing (production limits, PyAV)
# ---------------------------------------------------------------------------


class InputError(ValueError):
    """Raised when an uploaded video violates the model's input contract."""


@dataclass
class SourceVideo:
    """Frame index + presentation timestamps (µs, first frame = 0) for one upload."""

    path: Path
    rate: Fraction
    pts_us: tuple
    duration_us: int
    size: tuple

    @property
    def duration_s(self) -> float:
        return self.duration_us / 1_000_000


def probe(path) -> SourceVideo:
    """Index every source frame's presentation timestamp (first frame = 0 µs)."""
    path = Path(path)
    if path.stat().st_size > MAX_UPLOAD_BYTES:
        raise InputError("Video exceeds 256 MiB.")
    with av.open(str(path), options={"err_detect": "explode"}) as container:
        if len(container.streams.video) != 1:
            raise InputError("Use a file containing exactly one video stream.")
        stream = container.streams.video[0]
        stream.codec_context.thread_count = 1
        rate = stream.average_rate
        if not rate or rate < 4:
            raise InputError("The input needs at least 4 source frames per second.")
        size = (stream.codec_context.width, stream.codec_context.height)
        if min(size) <= 0 or max(size) > MAX_DIMENSION:
            raise InputError("Video dimensions must be at most 4096 pixels per side.")
        times, zero, last_length = [], None, None
        for frame in container.decode(stream):
            if frame.pts is None:
                raise InputError("Every source frame needs a presentation timestamp.")
            zero = frame.pts if zero is None else zero
            t = Fraction(frame.pts - zero) * stream.time_base * 1_000_000
            if times and t <= times[-1]:
                raise InputError("Source timestamps must be strictly increasing.")
            if len(times) >= MAX_FRAME_COUNT:
                raise InputError(f"Video exceeds {MAX_FRAME_COUNT} frames.")
            if t > MAX_DURATION_S * 1_000_000:
                raise InputError(f"Video exceeds {MAX_DURATION_S:g} seconds.")
            times.append(t)
            last_length = Fraction(frame.duration) * stream.time_base * 1_000_000 if frame.duration else None
        if len(times) < 2:
            raise InputError("The video contains too few source frames.")
        tail = last_length or Fraction(1_000_000, rate)
        duration_us = int(round(times[-1] + tail))
        if duration_us > MAX_DURATION_S * 1_000_000:
            raise InputError(f"Video exceeds {MAX_DURATION_S:g} seconds.")
    return SourceVideo(path, rate, tuple(times), duration_us, size)


def decode_selected(video: SourceVideo, wanted: set) -> dict:
    """Decode only the requested source frame indices as RGB PIL images."""
    wanted = set(wanted)
    result = {}
    if not wanted:
        return result
    if min(wanted) < 0 or max(wanted) >= len(video.pts_us):
        raise InputError("A requested source frame is outside the video.")
    with av.open(str(video.path), options={"err_detect": "explode"}) as container:
        stream = container.streams.video[0]
        stream.codec_context.thread_count = 1
        zero = None
        for i, frame in enumerate(container.decode(stream)):
            zero = frame.pts if zero is None else zero
            if i in wanted:
                if Fraction(frame.pts - zero) * stream.time_base * 1_000_000 != video.pts_us[i]:
                    raise InputError("Source frame timestamps changed while decoding.")
                result[i] = frame.to_image().convert("RGB")
            if i >= max(wanted):
                break
    if set(result) != wanted:
        raise InputError("A requested source frame could not be decoded.")
    return result


# ---------------------------------------------------------------------------
# Window scheduling (production scope scheduler, no coarse priors)
# ---------------------------------------------------------------------------


@dataclass
class Window:
    """One fine window over the source video."""

    window_index: int
    requested_start_us: int
    requested_end_us: int
    valid_start_tick: int
    valid_end_tick: int
    core_start_tick: int
    core_end_tick: int
    global_tick_offset: int  # requested_start_us // TICK_US


def _floor_div(a, b):
    return a // b


def _ceil_div(a, b):
    return -(-a // b)


def schedule_windows(duration_us: int) -> list:
    """Emit cores advance by 12 s; guards are clipped at the video edges."""
    windows = []
    index = 0
    while True:
        core_start_us = index * CORE_US
        if core_start_us >= duration_us:
            break
        core_end_us = min(core_start_us + CORE_US, duration_us)
        requested_start_us = core_start_us - LEFT_GUARD_US
        requested_end_us = requested_start_us + VIEW_US
        valid_start_us = max(0, requested_start_us)
        valid_end_us = min(duration_us, requested_end_us)
        windows.append(
            Window(
                window_index=index,
                requested_start_us=requested_start_us,
                requested_end_us=requested_end_us,
                valid_start_tick=_floor_div(valid_start_us - requested_start_us, TICK_US),
                valid_end_tick=_ceil_div(valid_end_us - requested_start_us, TICK_US),
                core_start_tick=_floor_div(core_start_us - requested_start_us, TICK_US),
                core_end_tick=_ceil_div(core_end_us - requested_start_us, TICK_US),
                global_tick_offset=requested_start_us // TICK_US,
            )
        )
        index += 1
    if not windows:
        raise InputError("Video is too short to segment.")
    return windows


# ---------------------------------------------------------------------------
# Frame plan: 64 slots @ 4 Hz, nearest source frame per slot, edge padding
# ---------------------------------------------------------------------------


def _nearest_index(requested_us: int, times, max_selection_error_us: int) -> int:
    """Index of the source frame nearest requested_us (within half a slot)."""
    requested = Fraction(requested_us, 1)
    position = bisect_left(times, requested)
    candidates = [i for i in (position - 1, position) if 0 <= i < len(times)]
    if not candidates:
        raise InputError("No decoded frame is available for an in-video slot.")
    selected = min(candidates, key=lambda i: (abs(times[i] - requested), times[i], i))
    if abs(times[selected] - requested) > max_selection_error_us:
        raise InputError("Nearest decoded frame is farther than half a profile slot.")
    return selected


def frame_plan(window: Window, video: SourceVideo) -> list:
    """Slot -> source frame index map (64 slots, nearest frame, edge padding)."""
    pts = video.pts_us
    if not pts or pts[0] != 0 or any(a >= b for a, b in zip(pts, pts[1:])):
        raise InputError("Source timestamps must start at zero and strictly increase.")
    if pts[-1] >= video.duration_us:
        raise InputError("The last frame lies outside the video duration.")
    requested = [window.requested_start_us + i * SLOT_SPACING_US for i in range(SLOT_COUNT)]
    pad = [not 0 <= t < video.duration_us for t in requested]
    indices = [None if padding else _nearest_index(t, pts, SLOT_SPACING_US // 2) for t, padding in zip(requested, pad)]
    real = [i for i, v in enumerate(indices) if v is not None]
    if not real:
        raise InputError("No decoded frame is available for this window.")
    first, last = real[0], real[-1]
    indices[:first] = [indices[first]] * first
    indices[last + 1:] = [indices[last]] * (SLOT_COUNT - last - 1)
    observed = [v for v, p in zip(indices, pad) if not p]
    if len(set(observed)) != len(observed):
        raise InputError("Real temporal slots must not duplicate evidence.")
    return indices


# ---------------------------------------------------------------------------
# Window prompt, instruction suffix, validation, projection
# ---------------------------------------------------------------------------

_ENVELOPE_RE = re.compile(r"(Canonical window envelope:\s*)(\{[^{}]+\})")


def window_prompt(window: Window, pool: str) -> str:
    """Trained pool prompt with this window's envelope JSON substituted."""
    template = POOL_PROMPTS[pool]
    matches = list(_ENVELOPE_RE.finditer(template))
    if len(matches) != 1:
        raise ValueError("pool template must contain exactly one envelope line")
    match = matches[0]
    keys = list(json.loads(match[2]))
    if set(keys) != {"valid_start_tick", "valid_end_tick", "core_start_tick", "core_end_tick"}:
        raise ValueError("envelope keys are invalid")
    envelope = json.dumps(
        {k: getattr(window, k) for k in keys},
        separators=(",", ":"),
    )
    return template[: match.start(2)] + envelope + template[match.end(2):]


def instruction_suffix(instruction: str) -> str:
    """Optional episode instruction, appended as sorted JSON."""
    value = (instruction or "").strip()
    if not value:
        return ""
    if len(value) > 2000:
        raise InputError("Keep the optional task instruction below 2000 characters.")
    return "\n" + json.dumps(
        {"episode_instruction": value}, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False
    )


def _object_without_duplicate_keys(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError(f"duplicate JSON key is forbidden: {key}")
        result[key] = value
    return result


def parse_prediction(text: str, window: Window) -> list:
    """Strict production validation of the model's rows against the envelope."""
    if not isinstance(text, str) or not text:
        raise ValueError("The model did not return a complete output for this input.")
    try:
        value = json.loads(text, object_pairs_hook=_object_without_duplicate_keys)
    except json.JSONDecodeError as exc:
        raise ValueError(f"Model output is not valid JSON: {exc}") from exc
    if not isinstance(value, dict) or list(value) != ["segments"]:
        raise ValueError("Output must be exactly one object with the single key segments.")
    rows = value["segments"]
    if not isinstance(rows, list):
        raise ValueError("segments must be a list.")
    if len(rows) > MAX_SEGMENTS:
        raise ValueError(f"prediction contains more than {MAX_SEGMENTS} segments")
    parsed, previous = [], None
    for index, row in enumerate(rows):
        if not isinstance(row, dict) or list(row) != ["s", "e", "lo", "ro"]:
            raise ValueError(f"segment {index}: keys must appear exactly as s,e,lo,ro")
        s, e, lo, ro = row["s"], row["e"], row["lo"], row["ro"]
        if isinstance(s, bool) or isinstance(e, bool) or not isinstance(s, int) or not isinstance(e, int):
            raise ValueError(f"segment {index}: ticks must be integers")
        if type(lo) is not bool or type(ro) is not bool:
            raise ValueError(f"segment {index}: open flags must be JSON booleans")
        if not window.valid_start_tick <= s < e <= window.valid_end_tick:
            raise ValueError(f"segment {index}: lies outside the valid tick envelope")
        if lo and s != window.valid_start_tick:
            raise ValueError(f"segment {index}: lo=true requires s==valid_start_tick")
        if ro and e != window.valid_end_tick:
            raise ValueError(f"segment {index}: ro=true requires e==valid_end_tick")
        key = (s, e, lo, ro)
        if previous is not None and key < previous:
            raise ValueError("segments are not in canonical order")
        previous = key
        core_intersects = s < window.core_end_tick and e > window.core_start_tick
        core_touches = e == window.core_start_tick or s == window.core_end_tick
        if not (core_intersects or core_touches):
            continue  # production drops guard-only fragments
        parsed.append(key)
    return parsed


def project_rows(rows, window: Window, duration_us: int) -> list:
    """Local ticks -> source-video microseconds (production projection)."""
    events = []
    for s, e, lo, ro in rows:
        start_us = max(0, (window.global_tick_offset + s) * TICK_US)
        end_us = min(duration_us, (window.global_tick_offset + e) * TICK_US)
        if end_us <= start_us:
            continue
        completeness = {
            (False, False): "complete",
            (True, False): "left_censored",
            (False, True): "right_censored",
            (True, True): "both_censored",
        }[(lo, ro)]
        events.append(
            {
                "start_s": round(start_us / 1_000_000, 3),
                "end_s": round(end_us / 1_000_000, 3),
                "completeness": completeness,
                "window": window.window_index,
            }
        )
    return events


def stitch_events(events: list) -> list:
    """Merge window events into video-level intervals (simplified vs production).

    Censored events that touch across a 12 s core seam are joined; the
    production stitcher additionally resolves overlapping components within a
    2-tick tolerance, which needs captioner support that is out of scope here.
    """
    ordered = sorted(events, key=lambda e: (e["start_s"], e["end_s"]))
    merged = []
    for event in ordered:
        if merged:
            prev = merged[-1]
            touches = event["start_s"] <= prev["end_s"] + 1e-9
            if touches and (
                prev["completeness"] in ("right_censored", "both_censored")
                or event["completeness"] in ("left_censored", "both_censored")
            ):
                prev["end_s"] = max(prev["end_s"], event["end_s"])
                prev["completeness"] = "stitched"
                prev["windows"].append(event["window"])
                continue
        merged.append({**event, "windows": [event["window"]]})
    return merged


# ---------------------------------------------------------------------------
# Preprocessing helpers
# ---------------------------------------------------------------------------


def regularized_resize_dims(height, width, min_pixels=256, max_pixels=40000):
    """Aspect-preserving pixel-area regularization (production helper)."""
    if min(height, width, min_pixels, max_pixels) <= 0 or min_pixels > max_pixels:
        raise ValueError("regularized resize arguments are invalid")
    area = height * width
    if area > max_pixels:
        scale = math.sqrt(max_pixels / area)
        return int(height * scale), int(width * scale)
    if area < min_pixels:
        scale = math.sqrt(min_pixels / area)
        return int(height * scale), int(width * scale)
    return height, width


def human_transport(image: Image.Image) -> Image.Image:
    """JPEG q98 4:4:4 round-trip the trained human pool applies to every frame."""
    import io

    stream = io.BytesIO()
    image.save(stream, format="JPEG", quality=98, subsampling=0, optimize=False)
    stream.seek(0)
    with Image.open(stream) as decoded:
        return decoded.convert("RGB")


# ---------------------------------------------------------------------------
# Model loading (module scope, ZeroGPU eager)
# ---------------------------------------------------------------------------
print("Loading Flora-TempSeg: Qwen3.5-9B base + FP32 GDN + FP32 LoRA adapter ...", flush=True)
_LOAD_STARTED = time.time()

from huggingface_hub import hf_hub_download
from peft import PeftConfig, get_peft_model, set_peft_model_state_dict
from safetensors.torch import load_file
from transformers import AutoModelForImageTextToText, AutoProcessor
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5PreTrainedModel
from transformers.models.qwen3_vl.video_processing_qwen3_vl import Qwen3VLVideoProcessor
from transformers.video_utils import VideoMetadata

# Keep Gated DeltaNet parameters FP32 (native precision contract).
Qwen3_5PreTrainedModel._keep_in_fp32_modules_strict = sorted(
    set(getattr(Qwen3_5PreTrainedModel, "_keep_in_fp32_modules_strict", None) or ()) | {"A_log", "dt_bias"}
)

_BASE = AutoModelForImageTextToText.from_pretrained(
    BASE_REPO,
    revision=BASE_REVISION,
    dtype=torch.bfloat16,
    attn_implementation="sdpa",
)
_ADAPTER_FILE = hf_hub_download(TEMPSEG_REPO, "adapter_model.safetensors", revision=TEMPSEG_REVISION)
MODEL = get_peft_model(
    _BASE,
    PeftConfig.from_pretrained(TEMPSEG_REPO, revision=TEMPSEG_REVISION),
    adapter_name="candidate",
    autocast_adapter_dtype=True,
)
set_peft_model_state_dict(MODEL, load_file(str(_ADAPTER_FILE), device="cpu"), adapter_name="candidate")
MODEL.to("cuda").eval()

# BF16 temporal compute autocast around generate (native contract).
_ORIGINAL_GENERATE = MODEL.generate


def _generate_bf16(*args, **kwargs):
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        return _ORIGINAL_GENERATE(*args, **kwargs)


MODEL.generate = _generate_bf16

PROCESSOR = AutoProcessor.from_pretrained(
    TEMPSEG_REPO,
    revision=TEMPSEG_REVISION,
    use_fast=True,
    min_pixels=1024,
    max_pixels=589824,
    video_processor=Qwen3VLVideoProcessor.from_pretrained(
        TEMPSEG_REPO, revision=TEMPSEG_REVISION, fps=SLOT_FPS, max_frames=128
    ),
)
PROCESSOR.chat_template = Path(
    hf_hub_download(TEMPSEG_REPO, "chat_template.jinja", revision=TEMPSEG_REVISION)
).read_text()

print(f"Model ready in {time.time() - _LOAD_STARTED:.1f}s", flush=True)


def build_inputs(frames, window: Window, pool: str, instruction: str):
    """Processor inputs for one window (do_sample_frames=False, 64-slot metadata)."""
    regularized = []
    for image in frames:
        image = image.convert("RGB")
        height, width = regularized_resize_dims(height=image.height, width=image.width)
        if image.size != (width, height):
            image = image.resize((width, height), Image.Resampling.BICUBIC)
        regularized.append(image)
    prompt = window_prompt(window, pool) + instruction_suffix(instruction)
    text = PROCESSOR.apply_chat_template(
        [{"role": "user", "content": prompt.replace("<video>", PROCESSOR.video_token)}],
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    inputs = PROCESSOR(
        text=[text],
        videos=[regularized],
        video_metadata=[VideoMetadata(total_num_frames=SLOT_COUNT, fps=SLOT_FPS, frames_indices=list(range(SLOT_COUNT)))],
        do_sample_frames=False,
        return_tensors="pt",
    )
    inputs.pop("video_metadata", None)
    return inputs, text


def run_window(window: Window, video: SourceVideo, pool: str, instruction: str):
    """Inference for one window: plan -> decode -> prompt -> greedy generate."""
    indices = frame_plan(window, video)
    images = decode_selected(video, set(indices))
    if pool == "human":
        images = {k: human_transport(v) for k, v in images.items()}
    frames = [images[i] for i in indices]
    inputs, prompt = build_inputs(frames, window, pool, instruction)
    n_tokens = inputs["input_ids"].shape[1]
    if n_tokens + MAX_NEW_TOKENS > CONTEXT_BUDGET:
        raise InputError("The selected window exceeds the native model context budget.")
    args = {k: v.to("cuda") if torch.is_tensor(v) else v for k, v in inputs.items()}
    eos = PROCESSOR.tokenizer.eos_token_id
    started = time.perf_counter()
    with torch.inference_mode():
        output = MODEL.generate(
            **args,
            do_sample=False,
            num_beams=1,
            max_new_tokens=MAX_NEW_TOKENS,
            eos_token_id=eos,
            pad_token_id=PROCESSOR.tokenizer.pad_token_id,
            use_cache=True,
            max_time=300,
        )
    elapsed = time.perf_counter() - started
    ids = output[0, n_tokens:].cpu().tolist()
    stop = bool(ids) and ids[-1] in ([eos] if isinstance(eos, int) else eos)
    text = PROCESSOR.tokenizer.decode(ids, skip_special_tokens=True).strip()
    if not stop or not text:
        raise InputError("The model did not return a complete output for this input.")
    return text, elapsed


# ---------------------------------------------------------------------------
# Timeline visualization
# ---------------------------------------------------------------------------


def timeline_html(events: list, duration_s: float) -> str:
    """SVG timeline with one colored bar per detected subtask interval."""
    if duration_s <= 0:
        return "<p>No timeline.</p>"
    width, track = 940, 46
    lanes = []
    for event in events:
        placed = None
        for lane in lanes:
            if event["start_s"] >= lane["end"]:
                placed = lane
                break
        if placed is None:
            placed = {"end": 0, "bars": []}
            lanes.append(placed)
        placed["end"] = max(placed["end"], event["end_s"])
        placed["bars"].append(event)
    n_lanes = max(1, min(len(lanes), 3))
    height = 12 + n_lanes * track + 22
    scale = width / duration_s
    palette = ["#5b6bd9", "#8f6bcb", "#4f9d8a", "#c96f9e", "#c98a5b", "#6bb1d9"]
    parts = [
        f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" '
        f'style="width:100%;background:#f8f8fb;border-radius:10px">'
    ]
    for li, lane in enumerate(lanes[:3]):
        y = 12 + li * track
        parts.append(
            f'<line x1="0" y1="{y + track // 2}" x2="{width}" y2="{y + track // 2}" stroke="#e2e2ee" stroke-width="1"/>'
        )
        for bi, bar in enumerate(lane["bars"]):
            x = bar["start_s"] * scale
            w = max(3.0, (bar["end_s"] - bar["start_s"]) * scale)
            color = palette[(li + bi) % len(palette)]
            title = f'{bar["start_s"]:.2f}s – {bar["end_s"]:.2f}s ({bar["completeness"]})'
            parts.append(
                f'<g><rect x="{x:.1f}" y="{y + 8}" width="{w:.1f}" height="{track - 16}" rx="6" fill="{color}" opacity="0.85">'
                f'<title>{title}</title></rect>'
                f'<text x="{x + w / 2:.1f}" y="{y + track / 2 + 4}" font-size="11" fill="#fff" text-anchor="middle">{bar["start_s"]:.1f}</text></g>'
            )
    step = max(1, int(duration_s // 8 or 1))
    for t in range(0, int(math.floor(duration_s)) + 1, step):
        x = t * scale
        parts.append(
            f'<line x1="{x:.1f}" y1="0" x2="{x:.1f}" y2="{12 + n_lanes * track}" stroke="#d5d5e4" stroke-width="1"/>'
        )
        parts.append(
            f'<text x="{x:.1f}" y="{height - 8}" font-size="12" fill="#666" text-anchor="middle">{t}s</text>'
        )
    parts.append("</svg>")
    return "".join(parts)


# ---------------------------------------------------------------------------
# Gradio interface
# ---------------------------------------------------------------------------

CSS = """
.gradio-container .prose {font-size: 15px;}
#timeline svg {max-width: 100%; height: auto;}
"""


def _estimate_duration(video_path, pool="human", instruction="", *args, **kwargs):
    """Reserve GPU time for one request, sized from measured per-window cost.

    Measured on ZeroGPU: 1.9-5.0 s of generation per 16-second window on
    10 fps clips; 1.4x margin plus an allowance for worst-case dense
    predictions (32 rows) and heavy source decode.
    """
    try:
        with av.open(str(video_path)) as container:
            stream = container.streams.video[0]
            duration_s = float(stream.duration * stream.time_base) if stream.duration else 0.0
        n_windows = max(1, math.ceil(duration_s / (CORE_US / 1_000_000)))
    except Exception:
        n_windows = 10  # 120 s worst case
    return min(240, 20 + 15 * n_windows)


@spaces.GPU(duration=_estimate_duration)
def segment_video(video_path, pool="human", instruction="", progress=gr.Progress(track_tqdm=True)):
    """Segment an uploaded video into subtask intervals and draw them on its timeline.

    Args:
        video_path: uploaded video file (≤120 s, ≤256 MiB, ≥4 fps, one stream).
        pool: trained prompt pool — "human" for egocentric human-action video,
            "yam" for single-view robot manipulation.
        instruction: optional free-text episode instruction (≤2000 characters).
        progress: Gradio progress bar.

    Returns:
        A dict of merged intervals, a markdown summary, a timeline SVG, and a
        per-window audit dict (envelope, raw model JSON, projected rows).
    """
    if not video_path:
        raise gr.Error("Upload a video first.")
    pool = "human" if pool in (None, "", "human", "Human") else "yam"
    video = probe(video_path)
    windows = schedule_windows(video.duration_us)
    progress(0.05, desc=f"{len(windows)} window(s) planned")
    raw_windows, events, timings = [], [], []
    for i, window in enumerate(windows):
        text, elapsed = run_window(window, video, pool, instruction)
        rows = parse_prediction(text, window)
        window_events = project_rows(rows, window, video.duration_us)
        events.extend(window_events)
        timings.append(elapsed)
        raw_windows.append(
            {
                "window": window.window_index,
                "requested_view_start_s": window.requested_start_us / 1_000_000,
                "envelope": {
                    "valid_start_tick": window.valid_start_tick,
                    "valid_end_tick": window.valid_end_tick,
                    "core_start_tick": window.core_start_tick,
                    "core_end_tick": window.core_end_tick,
                },
                "model_json": json.loads(text),
                "projected": window_events,
                "generation_s": round(elapsed, 2),
            }
        )
        progress((i + 1) / (len(windows) + 1), desc=f"Window {i + 1}/{len(windows)} done")
    merged = stitch_events(events)
    result = {
        "video_duration_s": round(video.duration_s, 3),
        "pool": pool,
        "num_windows": len(windows),
        "raw_windows": raw_windows,
        "segments": [
            {
                "start_s": s["start_s"],
                "end_s": s["end_s"],
                "completeness": s["completeness"],
                "windows": s["windows"],
            }
            for s in merged
        ],
        "total_generation_s": round(sum(timings), 2),
    }
    timeline = timeline_html(merged, video.duration_s)
    summary = (
        f"**{len(merged)}** interval(s) across {len(windows)} window(s) · "
        f"video {video.duration_s:.1f}s · pool `{pool}` · generation {sum(timings):.1f}s"
    )
    return result, summary, timeline, raw_windows


with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="Flora-TempSeg · video temporal segmentation") as demo:
    gr.HTML(
        """
        <div style="display:flex;align-items:baseline;gap:12px;flex-wrap:wrap">
          <h1 style="margin:0">Flora-TempSeg <span style="font-size:.5em;color:#888">· fine-grained video temporal segmentation</span></h1>
        </div>
        <p style="margin-top:6px;max-width:900px">
          A Qwen3.5-9B LoRA (hybrid linear attention) that emits <b>subtask interval boundaries</b> for long egocentric /
          robot-manipulation video. Each 16-second window is sampled as 64 ordered slots @ 4 Hz with a 12-second emit
          core; boundaries come back as 20 Hz local ticks, projected to video seconds and stitched across windows.
          The model emits <i>boundaries only</i> — it does not name actions (that is the separate Flora captioner).
        </p>
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            video_in = gr.Video(label="Video (≤120 s · ≤256 MiB · ≥4 fps · one stream)", sources=["upload"])
            pool_in = gr.Radio(
                ["human", "yam"],
                value="human",
                label="Pool prompt",
                info="human: egocentric human-action videos · yam: single-view robot manipulation",
            )
            instruction_in = gr.Textbox(
                label="Optional episode instruction (≤2000 chars)",
                placeholder="e.g. put the desktop items into the drawer and close the drawer",
            )
            run = gr.Button("Segment video", variant="primary")
        with gr.Column(scale=2):
            summary_md = gr.Markdown("Upload a video and press **Segment video**.")
            timeline = gr.HTML(elem_id="timeline", value="<div></div>")
            segments_json = gr.JSON(label="Segmented intervals (projected to video seconds)")
            with gr.Accordion("Per-window raw output (local ticks, envelope, model JSON)", open=False):
                windows_json = gr.JSON(label="Per-window audit")
    gr.Examples(
        examples=[
            [
                "assets/droid_drawer.mp4",
                "yam",
                "put the desktop items into the drawer and close the drawer",
            ],
            [
                "assets/galaxea_basket.mp4",
                "yam",
                "place the food in the right basket with the right gripper, and place the non food items in the left basket with the left gripper.",
            ],
        ],
        inputs=[video_in, pool_in, instruction_in],
        outputs=[segments_json, summary_md, timeline, windows_json],
        fn=segment_video,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples (WGO-Bench episodes, Apache-2.0 re-upload)",
    )
    run.click(
        segment_video,
        inputs=[video_in, pool_in, instruction_in],
        outputs=[segments_json, summary_md, timeline, windows_json],
        api_name="segment",
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True)