#!/usr/bin/env python3
"""Proact-VL-style streaming inference with off-the-shelf VLMs (Colab friendly).

The official Proact-VL checkpoint is currently restricted, so this module
approximates its proactive streaming behavior (a per-second SPEAK/SILENCE
decision plus a short commentary line, like `MultiAssistantStreamInference.infer_one_chunk`)
using a Qwen2.5-VL / Qwen3-VL backbone.

Two tasks:
  --task commentary (default)
      Per-second proactive commentary. Silent seconds are omitted from `pred`
      (matching the official quickstart), which `f1_only` / `pauc` treat as
      silence. Output follows the official evaluation JSONL format so the
      results can be scored with `proactvl_eval` / the official scripts.
  --task count
      Per-second event detection: the VLM decides each second whether a NEW
      item (e.g. a piece of clothing) is being dropped into the machine,
      with client-side merge dedup. Output record adds `events`, `count`,
      and optionally `gt_count` for immediate accuracy check.
      Use this for 衣物投放视频 ("how many clothes did the user drop").

Output record (one line per video segment):
    {"video_id": "myclip.mp4", "dataset_name": "custom", "begin": 0, "end": 60,
     "idx": 0, "tag": "custom", "model_id": "...", "task": "count",
     "scene": "衣物投放", "pred": {"5": "DROP: red shirt", "12": "DROP: blue jeans"},
     "events": [{"sec": 5, "desc": "red shirt"}, {"sec": 12, "desc": "blue jeans"}],
     "count": 2, "gt_count": 3}

Examples:
  # single video, count dropped clothing items
  python proactvl_stream_infer.py --video clip.mp4 --task count \
      --begin 0 --end 60 --gt_count 3 --out count.jsonl
  # Proact-VL-style commentary (official eval pipeline input)
  python proactvl_stream_infer.py --video clip.mp4 --task commentary \
      --begin 0 --end 30 --tag "Solo commentators" --out pred.jsonl
  # LiveGamingBenchmark segments
  python proactvl_stream_infer.py --dataset oaaoaa/LiveGamingBenchmark \
      --test_name all_in_one --video_dir ./data/LiveGamingBenchmark/DATA \
      --num_samples 5 --out pred.jsonl

Hardware note: the 3B model fits a 16GB GPU in fp16/bf16; 7B/8B models need
`--load_in_4bit` on Colab T4/P100.
"""

import argparse
import json
import os
from dataclasses import dataclass
from pathlib import Path

COMMENTARY_SYSTEM_PROMPT = (
    "You are a live commentator for a game video. "
    "Your role is to independently analyze and narrate the game, delivering "
    "insightful, engaging, and natural commentary just like a human expert. "
    "It is not necessary to speak continuously; during uneventful or "
    "transitional parts, remain silent."
)

COUNT_SYSTEM_PROMPT = (
    "You are a careful video analyst counting physical objects being placed "
    "into a machine or container in real time. You watch a live stream "
    "second by second and detect the exact moment each NEW object is dropped. "
    "An object already in progress from the previous second is not a new one."
)

DEFAULT_COUNT_SCENE = "衣物投放 (clothing items being dropped into a washing machine / drop bin)"


@dataclass
class InferenceConfig:
    model: str = "Qwen/Qwen2.5-VL-3B-Instruct"
    device: str = "cuda"
    dtype: str = "auto"          # auto | float16 | bfloat16
    load_in_4bit: bool = False
    task: str = "commentary"     # commentary | count
    scene: str = ""              # used for count task, e.g. 衣物投放
    merge_gap: int = 1          # event seconds within this gap are merged (dedup)
    gt_count: int = -1          # -1 = unknown; else used for accuracy check
    frames_per_sec: int = 2
    max_new_tokens: int = 12    # mirrors the official generate_config
    temperature: float = 0.7
    top_p: float = 0.9
    repetition_penalty: float = 1.15
    do_sample: bool = True
    history_lines: int = 8      # spoken lines kept in context (commentary)
    model_id: str = "stream_infer"


