"""Stage 2 of the VIPER-19K pipeline: three-axis physics annotation.

The paper annotates every surviving clip on three axes -- **material**,
**trajectory** and **physical impact**.  Those labels are what the pairing stage
buckets on, so they are drawn from a closed taxonomy (free-form text would not
bucket).  We additionally keep a one-sentence physics summary, which is used as
the target prompt during training.

Run with ``--shard i --num_shards N`` to spread over GPUs.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import numpy as np
import torch

MATERIALS = [
    "rigid_solid", "elastic_deformable", "plastic_deformable", "brittle",
    "granular", "liquid", "viscous_liquid", "soft_body_cloth", "gas_smoke",
]
TRAJECTORIES = [
    "free_fall", "projectile_arc", "rolling", "sliding", "swinging_pendulum",
    "rotating_spinning", "flowing_pouring", "oscillating_vibrating",
    "expanding_dispersing", "in_place_deformation",
]
IMPACTS = [
    "collision_impact", "compression_squeeze", "fracture_shatter",
    "dent_deformation", "bounce_rebound", "splash_dispersal",
    "melting_phase_change", "tearing_rupture", "crumbling", "no_impact",
]

PROMPT = f"""Analyse the physical process in this video and answer STRICTLY as JSON.

Choose exactly one value for each axis from these closed sets:
material: {MATERIALS}
trajectory: {TRAJECTORIES}
physical_impact: {IMPACTS}

Return this JSON object and nothing else:
{{
  "material": "<one of material>",
  "trajectory": "<one of trajectory>",
  "physical_impact": "<one of physical_impact>",
  "physics_summary": "<one sentence describing the physical dynamics: how the \
object moves, deforms and interacts. Describe the PHYSICS, not the appearance \
or identity of the objects.>"
}}"""


def sample_frames(path: str, n: int = 8) -> np.ndarray:
    import imageio.v3 as iio

    v = iio.imread(path, plugin="pyav")            # [T,H,W,3] uint8
    idx = np.linspace(0, len(v) - 1, n).astype(int)
    return v[idx]


def parse_json(text: str) -> dict | None:
    text = text.strip()
    if "```" in text:
        text = text.split("```")[1]
        if text.startswith("json"):
            text = text[4:]
    s, e = text.find("{"), text.rfind("}")
    if s < 0 or e < 0:
        return None
    try:
        return json.loads(text[s:e + 1])
    except json.JSONDecodeError:
        return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--clips", default="data/filtered_clips.jsonl")
    ap.add_argument("--out", default="data/annotated.jsonl")
    ap.add_argument("--mllm", default="models/Qwen3-VL-4B-Instruct")
    ap.add_argument("--num_frames", type=int, default=8)
    ap.add_argument("--shard", type=int, default=0)
    ap.add_argument("--num_shards", type=int, default=1)
    ap.add_argument("--limit", type=int, default=0)
    args = ap.parse_args()

    from transformers import AutoProcessor, Qwen3VLForConditionalGeneration

    rows = [json.loads(l) for l in open(args.clips)]
    if args.limit:
        rows = rows[: args.limit]
    rows = rows[args.shard::args.num_shards]

    proc = AutoProcessor.from_pretrained(args.mllm)
    model = Qwen3VLForConditionalGeneration.from_pretrained(
        args.mllm, dtype=torch.bfloat16, device_map="cuda",
        attn_implementation="sdpa",
    ).eval()

    out_path = args.out if args.num_shards == 1 else f"{args.out}.{args.shard}"
    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
    n_ok = 0
    with open(out_path, "w") as f:
        for i, r in enumerate(rows):
            try:
                frames = sample_frames(r["video"], args.num_frames)
            except Exception as e:
                print(f"[skip decode] {r['id']}: {e}", flush=True)
                continue

            messages = [{"role": "user", "content": [
                {"type": "video"}, {"type": "text", "text": PROMPT}]}]
            text = proc.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True)
            inputs = proc(text=[text], videos=[frames], return_tensors="pt").to("cuda")

            with torch.inference_mode():
                gen = model.generate(**inputs, max_new_tokens=256, do_sample=False)
            reply = proc.batch_decode(
                gen[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0]

            ann = parse_json(reply)
            if not ann or not all(k in ann for k in
                                  ("material", "trajectory", "physical_impact")):
                print(f"[skip parse] {r['id']}", flush=True)
                continue
            # Reject hallucinated out-of-taxonomy values.
            if (ann["material"] not in MATERIALS
                    or ann["trajectory"] not in TRAJECTORIES
                    or ann["physical_impact"] not in IMPACTS):
                print(f"[skip taxonomy] {r['id']}: {ann}", flush=True)
                continue

            r.update(ann)
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
            f.flush()
            n_ok += 1
            if i % 25 == 0:
                print(f"[{args.shard}] {i}/{len(rows)} ok={n_ok}", flush=True)

    print(f"[{args.shard}] DONE ok={n_ok}/{len(rows)} -> {out_path}", flush=True)


if __name__ == "__main__":
    main()