"""PadCaptioner: omni-modal dense video captioning with parallelized autoregressive decoding.

Paper: "Parallelized Autoregressive Decoding for Omni-Modal Dense Video Captioning" (ECCV 2026).
Code:   https://github.com/showlab/PadCaptioner
Model:  wenzhengzeng/PadCaptioner-3B

The app mirrors the authors' reference test loop (`qwenvl/train/train_qwen_modified.py`,
`scripts_run/test.sh`) exactly: the same preprocessing (LazySupervisedDataset with
base_interval=0.5s, 64..256 frames, per-frame pixel budgets), the same inference options
(grounding threshold 0.7, snap window 1, coverage mask, branch interval attention, dedup
IoU 0.75), and the same post-processing (interval insertion into the <G> placeholders,
NMS over records).
"""

import os
import subprocess
import sys

# Vendored transformers 4.51.3 (the version PadCaptioner was trained/evaluated with).
# PyPI's transformers 4.51.3 metadata pins huggingface-hub<1.0, which conflicts with
# gradio 6.x's huggingface-hub>=1.16 at pip-resolution time, so the vendored wheel's
# metadata relaxes the cap to >=0.30.0 — the code itself is untouched. Installed with
# --no-deps because its real deps are listed in requirements.txt.
_HERE = os.path.dirname(os.path.abspath(__file__))
_VENDORED_WHEEL = os.path.join(_HERE, "vendor", "transformers-4.51.3-py3-none-any.whl")
try:
    import transformers  # noqa: F401
    from transformers import __version__ as _tf_version  # noqa: F401
    if _tf_version != "4.51.3":
        raise ImportError("wrong transformers version")
except Exception:
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "--no-deps", "--quiet", _VENDORED_WHEEL],
        check=True,
    )

import spaces  # MUST come before torch / transformers CUDA init

import json
import re
import tempfile
import time
from types import SimpleNamespace

import gradio as gr

# The bundled qwenvl package (from github.com/showlab/PadCaptioner) sits next to app.py.
sys.path.insert(0, _HERE)

import torch  # noqa: E402  (after spaces)
import transformers  # noqa: E402
from transformers import AutoTokenizer, WhisperFeatureExtractor  # noqa: E402

from qwenvl.data.dataset import LazySupervisedDataset  # noqa: E402
from qwenvl.data.image_processing_qwen2_vl_fast import Qwen2VLImageProcessorFast  # noqa: E402
from qwenvl.model.modeling_qwen2_5_vl import video_SALMONN2_plus  # noqa: E402

MODEL_ID = "wenzhengzeng/PadCaptioner-3B"

# ---- Inference settings: the released-model evaluation settings from scripts_run/test.sh ----
MAX_PIXELS = 61250
MIN_PIXELS = 784
MIN_FRAMES = 64
MAX_FRAMES = 256
BASE_INTERVAL = 0.5
GROUNDING_THRESHOLD = 0.7
GROUNDING_MERGE_GAP = 0
GROUNDING_SNAP_WINDOW = 1
MAX_PLANNING_EVENTS = 24
DVC_DEDUP_IOU = 0.75
MAX_NEW_TOKENS = 1026  # task_body_cap default 1024 + 2 (<G> and EOS)

# The dvc prompt used by the authors' test annotations (DATASET.md / test_dvc2000.json).
DVC_PROMPT = (
    "<video>\nCould you please detail the events that took place during different time "
    "segments in the video? List the events in the format: From second{} to second{}, "
    "event1. \n From second{} to second{}, event2. \n ..."
)

# EVIDENCE_GEN_KEYS: non-tensor grounding inputs the model needs alongside the tensors
# (from the reference test loop).
EVIDENCE_GEN_KEYS = (
    "video_token_indices_list_src",
    "audio_token_indices_list_src",
    "second_per_grid_ts",
    "last_text_token_indices",
    "video_token_indices_per_timestep",
    "audio_token_indices_per_timestep",
)


# ---------- Post-processing helpers (verbatim from qwenvl/train/train_qwen_modified.py) ----------

def format_intervals(intervals):
    return [f"{subinterval[0]} - {subinterval[1]} " for subinterval in intervals]