def sample_frames(video_path: str, begin: int, end: int, fps: int,
                  out_hw=(288, 512)) -> dict:
    """Return {second: [PIL.Image, ...]} for t in [begin, end)."""
    from PIL import Image

    try:
        import decord
        vr = decord.VideoReader(video_path, num_threads=1)
        fps_video = vr.get_avg_fps()
        frames = {}
        for t in range(begin, end):
            i0 = int(t * fps_video)
            idxs = [min(i, len(vr) - 1) for i in range(i0, i0 + fps)]
            frames[t] = [Image.fromarray(vr[i].asnumpy()).convert("RGB").resize(out_hw)
                         for i in idxs]
        return frames
    except ImportError:
        import cv2
        cap = cv2.VideoCapture(video_path)
        fps_video = cap.get(cv2.CAP_PROP_FPS) or 30.0
        frames = {}
        for t in range(begin, end):
            imgs = []
            for k in range(fps):
                cap.set(cv2.CAP_PROP_POS_FRAMES, int((t + (k + 0.5) / fps) * fps_video))
                ok, fr = cap.read()
                if ok:
                    fr = cv2.resize(fr, out_hw)
                    imgs.append(Image.fromarray(fr[:, :, ::-1]).convert("RGB"))
            frames[t] = imgs
        cap.release()
        return frames


