"""Run inference with the judge LoRA adapter.

The script can either load files from a local directory or pull them
directly from the Hugging Face Hub. By default it points at the
companion repository ``NU-World-Model-Embodied-AI/phyjudge-9B``:

  # From the Hub (no clone needed):
  python infer.py --video demo.mp4 --caption "A ball rolls down a ramp." --metric SA
  python infer.py --video demo.mp4 --caption "A ball rolls down a ramp." --law gravity

  # From a local clone of the model repo:
  python infer.py --adapter-dir /path/to/local/clone --video demo.mp4 \
      --caption "A ball rolls down a ramp." --law gravity

It loads:
  - adapter_config.json to find the base model
  - adapter_model.safetensors through PEFT
  - subq+human.yaml to render the scoring prompt
"""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Any

import torch
import yaml
from peft import PeftModel
from transformers import AutoProcessor


# The tables below mirror evals/prompts/__init__.py (GENERAL_SUB_QUESTIONS),
# evals/physics_criteria.py (CRITERIA) and evals/sub_questions.py (SUB_QUESTIONS)
# in https://github.com/NU-World-Model-Embodied-AI/PhyGround, so the rendered
# prompts match `python -m evals.vlm_eval --prompt_config subq+human.yaml`.
GENERAL_SUB_QUESTIONS: dict[str, list[str]] = {
    "SA": [
        "Are the main objects in the caption present in the video?",
        "Are the key actions or interactions from the caption visible?",
        "Are important scene attributes and relationships preserved?",
        "Does the video avoid major contradictions to the caption?",
    ],
    "PTV": [
        "Do causes appear before their effects?",
        "Do physical events unfold in a plausible temporal order?",
        "Are motion transitions continuous rather than abrupt jumps or loops?",
        "Does the sequence avoid impossible reversals or repeated resets?",
    ],
    "persistence": [
        "Do objects maintain consistent existence throughout the video?",
        "Do objects keep a stable shape, size, color, and texture?",
        "Do objects avoid disappearing, appearing, or transforming unexpectedly?",
        "Do objects preserve identity through motion and brief occlusion?",
    ],
}


PHYSICAL_CRITERIA: dict[str, str] = {
    "gravity": "Do unsupported objects fall downward? Do thrown objects follow a curved trajectory? Does poured liquid fall with gravity?",
    "inertia": "Do stationary objects remain still unless acted upon? Do moving objects maintain their motion unless stopped by friction, collision, or an obstacle?",
    "momentum": "After collision, push, or pull, is the direction of motion reasonable? Ignore speed magnitude.",
    "impenetrability": "Do objects maintain impenetrability — no passing through each other?",
    "collision": "After impact, is there reasonable bounce/shatter/deformation? Does response match impact force?",
    "material": "Does each material respond according to its properties? (glass shatters, rubber bounces, metal is rigid, cloth deforms softly, etc.)",
    "buoyancy": "Do dense objects sink? Do wood/plastic float?",
    "displacement": "When you add more liquid or put an object into it, does the liquid level rise in a realistic way? Does it overflow when full?",
    "flow_dynamics": "Does the liquid's overall motion behave realistically over time — flowing along surfaces, spreading, draining naturally?",
    "boundary_interaction": "When the liquid hits a boundary such as a rock face, container wall, or floor, does it respond realistically? Do local splash, rebound, or split patterns on impact look physically plausible?",
    "fluid_continuity": "Does the liquid avoid disappearing or appearing out of nowhere? Small splashes that briefly break apart are okay.",
    "reflection": "Does the reflection roughly match objects and colors in the scene, and avoid completely unrelated content?",
    "shadow": "Are shadow directions consistent with light source? Do shadows move with objects?",
}


PHYSICAL_SUB_QUESTIONS: dict[str, list[str]] = {
    "gravity": [
        "Do any unsupported objects float upward or hover in mid-air?",
        "Does any heavy object drift or float down unrealistically slowly, as if it had no weight?",
        "Does any object fall sideways or in an unnatural direction instead of downward?",
    ],
    "inertia": [
        "Does any object spontaneously start moving without a visible cause?",
        "Does any moving object suddenly stop or reverse direction with no visible contact or force?",
    ],
    "momentum": [
        "Does any object fly off in a direction completely unrelated to how it was hit or pushed?",
        "Does the recoil direction contradict the direction of the applied force?",
    ],
    "impenetrability": [
        "Does any object pass through another solid object as if it were not there?",
        "Does any part of an object clip into a wall, floor, or another object's interior?",
    ],
    "collision": [
        "Does any object remain completely unaffected after being clearly hit by another object?",
        "Does the collision response look wildly too weak or too strong for the visible impact?",
        "Does any object shatter or deform dramatically from a very light touch with no reasonable force?",
    ],
    "material": [
        "Does any rigid material (glass, metal, stone) bend or stretch like rubber?",
        "Does any soft material (cloth, rubber, rope) behave as if it were completely rigid and unbending?",
    ],
    "buoyancy": [
        "Does any heavy object (metal, stone) float on the liquid surface?",
        "Does any light object (wood, cork, plastic) sink to the bottom?",
    ],
    "displacement": [
        "Does the liquid level remain completely unchanged when a large object is submerged?",
        "Does a full container spill or overflow when more liquid or an object is added?",
        "Does the liquid level behave in a clearly impossible way, such as dropping when volume is added?",
    ],
    "flow_dynamics": [
        "Does liquid flow uphill or against gravity without any force pushing it?",
        "On a flat surface, does liquid spread outward and become thinner over time?",
        "Does liquid suddenly stop flowing or freeze in place without an obvious reason?",
    ],
    "boundary_interaction": [
        "Does liquid ignore a solid boundary and continue moving as if nothing were there?",
        "Does liquid striking a surface produce visible droplets, spray, or ripples?",
        "Does liquid accumulate or pool on the wrong side of a barrier?",
    ],
    "fluid_continuity": [
        "Does a continuous pour or flow stay connected as a stream without sudden gaps?",
        "Does liquid disappear into nothing (not counting brief splashes)?",
        "Does liquid appear from nowhere with no visible source?",
    ],
    "reflection": [
        "Do reflections show completely unrelated content not present in the scene?",
        "Does a reflection remain completely static while the scene clearly changes around it?",
    ],
    "shadow": [
        "Do different shadows in the same scene point in contradictory directions?",
        "Does any shadow remain fixed in place while its object clearly moves?",
    ],
}