def timestep_intervals_to_seconds(intervals, timestep_duration, round_digits=2):
    """Convert timestep-index intervals [[a, b], ...] to seconds."""
    if timestep_duration is None or timestep_duration <= 0:
        return intervals
    out = []
    for a, b in intervals:
        start_sec = a * timestep_duration
        end_sec = (b + 1) * timestep_duration  # half-open interval: end time of step b
        out.append([round(start_sec, round_digits), round(end_sec, round_digits)])
    return out


def is_degenerate(body):
    """True if the body is too short, mostly non-ASCII, or leaks the template / modality tags."""
    b = body.strip()
    if len(b.split()) < 6:
        return True
    non_ascii = sum(1 for ch in b if ord(ch) > 127) / max(1, len(b))
    return (non_ascii > 0.3 or "Analyze the" in b or "</audio>" in b or "<audio>" in b
            or "</video>" in b or "<video>" in b)


def trim_to_sentence(body, min_keep=0.5):
    """Drop the fragment after the last sentence end if at least min_keep of the text remains."""
    m = list(re.finditer(r'[.!?]["”]?(?=\s|$)', body))
    if not m:
        return body
    cut = m[-1].end()
    return body[:cut] if cut >= min_keep * len(body) else body


def dedup_dvc_records(text, iou_thresh, drop_degenerate=False):
    """NMS over deinterleaved dvc records ("start - end caption" per line)."""
    head = re.compile(r"^\s*(\d+(?:\.\d+)?)\s*-\s*(\d+(?:\.\d+)?)\s*(.*)$")
    records = []  # [start, end, caption, raw_lines]
    for line in text.split("\n"):
        m = head.match(line)
        if m:
            records.append([float(m.group(1)), float(m.group(2)), m.group(3), line])
        elif records:
            records[-1][2] += " " + line.strip()
            records[-1][3] += "\n" + line
        else:
            return text  # leading non-record text: not a deinterleaved dvc output
    if len(records) < 2:
        return text
    kept = []
    for s, e, cap, raw in records:
        if drop_degenerate and is_degenerate(cap):
            continue
        ncap = " ".join(cap.split()).lower()
        redundant = False
        for ks, ke, kcap, _ in kept:
            inter = max(0.0, min(e, ke) - max(s, ks))
            union = max(e, ke) - min(s, ks)
            iou = inter / union if union > 0 else 1.0
            if iou >= iou_thresh or (ncap and ncap == kcap and iou > 0.1):
                redundant = True
                break
        if not redundant:
            kept.append((s, e, ncap, raw))
    return "\n".join(r[3] for r in kept)


# ---------- Model / processor load (module scope, eager .cuda()) ----------

print("[app] loading tokenizer / processors ...", flush=True)
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    model_max_length=131072,
    padding_side="right",
    use_fast=False,
)
image_processor = Qwen2VLImageProcessorFast.from_pretrained(MODEL_ID)
audio_processor = WhisperFeatureExtractor(
    feature_size=128,
    sampling_rate=16000,
    hop_length=160,
    chunk_length=30,
)

print("[app] loading model ...", flush=True)
model = video_SALMONN2_plus.from_pretrained(
    MODEL_ID,
    attn_implementation="sdpa",
    torch_dtype=torch.bfloat16,
)
# Inference-time config overrides (the reference test loop's settings for the released model).
model.config.grounding_threshold = GROUNDING_THRESHOLD
model.config.grounding_merge_gap = GROUNDING_MERGE_GAP
model.config.grounding_snap_window = GROUNDING_SNAP_WINDOW
model.config.max_planning_events = MAX_PLANNING_EVENTS
model.config.audio_mean_weight = 1.0  # dvc: audio weighted into the aggregated <G> feature
model.eval()
model.cuda()
print("[app] model ready", flush=True)


# ---------- Preprocessing: one-element annotation -> LazySupervisedDataset item ----------