class OpenVLStreamInfer:
    """Per-second proactive decisions using an open VLM as the decision head."""

    def __init__(self, cfg: InferenceConfig):
        self.cfg = cfg
        import torch
        from transformers import AutoProcessor, AutoModelForVision2Seq

        self.torch = torch
        kwargs = {}
        if not cfg.load_in_4bit:
            kwargs["torch_dtype"] = cfg.dtype
            kwargs["attn_implementation"] = "sdpa"
        else:
            from transformers import BitsAndBytesConfig
            kwargs["quantization_config"] = BitsAndBytesConfig(
                load_in_4bit=True, bnb_4bit_compute_dtype="float16")
        self.model = AutoModelForVision2Seq.from_pretrained(cfg.model, **kwargs)
        self.model = self.model.to(cfg.device).eval()
        self.processor = AutoProcessor.from_pretrained(cfg.model)

    # ---------- commentary task ----------

    def _history_text(self, spoken: list) -> str:
        if not spoken:
            return "(no prior commentary yet)"
        return "\n".join(f"[t={t}] {txt}" for t, txt in spoken[-self.cfg.history_lines:])

    def chunk(self, second: int, imgs, spoken: list) -> str:
        """Return commentary text, or '' for silence."""
        cfg = self.cfg
        prompt = (
            f"Video stream up to second {second}:\n"
            f"Prior commentary:\n{self._history_text(spoken)}\n"
            f"Current frames for second {second} are attached below.\n"
            "Decide now: if something notable is happening in these frames "
            "and you have not just said it, reply with one short commentary "
            "line (max ~12 words). Otherwise reply exactly SILENCE."
        )
        messages = [
            {"role": "system", "content": COMMENTARY_SYSTEM_PROMPT},
            {"role": "user", "content": [
                *[{"type": "image", "image": im} for im in imgs],
                {"type": "text", "text": prompt},
            ]},
        ]
        inputs = self._generate(messages)
        text = self._decode(inputs)
        if text.upper().startswith("SILENCE") or not text:
            return ""
        return text

    # ---------- count task ----------

    def chunk_event(self, second: int, imgs, events: list) -> str:
        """Return event description ('DROP: ...' style), or '' for no new event."""
        cfg = self.cfg
        scene = cfg.scene or DEFAULT_COUNT_SCENE
        secs = ", ".join(str(e["sec"]) for e in events) if events else "none"
        prompt = (
            f"You are watching a live video of {scene}.\n"
            f"New items confirmed so far: {len(events)} (at seconds: {secs}).\n"
            f"Attached frames show second {second} of the stream.\n"
            "If a NEW item is clearly being dropped/placed in during second "
            f"{second} (it was not already in progress in the previous second), "
            "respond in exactly this format:\n"
            "DROP: <short description of the item, max 6 words>\n"
            "If no new item is dropped in this second, respond exactly:\n"
            "NONE"
        )
        messages = [
            {"role": "system", "content": COUNT_SYSTEM_PROMPT},
            {"role": "user", "content": [
                *[{"type": "image", "image": im} for im in imgs],
                {"type": "text", "text": prompt},
            ]},
        ]
        inputs = self._generate(messages)
        text = self._decode(inputs).strip()
        if "DROP" in text.upper():
            desc = text.split(":", 1)[1].strip() if ":" in text else text
            return desc
        return ""

    # ---------- shared ----------

    def _generate(self, messages: list):
        cfg = self.cfg
        inputs = self.processor.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=True,
            return_dict=True, return_tensors="pt")
        inputs = {k: v.to(cfg.device) for k, v in inputs.items()}
        with self.torch.inference_mode():
            out = self.model.generate(
                **inputs, max_new_tokens=cfg.max_new_tokens,
                do_sample=cfg.do_sample, temperature=cfg.temperature,
                top_p=cfg.top_p, repetition_penalty=cfg.repetition_penalty)
        return inputs, out

    def _decode(self, generated: tuple) -> str:
        inputs, out = generated
        gen = out[0][inputs["input_ids"].shape[1]:]
        return self.processor.decode(gen, skip_special_tokens=True).strip()

    def run_video(self, video_path: str, begin: int, end: int) -> dict:
        """Run the configured task over [begin, end).

        commentary -> {"pred": {sec: text}}
        count      -> {"pred": {sec: "DROP: ..."}, "events": [...], "count": int}
        """
        frames = sample_frames(video_path, begin, end, self.cfg.frames_per_sec)
        out: dict = {}
        if self.cfg.task == "count":
            events = []
            pred = {}
            for t in range(begin, end):
                desc = self.chunk_event(t, frames[t], events)
                if desc.upper().startswith('DROP:'):
                    desc = desc.split(':', 1)[1].strip()
                if desc:
                    last = events[-1]["sec"] if events else -10**9
                    if t - last <= self.cfg.merge_gap:
                        # continuation of the same drop motion: merge
                        events[-1]["desc"] = desc
                        pred[str(t)] = f"DROP: {desc}"
                        print(f"[Sec {t}]: (merged) DROP: {desc}")
                        continue
                    events.append({"sec": t, "desc": desc})
                    pred[str(t)] = f"DROP: {desc}"
                    print(f"[Sec {t}]: DROP: {desc}  (total {len(events)})")
                else:
                    print(f"[Sec {t}]: <none>")
            out["pred"] = pred
            out["events"] = events
            out["count"] = len(events)
            if self.cfg.gt_count >= 0:
                out["gt_count"] = self.cfg.gt_count
                out["count_correct"] = len(events) == self.cfg.gt_count
                print(f"Predicted {len(events)}, ground truth {self.cfg.gt_count} -> "
                      f"{'CORRECT' if out['count_correct'] else 'MISMATCH'}")
        else:
            pred = {}
            spoken = []
            for t in range(begin, end):
                txt = self.chunk(t, frames[t], spoken)
                if txt:
                    pred[str(t)] = txt
                    spoken.append((t, txt))
                print(f"[Sec {t}]: {txt if txt else '<SILENCE>'}")
            out["pred"] = pred
        return out