def load_json(path: Path) -> dict[str, Any]:
    with path.open() as f:
        return json.load(f)


def load_yaml(path: Path) -> dict[str, Any]:
    with path.open() as f:
        return yaml.safe_load(f)


def questions_block(questions: list[str]) -> str:
    return "\n".join(f"q{idx}: {question}" for idx, question in enumerate(questions, 1))


def build_prompt(
    cfg: dict[str, Any],
    caption: str,
    *,
    metric: str | None = None,
    law: str | None = None,
    criteria: str | None = None,
) -> tuple[str, str, str]:
    if metric:
        if metric not in GENERAL_SUB_QUESTIONS:
            raise ValueError(f"unknown metric: {metric}")
        prompt = cfg["eval_prompts"][metric].format(
            prompt=caption,
            questions_block=questions_block(GENERAL_SUB_QUESTIONS[metric]),
        )
        return cfg["system_prompt"], prompt, metric

    if not law:
        raise ValueError("either --metric or --law is required")
    if law not in PHYSICAL_CRITERIA:
        raise ValueError(f"unknown law: {law}")
    prompt = cfg["physical_template"].format(
        prompt=caption,
        law=law,
        criteria=criteria or PHYSICAL_CRITERIA[law],
        questions_block=questions_block(PHYSICAL_SUB_QUESTIONS[law]),
    )
    return cfg["system_prompt"], prompt, law


def load_base_model(base_id: str, dtype: torch.dtype, device_map: str):
    errors: list[str] = []
    for class_name in (
        "AutoModelForImageTextToText",
        "AutoModelForVision2Seq",
        "AutoModelForCausalLM",
    ):
        try:
            module = __import__("transformers", fromlist=[class_name])
            model_cls = getattr(module, class_name)
            return model_cls.from_pretrained(
                base_id,
                torch_dtype=dtype,
                device_map=device_map,
                trust_remote_code=True,
            )
        except Exception as exc:  # pragma: no cover - depends on local transformers version
            errors.append(f"{class_name}: {exc}")
    raise RuntimeError("failed to load base model:\n" + "\n".join(errors))


def resolve_adapter_dir(source: str) -> Path:
    """Return a local directory holding the adapter files.

    If ``source`` is a directory containing ``adapter_config.json`` it is used
    as-is. Otherwise ``source`` is interpreted as a HF Hub repo id and the
    snapshot is downloaded into the local cache.
    """
    candidate = Path(source)
    if candidate.is_dir() and (candidate / "adapter_config.json").exists():
        return candidate
    try:
        from huggingface_hub import snapshot_download
    except ImportError as exc:
        raise ImportError(
            "huggingface_hub is required to fetch the adapter from the Hub. "
            "Install it with: pip install huggingface_hub"
        ) from exc
    return Path(snapshot_download(repo_id=source))


def load_model(adapter_source: str, dtype: torch.dtype, device_map: str) -> tuple[Any, Any, Path]:
    adapter_dir = resolve_adapter_dir(adapter_source)
    adapter_cfg = load_json(adapter_dir / "adapter_config.json")
    base_id = adapter_cfg["base_model_name_or_path"]
    processor = AutoProcessor.from_pretrained(base_id, trust_remote_code=True)
    base = load_base_model(base_id, dtype=dtype, device_map=device_map)
    model = PeftModel.from_pretrained(base, adapter_dir)
    model.eval()
    return processor, model, adapter_dir


def build_messages(system_prompt: str, user_prompt: str, video_path: Path) -> list[dict[str, Any]]:
    return [
        {"role": "system", "content": system_prompt},
        {
            "role": "user",
            "content": [
                {"type": "video", "video": str(video_path)},
                {"type": "text", "text": user_prompt},
            ],
        },
    ]