def build_gen_kwargs(video_path, use_audio=True):
    """Run the authors' preprocessing for a single video and assemble generate_parallel kwargs."""
    annotation = {
        "id": "demo",
        "video": os.path.abspath(video_path),
        "task": "dvc",
        "use_audio": bool(use_audio),
        "video_id": "dvc_demo",
        "conversations": [
            {"from": "human", "value": DVC_PROMPT},
            {"from": "gpt", "value": "placeholder"},  # truncated away in run_test mode
        ],
    }
    with tempfile.NamedTemporaryFile("w", suffix=".json", delete=False) as f:
        json.dump([annotation], f)
        ann_path = f.name

    data_args = SimpleNamespace(
        dataset_use=ann_path,
        model_type="qwen2.5vl",
        image_processor=image_processor,
        audio_processor=audio_processor,
        max_pixels=MAX_PIXELS,
        min_pixels=MIN_PIXELS,
        video_max_frame_pixels=MAX_PIXELS,
        video_min_frame_pixels=MIN_PIXELS,
        video_min_frames=MIN_FRAMES,
        video_max_frames=MAX_FRAMES,
        base_interval=BASE_INTERVAL,
        run_test=True,
        aggregation_jitter_ratio=0.0,
        do_sample=False,
        num_sample=1,
        deinterleave_parallel_output=True,
        feature_size=128,
        chunk_length=30,
        hop_length=160,
        sampling_rate=16000,
    )
    model_args = SimpleNamespace(use_evidence=True)

    dataset = LazySupervisedDataset(
        tokenizer=tokenizer, data_args=data_args, model_args=model_args
    )
    inputs = dataset[0]  # identity collate: batch_size=1 -> the raw sample dict
    os.unlink(ann_path)

    # Non-tensor keys the reference loop keeps aside (metadata + grounding inputs).
    res = {
        "video": inputs.pop("video", None),
        "use_audio": inputs.pop("use_audio", False),
        "task": inputs.pop("task", None),
    }
    for k in EVIDENCE_GEN_KEYS:
        res[k] = inputs.get(k, None)
    res["timestep_duration"] = inputs.get("timestep_duration", None)

    gen_kwargs = {
        k: v.to("cuda") for k, v in inputs.items() if isinstance(v, torch.Tensor)
    }
    for k in EVIDENCE_GEN_KEYS:
        gen_kwargs[k] = res[k]
    gen_kwargs["timestep_duration"] = res["timestep_duration"]
    return gen_kwargs


# ---------- Inference ----------

@spaces.GPU(duration=30)  # measured: 23.2-23.7s on the bundled examples + margin
def caption_video(video: str, use_audio: bool = True,
                  progress=gr.Progress(track_tqdm=True)) -> str:
    """Produce temporally grounded event captions for a video (with its audio track).

    Args:
        video: uploaded video file. Its audio track is used when present.
        use_audio: feed the audio track to the model along with the video frames.

    Returns:
        One "start - end caption" record per detected event, in temporal order.
    """
    if video is None:
        raise gr.Error("Please upload a video first.")
    t0 = time.perf_counter()

    gen_kwargs = build_gen_kwargs(video, use_audio=use_audio)

    newline_token_id = tokenizer.encode("\n", add_special_tokens=False)[0]
    with torch.no_grad():
        outputs = model.generate_parallel(
            **gen_kwargs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            top_p=0.9,
            constrain_planning=True,   # dvc planning vocab constraint
            min_body_tokens=0,
            single_event_no_replan=False,
            record_separator_id=newline_token_id,
            deinterleave=True,
        )

    output_trimmed = outputs[0, len(gen_kwargs["input_ids"][0]):]
    output_text = tokenizer.decode(
        output_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )

    # Grounding readout: replace each <G> placeholder with its predicted interval.
    tgt = getattr(model, "tgt", None)
    if tgt is not None and len(tgt[0]) == len(output_trimmed):
        evi_token_id = getattr(model.config, "evi_token_id", None)
        match_inds = torch.where(output_trimmed == evi_token_id)[0].tolist()
        if len(match_inds) > 0:
            timestep_duration = gen_kwargs["timestep_duration"]
            evi_strs = []
            for i in match_inds:
                intervals_sec = timestep_intervals_to_seconds(
                    tgt[0][i], timestep_duration)
                evi_strs.append(", ".join(s.strip() for s in format_intervals(intervals_sec)))

            parts = output_text.split("<G>")
            if len(parts) == len(evi_strs) + 1:
                pieces = [parts[0]]
                for i, (iv, rest) in enumerate(zip(evi_strs, parts[1:])):
                    pieces.append(iv)
                    if rest == "":
                        if i < len(evi_strs) - 1:
                            pieces.append("; ")
                    elif rest[0] not in " ;,.<":
                        pieces.append(" ")
                    pieces.append(rest)
                output_text = "".join(pieces)

        # NMS over the deinterleaved records (reference: dvc_dedup_iou=0.75, drop degenerate).
        output_text = dedup_dvc_records(
            output_text, DVC_DEDUP_IOU, drop_degenerate=True)

        # Length cap hit: trim to the last sentence end.
        if len(output_trimmed) >= MAX_NEW_TOKENS - 1:
            m_iv = re.match(r"^[\d\.\s\-,]*", output_text)
            head, body_ = output_text[:m_iv.end()], output_text[m_iv.end():]
            output_text = head + trim_to_sentence(body_)

    if not output_text.strip():
        output_text = "(the model produced no captions for this video)"

    elapsed = time.perf_counter() - t0
    print(f"[app] inference took {elapsed:.1f}s", flush=True)
    return output_text.strip()