def iter_records(args, infer: OpenVLStreamInfer):
    """Yield prediction/analysis records (official eval JSONL compatible)."""
    if args.dataset:
        from datasets import load_dataset
        ds = load_dataset(args.dataset, name=args.test_name, split="test")
        n = min(args.num_samples, len(ds)) if args.num_samples > 0 else len(ds)
        for idx in range(n):
            rec = ds[idx]
            video_path = os.path.join(args.video_dir, rec["video_path"])
            result = infer.run_video(video_path, rec["video_begin"], rec["video_end"])
            yield {
                "video_id": os.path.basename(video_path),
                "dataset_name": rec["dataset_name"],
                "begin": rec["video_begin"],
                "end": rec["video_end"],
                "idx": rec.get("idx", idx),
                "tag": rec.get("tag", ""),
                "model_id": args.model_id,
                "task": args.task,
                **result,
            }
    else:
        result = infer.run_video(args.video, args.begin, args.end)
        yield {
            "video_id": os.path.basename(args.video),
            "dataset_name": args.dataset_name,
            "begin": args.begin,
            "end": args.end,
            "idx": 0,
            "tag": args.tag,
            "model_id": args.model_id,
            "task": args.task,
            **result,
        }


def _cuda_available() -> bool:
    try:
        import torch
        return torch.cuda.is_available()
    except ImportError:
        return False


def parse_args():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--model", default="Qwen/Qwen2.5-VL-3B-Instruct")
    p.add_argument("--model_id", default=None, help="Label written into outputs")
    p.add_argument("--device", default=None)
    p.add_argument("--dtype", default="auto", choices=["auto", "float16", "bfloat16"])
    p.add_argument("--load_in_4bit", action="store_true")
    p.add_argument("--fps", type=int, default=2, dest="frames_per_sec")
    p.add_argument("--history_lines", type=int, default=8)
    p.add_argument("--task", default="commentary", choices=["commentary", "count"])
    p.add_argument("--scene", default="", help="Scene description for count task, e.g. 衣物投放")
    p.add_argument("--merge_gap", type=int, default=1,
                   help="Drops within this many seconds are merged as one event")
    p.add_argument("--gt_count", type=int, default=-1,
                   help="Ground-truth count; -1 to skip comparison")
    p.add_argument("--out", required=True)
    g = p.add_argument_group("input: single video")
    g.add_argument("--video")
    g.add_argument("--begin", type=int, default=0)
    g.add_argument("--end", type=int, default=30)
    g.add_argument("--tag", default="Solo commentators")
    g.add_argument("--dataset_name", default="custom")
    g = p.add_argument_group("input: LiveGamingBenchmark segments")
    g.add_argument("--dataset", default=None, help="HF dataset, e.g. oaaoaa/LiveGamingBenchmark")
    g.add_argument("--test_name", default="all_in_one")
    g.add_argument("--video_dir", default=None)
    g.add_argument("--num_samples", type=int, default=5)
    return p.parse_args()


def main():
    args = parse_args()
    if not args.video and not args.dataset:
        raise SystemExit("provide --video ... or --dataset ...")
    if args.model_id is None:
        args.model_id = os.path.basename(args.model) + "_stream"

    cfg = InferenceConfig(
        model=args.model,
        device=args.device or ("cuda" if _cuda_available() else "cpu"),
        dtype=args.dtype,
        load_in_4bit=args.load_in_4bit,
        task=args.task,
        scene=args.scene,
        merge_gap=args.merge_gap,
        gt_count=args.gt_count,
        frames_per_sec=args.frames_per_sec,
        history_lines=args.history_lines,
        model_id=args.model_id,
    )
    infer = OpenVLStreamInfer(cfg)
    records = list(iter_records(args, infer))
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    with open(out, "w", encoding="utf-8") as f:
        for rec in records:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    print(f"Wrote {len(records)} record(s) to {out}")
    for rec in records:
        if rec.get("task") == "count":
            print(f"  {rec['video_id']}: predicted count = {rec['count']}"
                  + (f" (gt={rec['gt_count']}, correct={rec['count_correct']})"
                     if "gt_count" in rec else ""))


if __name__ == "__main__":
    main()