def prepare_inputs(
    processor: Any,
    messages: list[dict[str, Any]],
    device: torch.device,
    *,
    fps: float,
    max_pixels: int,
) -> dict[str, Any]:
    # Qwen3.5 thinks by default; the judge answers with bare JSON, so turn it off
    # (same as `evals.vlm_eval --no-thinking`).
    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )

    try:
        from qwen_vl_utils import process_vision_info
    except ImportError as exc:
        raise ImportError(
            "qwen-vl-utils>=0.0.14 is required for local video inference. "
            "Install it with: pip install 'qwen-vl-utils[decord]>=0.0.14'"
        ) from exc

    for msg in messages:
        content = msg.get("content")
        if isinstance(content, list):
            for item in content:
                if item.get("type") == "video":
                    item.setdefault("fps", fps)
                    item.setdefault("max_pixels", max_pixels)

    # Qwen3.5 uses the Qwen3-VL video processor: frames are resized to multiples
    # of its patch size (16), and the processor needs per-video metadata to put
    # the right frame timestamps in the prompt (without it, it assumes 24 fps).
    patch_size = getattr(getattr(processor, "video_processor", None), "patch_size", 16)
    image_inputs, video_inputs, video_kwargs = process_vision_info(
        messages,
        return_video_kwargs=True,
        return_video_metadata=True,
        image_patch_size=patch_size,
    )
    video_metadata = None
    if video_inputs is not None:
        video_inputs, video_metadata = (list(x) for x in zip(*video_inputs))

    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        video_metadata=video_metadata,
        do_resize=False,
        padding=True,
        return_tensors="pt",
        **video_kwargs,
    )
    return inputs.to(device)


def decode_generated(processor: Any, inputs: dict[str, Any], generated_ids: torch.Tensor) -> str:
    input_len = inputs["input_ids"].shape[1]
    generated_ids = generated_ids[:, input_len:]
    return processor.batch_decode(
        generated_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0].strip()


def parse_score(text: str, key: str) -> int | None:
    match = re.search(r"\{.*?\}", text, flags=re.S)
    if match:
        try:
            obj = json.loads(match.group(0))
            value = obj.get(key)
            if isinstance(value, int) and 1 <= value <= 5:
                return value
        except json.JSONDecodeError:
            pass
    match = re.search(rf'"?{re.escape(key)}"?\s*:\s*([1-5])', text)
    if match:
        return int(match.group(1))
    return None


def dtype_from_name(name: str) -> torch.dtype:
    if name == "bfloat16":
        return torch.bfloat16
    if name == "float16":
        return torch.float16
    if name == "float32":
        return torch.float32
    raise ValueError(f"unsupported dtype: {name}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Infer with the judge adapter.")
    parser.add_argument(
        "--adapter-dir",
        default="NU-World-Model-Embodied-AI/phyjudge-9B",
        help=(
            "Local directory with adapter_config.json + adapter_model.safetensors "
            "+ subq+human.yaml, or a HF Hub repo id "
            "(default: NU-World-Model-Embodied-AI/phyjudge-9B)."
        ),
    )
    parser.add_argument("--video", required=True, type=Path)
    parser.add_argument("--caption", required=True)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--metric", choices=["SA", "PTV", "persistence"])
    group.add_argument("--law", choices=sorted(PHYSICAL_CRITERIA))
    parser.add_argument("--criteria", help="Override physical-law criterion text.")
    parser.add_argument("--max-new-tokens", type=int, default=64)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--fps", type=float, default=2.0)
    parser.add_argument("--max-pixels", type=int, default=360 * 640)
    parser.add_argument("--dtype", choices=["bfloat16", "float16", "float32"], default="bfloat16")
    parser.add_argument("--device-map", default="auto")
    parser.add_argument("--print-prompt", action="store_true")
    args = parser.parse_args()

    if not args.video.is_file():
        raise FileNotFoundError(args.video)

    dtype = dtype_from_name(args.dtype)
    processor, model, adapter_dir = load_model(
        args.adapter_dir, dtype=dtype, device_map=args.device_map
    )

    prompt_cfg = load_yaml(adapter_dir / "subq+human.yaml")
    system_prompt, user_prompt, score_key = build_prompt(
        prompt_cfg,
        args.caption,
        metric=args.metric,
        law=args.law,
        criteria=args.criteria,
    )

    if args.print_prompt:
        print("SYSTEM:")
        print(system_prompt)
        print("\nUSER:")
        print(user_prompt)
        print()
    device = next(model.parameters()).device
    messages = build_messages(system_prompt, user_prompt, args.video)
    inputs = prepare_inputs(
        processor,
        messages,
        device,
        fps=args.fps,
        max_pixels=args.max_pixels,
    )

    generation_kwargs: dict[str, Any] = {
        "max_new_tokens": args.max_new_tokens,
        "do_sample": args.temperature > 0,
        "temperature": args.temperature if args.temperature > 0 else None,
    }
    generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None}

    with torch.inference_mode():
        generated_ids = model.generate(**inputs, **generation_kwargs)

    raw = decode_generated(processor, inputs, generated_ids)
    score = parse_score(raw, score_key)
    print(json.dumps({"key": score_key, "score": score, "raw": raw}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()