# ---------- UI ----------

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES_DIR = os.path.join(_HERE, "examples")
EXAMPLE_VIDEOS = [
    os.path.join(EXAMPLES_DIR, "example_1.mp4"),
    os.path.join(EXAMPLES_DIR, "example_3.mp4"),
    os.path.join(EXAMPLES_DIR, "example_10.mp4"),
]

with gr.Blocks() as demo:  # theme/css moved to launch() in Gradio 6.0
    gr.Markdown(
        """
        # PadCaptioner: Omni-Modal Dense Video Captioning
        **Parallelized Autoregressive Decoding for Omni-Modal Dense Video Captioning** (ECCV 2026)

        Upload a video (with sound) and the model lists **temporally grounded event captions** —
        one `"start - end caption"` record per event, watched *and* heard. PadCaptioner first
        plans event boundaries with a grounding head (the `<G>` tokens), then decodes all event
        captions **in parallel** in a single pass.

        - Paper: [hf.co/papers/2607.02963](https://huggingface.co/papers/2607.02963)
        - Code: [github.com/showlab/PadCaptioner](https://github.com/showlab/PadCaptioner)
        - Model: [wenzhengzeng/PadCaptioner-3B](https://huggingface.co/wenzhengzeng/PadCaptioner-3B)
        """
    )
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column():
                video_in = gr.Video(label="Video (with audio track)", sources=["upload"])
                use_audio = gr.Checkbox(
                    value=True,
                    label="Use the audio track",
                    info="Off = frames only (video captions without sound context).",
                )
                caption_btn = gr.Button("Caption events", variant="primary")
            with gr.Column():
                caption_out = gr.Textbox(
                    label="Grounded event captions",
                    lines=14,
                    buttons=["copy"],  # `show_copy_button` was removed in Gradio 6.0
                )

        gr.Examples(
            examples=[[v, True] for v in EXAMPLE_VIDEOS],
            inputs=[video_in, use_audio],
            fn=caption_video,
            outputs=caption_out,
            cache_examples=True,
            cache_mode="lazy",
            label="Examples — click one, then run it (first click is cached afterwards)",
        )

        # Primary interaction: the button must be bound or the app exposes no endpoint at all.
        caption_btn.click(
            caption_video,
            inputs=[video_in, use_audio],
            outputs=caption_out,
            api_name="caption_video",
        )

        gr.Markdown(
            """
            <details>
            <summary>About the settings</summary>
            The demo runs the authors' released-model evaluation configuration: 0.5 s event
            grid, 64–256 frames, grounding threshold 0.7 with boundary snapping, coverage
            masking, parallel branch interval attention, and record NMS at IoU 0.75.
            Generation is greedy, exactly as in <code>scripts_run/test.sh</code>.
            </details>

            Example videos are from the
            [video-SALMONN_2 test set](https://huggingface.co/datasets/tsinghua-ee/video-SALMONN_2_testset)
            (Apache-2.0).
            """
        )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)