"""
vlm_judge.py — uses Qwen3-VL-4B-Instruct (a real vision-language model)
to automatically judge a sketch's generated keyframe sequence.

Runs in its OWN separate environment (Python 3.9+, transformers>=4.57.0),
NOT the main "mosketch" environment used for classify/narrate/deform —
Qwen3-VL requires transformers>=4.57.0, which requires Python>=3.9, and
the main environment is confirmed pinned to Python 3.8 (max installable
transformers there is 4.46.3). pipeline.py calls this script as a
SEPARATE PROCESS using the judge environment's python interpreter, not
as an in-process import, specifically so the two environments never need
to coexist in the same Python process.

Takes the frames/{sketch}/kf0.png ... kf4.png images produced by
pipeline.py and asks the VLM to score them on:
  1. Plausibility — does the pose progression look like coherent,
     physically believable motion (not static, not erratic)?
  2. Faithfulness — if a caption is available, does the sequence match
     what was supposed to happen?

Connectivity was deliberately dropped as a criterion: ARAP guarantees it
by construction (that's what the seam-constraint work was for), so asking
a VLM to check for it is redundant. A "completeness" criterion (missing
limbs from incomplete source strokes) was also tried and removed — see
the pipeline discussion for why that's a structural limitation of
stroke-deformation rather than something the judge can meaningfully score.

Requires transformers >= 4.57.0 and Python >= 3.9 — run this from the
separate judge environment (see docstring above), not the mosketch env.

Usage:
    python vlm_judge.py dog9
    # frames auto-derived from frames/dog9/, caption auto-pulled by sketch name

    python vlm_judge.py dog9 --frames-dir custom/path --caption "custom text"
    # override either default individually

    python vlm_judge.py dog9 --no-caption
    # skip caption lookup entirely; faithfulness scored as N/A
"""

import argparse
import json
import os
import re

# Same default as pipeline.py — so the caption is auto-pulled by sketch
# name alone, no need to pass --caption-file every time.
CAPTION_FILE_DEFAULT = "/user/HS400/rk01499/my_scratch/sketch/data/raw/60sketches/caption.txt"


def build_joint_legend_text(objects_info):
    """
    Rest-pose joint legend per ARAP object — index, rest (x,y), anchor flag.
    Gives the judge a stable vocabulary ("joint_2") to point at in
    joint_feedback, grounded against the SAME indices the generator itself
    reads/writes (deform_outputs[obj][kf]["joint_i"]) — so a judge complaint
    about joint_2 is directly actionable by the next narrate+deform call
    without any translation step.
    """
    sections = []
    for name, info in objects_info.items():
        lines = []
        for i, j in enumerate(info["joints"]):
            anchor_tag = "  <-- ANCHOR (should stay near rest every keyframe)" if i == info["anchor_idx"] else ""
            lines.append(f"    joint_{i}: rest (x={j[0]:.1f}, y={j[1]:.1f}){anchor_tag}")
        sections.append(f'Object "{name}" joints:\n' + "\n".join(lines))
    return "\n\n".join(sections)


def build_joint_targets_text(objects_info, deform_outputs, n_keyframes=5):
    """
    Per-keyframe joint TARGET coordinates (post-clip, what the generator
    actually asked for this attempt) for each ARAP object — NOT the
    deformed stroke geometry. Deliberately excludes per-stroke coordinate
    text: that format gave the judge no joint identity to anchor feedback
    to, which is the specific problem this rework addresses. Per-keyframe
    deformed stroke text may be added back later as a separate experiment
    once joint-level feedback quality is evaluated on its own.
    """
    sections = []
    for name, info in objects_info.items():
        obj_targets = (deform_outputs or {}).get(name)
        if not obj_targets:
            continue
        kf_lines = []
        for kf in range(n_keyframes):
            kf_key = f"kf{kf}"
            if kf_key not in obj_targets:
                continue
            parts = ", ".join(f"{j}=({v[0]:.1f},{v[1]:.1f})" for j, v in obj_targets[kf_key].items())
            kf_lines.append(f"    {kf_key}: {parts}")
        if kf_lines:
            sections.append(f'Object "{name}" joint TARGETS this attempt:\n' + "\n".join(kf_lines))
    return "\n\n".join(sections)


def build_bbox_trajectory_text(object_names, real_trajectories, n_keyframes=5):
    """
    Per-keyframe bounding-box position for EVERY object (ARAP and
    TRAJ_ONLY alike) — sourced from preprocessed trajectory data, i.e.
    FIXED ground truth, not something the judge should critique or the
    generator can change. Included purely as spatial context so the judge
    isn't reasoning about joint deformation in a positional vacuum (e.g.
    "is the dog's head reaching toward the frisbee" requires knowing
    where the frisbee bbox actually is at that keyframe).
    """
    from lib import bbox_deltas
    sections = []
    for name in object_names:
        traj = real_trajectories.get(name)
        if traj is None:
            continue
        dx_vals, dy_vals = bbox_deltas(traj)
        parts = ", ".join(f"kf{kf}=(dx={dx_vals[kf]:.1f},dy={dy_vals[kf]:.1f})" for kf in range(n_keyframes))
        sections.append(f'Object "{name}" bbox trajectory (FIXED, given, do not critique or ask to change): {parts}')
    return "\n".join(sections)


def compute_joint_feedback_deltas(joint_feedback, deform_outputs):
    """
    Turns the judge's target_coords into an actual pixel delta by
    SUBTRACTION — not model math. The judge only has to say "the mouth
    stroke is at this coordinate" (a lookup/grounding task); this function
    does target - actual arithmetic exactly, so the number the generator
    receives can't inherit any model arithmetic error, only whatever
    error is in the judge's chosen target_coords itself.

    Mutates each entry in place, adding:
      - "current_coords": {"x":..., "y":...} — where the joint actually
        is right now (from deform_outputs), for reference
      - "delta_px": {"dx":..., "dy":...} — target - current, or None if
        target_coords was null/missing/malformed (ungrounded feedback is
        left as prose-only, NOT silently defaulted to a zero or guessed
        delta)

    Entries with joint=None handling, missing objects, or missing
    keyframes in deform_outputs are left with delta_px=None rather than
    raising — a judge referencing a keyframe/joint that doesn't exist in
    deform_outputs is a validation problem to surface, not something to
    paper over with a computed number.
    """
    for entry in joint_feedback or []:
        entry["current_coords"] = None
        entry["delta_px"] = None

        obj_name = entry.get("object")
        joint_idx = entry.get("joint")
        kf = entry.get("keyframe")
        target = entry.get("target_coords")

        if target is None or target.get("x") is None or target.get("y") is None:
            continue  # judge explicitly said it couldn't ground this — respect that, don't invent a delta

        # anchored_to must disambiguate WHICH geometry block the coordinate came from — both the
        # per-keyframe deformed block and the rest-pose block reuse "stroke_N" numbering, so a bare
        # "stroke_2" can't be trusted to be the right one. Rather than trust an ambiguous anchor and
        # silently compute a delta that might be using a stale rest-pose coordinate for a moved
        # keyframe, treat it the same as ungrounded: no delta, falls back to prose-only feedback.
        anchor = entry.get("anchored_to", "")
        disambiguated = isinstance(anchor, str) and any(
            term in anchor.lower() for term in ("kf", "keyframe", "deformed", "rest-pose", "rest pose"))
        if not disambiguated:
            continue

        obj_targets = (deform_outputs or {}).get(obj_name)
        if obj_targets is None:
            continue
        kf_key = f"kf{kf}"
        joint_key = f"joint_{joint_idx}"
        current = obj_targets.get(kf_key, {}).get(joint_key)
        if current is None:
            continue

        entry["current_coords"] = {"x": current[0], "y": current[1]}
        entry["delta_px"] = {
            "dx": round(target["x"] - current[0], 1),
            "dy": round(target["y"] - current[1], 1),
        }
    return joint_feedback


def build_judge_prompt(sketch_name, caption=None, dino_stagnation=None, dino_temporal=None, clip_scores=None,
                        cap_utilization=None, objects_info=None, deform_outputs=None,
                        real_trajectories=None, all_object_names=None, n_keyframes=5,
                        deformed_geometry_text=None, narratives=None):
    """
    dino_stagnation: the dict returned by dino_similarity.stagnation_score()
        comparing this attempt to the previous one (None on attempt 1).
    dino_temporal: list of consecutive-keyframe DINOv2 similarities within
        THIS attempt (from dino_similarity.temporal_consistency()).
    clip_scores: the dict returned by clip_score.compute_sequence_clip_scores()
        (None if no caption was available to score against).
    cap_utilization: {obj_name: mean_fraction_0_to_1} — how much of the
        allowed movement range each object's joints actually used.
        CONFIRMED on real hardware (horsecar5's walking person) that Qwen
        can propose displacement well within its allowance without any
        signal catching it — a leg moved only 18% of its permitted range,
        invisible at render scale despite technically being "movement."

    narratives: {obj_name: [n_keyframes pose description strings]} — the
        GENERATOR's OWN stated intent for each keyframe (same data written
        to {name}_narratives.json). When given, FAITHFULNESS is judged as
        narrative-vs-EXECUTION (does what the generator claimed happen
        actually match the rendered images/coordinates?) instead of
        narrative-vs-caption. This catches a failure mode the caption-only
        check cannot: a generator that writes a plausible story but
        produces timid or wrong joint targets that don't execute it — the
        old check only ever looked at the final image against the
        caption, never verified the generator's own stated intent was
        actually carried out. If omitted, falls back to the original
        image-vs-caption faithfulness check.

    objects_info / deform_outputs / real_trajectories / all_object_names:
        if given, the judge is shown (a) rest-pose stroke geometry per ARAP
        object, (b) a rest-pose joint index legend, (c) this attempt's
        joint TARGET coordinates per keyframe, and (d) every object's fixed
        bbox trajectory. This is what makes joint_feedback possible: the
        judge has a joint vocabulary to point at instead of only prose
        description. If none of these are given, the prompt falls back to
        the original image-only behavior and joint_feedback is not requested.

    deformed_geometry_text: PER-KEYFRAME actual deformed stroke geometry
        (from pipe.py's build_all_keyframes_deformed_geometry_text) — the
        judge is instructed to anchor target_coords against THIS, not the
        rest-pose geometry, because a stroke's rest coordinate is only
        correct for whichever keyframe happens to match rest; every other
        keyframe needs that stroke's TRUE position at that specific
        keyframe. If omitted, falls back to rest-pose-only anchoring
        (less accurate per-keyframe, but still grounded rather than a
        free guess).
    """
    caption_block = ""
    if caption:
        caption_block = (f'\nThe sketch is supposed to depict: "{caption}" (background context for '
                          f'PLAUSIBILITY only — FAITHFULNESS below is judged against the narrative, not this)\n')

    narrative_block = ""
    faithfulness_criterion = (
        'FAITHFULNESS (1-5, or "N/A" if no caption was given above): Does the sequence match what the caption '
        'describes should be happening? 5 = clearly matches, 1 = unrelated to the description.'
    )
    if narratives:
        narrative_parts = []
        for obj_name, kf_list in narratives.items():
            kf_lines = "\n".join(f"    kf{i}: {desc}" for i, desc in enumerate(kf_list))
            narrative_parts.append(f'Object "{obj_name}":\n{kf_lines}')
        narrative_block = (
            "\nThe GENERATOR's OWN stated intent for each keyframe (what it claims is happening, written "
            "BEFORE rendering):\n" + "\n\n".join(narrative_parts) + "\n"
        )
        faithfulness_criterion = (
            'FAITHFULNESS (1-5): Does what ACTUALLY happened (the rendered images and joint coordinates above) '
            'match what the generator ITSELF claimed would happen (the stated narrative above) — NOT whether '
            'it matches the caption. 5 = the narrative\'s claims are clearly executed in the actual poses/'
            'coordinates. 1 = the narrative describes an action but the actual joint positions barely differ '
            'from rest, or move somewhere unrelated to what was claimed. This catches a generator that writes '
            'a plausible story without actually executing it.'
        )

    metrics_lines = []
    if dino_stagnation is not None:
        sims = ", ".join(f"kf{i}={s:.3f}" for i, s in enumerate(dino_stagnation["per_keyframe_similarity"]))
        metrics_lines.append(
            f"- DINOv2 similarity to the PREVIOUS attempt, per keyframe (1.0 = identical, this attempt "
            f"made no change at all; lower = more different): {sims} (mean {dino_stagnation['mean_similarity']:.3f})"
        )
    if dino_temporal is not None:
        temp_str = ", ".join(f"kf{i}->kf{i+1}={s:.3f}" for i, s in enumerate(dino_temporal))
        metrics_lines.append(
            f"- DINOv2 similarity between CONSECUTIVE keyframes within this attempt (lower = more visual "
            f"change between those two frames, higher = little change): {temp_str}"
        )
    if cap_utilization is not None:
        util_str = ", ".join(f"{obj}={frac*100:.0f}%" for obj, frac in cap_utilization.items())
        metrics_lines.append(
            f"- Movement allowance used, per object (how much of the permitted displacement range each "
            f"object's joints actually used — low values mean technically-nonzero but visually negligible "
            f"motion): {util_str}"
        )
    if clip_scores is not None:
        clip_str = ", ".join(f"kf{i}={s:.3f}" for i, s in enumerate(clip_scores["per_keyframe_clip_score"]))
        metrics_lines.append(
            f"- CLIP image-caption similarity, per keyframe (higher = the image better matches the caption "
            f"text; CLIP was trained on photographs, not sketches, so treat this as a rough signal, not "
            f"ground truth): {clip_str} (mean {clip_scores['mean_clip_score']:.3f})"
        )
    metrics_block = ""
    if metrics_lines:
        metrics_block = "\nObjective measurements computed for this attempt (use these to INFORM your " \
                         "reasoning and cite them in your notes where relevant, but judge primarily from " \
                         "what you actually see in the images):\n" + "\n".join(metrics_lines) + "\n"

    # coordinate/geometry context — only built if the caller supplied it
    has_coord_context = objects_info is not None
    geometry_block = ""
    joint_feedback_instruction = ""
    joint_feedback_schema = ""
    if has_coord_context:
        rest_geometry_parts = []
        for name, info in objects_info.items():
            strokes = info.get("strokes")
            if strokes:
                from lib import build_stroke_geometry_text
                rest_geometry_parts.append(
                    f'Object "{name}" rest-pose strokes (undeformed, for identifying which body part each '
                    f'joint is near):\n{build_stroke_geometry_text(strokes, n_points=2)}'
                )
        rest_geometry_text = "\n\n".join(rest_geometry_parts)

        joint_legend_text = build_joint_legend_text(objects_info)
        joint_targets_text = build_joint_targets_text(objects_info, deform_outputs, n_keyframes=n_keyframes)
        bbox_text = ""
        if real_trajectories is not None and all_object_names is not None:
            bbox_text = build_bbox_trajectory_text(all_object_names, real_trajectories, n_keyframes=n_keyframes)

        if deformed_geometry_text:
            anchor_source_block = f"""
PER-KEYFRAME deformed stroke geometry — each stroke's ACTUAL position AT THAT SPECIFIC KEYFRAME (not rest).
THIS is the correct anchor source for target_coords: a stroke's position can differ keyframe to keyframe, so
when grounding a target for kf3, use kf3's coordinates below, not the rest-pose ones:
{deformed_geometry_text}

Rest-pose stroke geometry (undeformed — for orientation/context only; do NOT use these coordinates as a
target_coords anchor for a specific keyframe, use the per-keyframe geometry above instead):
{rest_geometry_text}"""
            anchor_instruction = (
                "find that stroke's EXACT coordinates AT THE SAME KEYFRAME you are giving feedback for, in the "
                "per-keyframe deformed geometry above — NOT the rest-pose geometry below, which is only correct "
                "for keyframes that happen to match rest and is included for orientation only. When you cite "
                "anchored_to, you MUST state which block you took the coordinate from, e.g. 'kf4 deformed "
                "geometry, stroke_2' — NOT just 'stroke_2', since both blocks use the same stroke numbering and "
                "a bare stroke number does not say which one you actually used"
            )
        else:
            anchor_source_block = f"""
Rest-pose stroke geometry, with EXACT coordinates (use the IMAGES to identify what each stroke/region actually
is — e.g. "that stroke is the mouth", "that stroke is the spoon" — then use these coordinates as the real,
exact position of that thing. Do NOT estimate a position by eye from the image when a matching coordinate is
available here. NOTE: these are REST-POSE coordinates only — no per-keyframe geometry was provided, so treat
any target_coords derived from these as approximate for keyframes that have moved far from rest):
{rest_geometry_text}"""
            anchor_instruction = (
                "find that stroke's EXACT coordinates in the rest-pose stroke geometry above (no per-keyframe "
                "geometry was provided this call, so this is an approximation for keyframes far from rest). "
                "State 'rest-pose geometry, stroke_N' in anchored_to"
            )

        geometry_block = f"""
The FIRST attached image is the rest pose (undeformed). The remaining {n_keyframes} attached images are this
attempt's actual rendered keyframes, in order.
{anchor_source_block}

Joint index legend (rest positions, same coordinate space as the strokes above, y increases DOWNWARD):
{joint_legend_text}

This attempt's joint TARGET positions per keyframe (what was actually requested — compare against rest to see
what moved and how far):
{joint_targets_text}

Every object's bounding-box trajectory (FIXED, preprocessed ground truth — this is GIVEN, not something the
generator controls or you should critique; use it only as spatial context, e.g. to check whether a reaching
pose is plausible given where the other object actually is):
{bbox_text}
"""
        arap_object_names = list(objects_info.keys())
        arap_list_str = ", ".join(f'"{n}"' for n in arap_object_names)
        # REVERTED from a mandatory status/exactly-one-entry-per-object schema — CONFIRMED on real
        # hardware (football7, real 3-object scene) that forcing a binary needs_correction/
        # no_change_needed classification for every object, every time, caused the model to default
        # to "no_change_needed" across the board even when its OWN overall_verdict/faithfulness_score
        # said something was wrong — a worse failure than the original "silently skips an object"
        # problem this was built to fix, since it produces confident-looking but false coverage
        # instead of an honest gap. Back to an open list: report whatever objects actually have a
        # real problem, omit the rest, no forced per-object judgment call.
        joint_feedback_instruction = (
            "\n4. For any ARAP object whose deformation looks wrong, give PRECISE corrective feedback as a "
            f"coordinate, not a description: name the object (from this scene's ARAP objects: {arap_list_str}), "
            "the specific joint_i index (from the legend above), and the SINGLE worst keyframe (not every "
            "keyframe it's visible in). To decide the target coordinate: look at the images to identify which "
            "stroke(s) the joint should move toward or align with (e.g. the mouth, another object's edge), "
            f"then {anchor_instruction}, and use that coordinate (or a point clearly interpolated between two "
            "nearby stroke coordinates) as target_coords. Do NOT invent a coordinate you cannot trace back to "
            "the geometry text or joint legend — if you cannot find a grounded coordinate, describe the issue "
            "in words instead of guessing a number, and state that no grounded target was found. Always name "
            "what you anchored the target to in anchored_to. Do NOT create an entry for any object name other "
            "than this scene's ARAP objects listed above — appearing elsewhere in this prompt (caption, "
            "narrative, bbox trajectory) does NOT mean it has joints. If MULTIPLE objects have a real problem, "
            "include an entry for EACH of them. If nothing specific is wrong anywhere, return an empty list.\n"
        )
        joint_feedback_schema = (
            ',\n  "joint_feedback": [\n    {"object": "<name>", "joint": <joint index int>, '
            '"keyframe": <single int, the worst one>, "issue": "<one sentence>", '
            '"target_coords": {"x": <number or null if ungrounded>, "y": <number or null if ungrounded>}, '
            '"anchored_to": "<what stroke/coordinate you traced this to, or \'none - description only\'>"}\n  ]'
        )

    return f"""You are judging a sequence of {n_keyframes} keyframe images generated for an animated sketch named "{sketch_name}". The images are provided in order (keyframe 0 through keyframe {n_keyframes - 1}), showing an object's pose changing over time.

IMPORTANT: these are SPARSE keyframes sampled across the ENTIRE action from start to finish — NOT consecutive video frames a fraction of a second apart. Large, dramatic differences between consecutive keyframes are normal and correct; the actual in-between motion frames are generated separately later by a different model. Do NOT judge these like near-continuous animation frames, and do NOT penalize a keyframe for looking very different from its neighbor — that is expected. Judge whether the SEQUENCE of poses tells a coherent story of the action from start to finish.
{caption_block}{narrative_block}{metrics_block}{geometry_block}
Evaluate the sequence on these criteria:

1. PLAUSIBILITY (1-5): Does the sequence of poses, taken as sparse waypoints across the whole action, tell a coherent and physically believable story from start to finish? 5 = each keyframe is a sensible, meaningfully different stage of the action, in a believable order. 1 = static (no visible change between keyframes at all) or the poses themselves are physically impossible/nonsensical — NOT simply "very different from the previous keyframe," which is expected and correct.

2. {faithfulness_criterion}

3. QUALITY (1-5): Does each frame still look like a clean, recognizable line drawing of the object — NOT the motion, just the visual rendering itself? 5 = clean, coherent lines throughout, 1 = garbled, noisy, or unrecognizable as the object in any frame. THIS INCLUDES checking each joint marked <-- ANCHOR in the legend above: an anchor is supposed to stay near its rest position in EVERY keyframe, so also check whether the region around each anchor (and any other part of the object that should stay visually stable while only the moving joints change) has held its shape correctly across the sequence — not just whether lines are crisp. Unintended distortion of a part that was supposed to remain stable (e.g. a face warping, a torso shifting shape, when only a hand/limb was meant to move) is a QUALITY failure even if the lines themselves are still clean and non-garbled — score it low (1-2) if you see this, even if nothing looks "noisy."
{joint_feedback_instruction}
IMPORTANT — overall_verdict is a single passage covering the WHOLE scene, used for logging/summary only. The
generator receives the SAME global feedback string for every object (per-object text scoping was tried and
REMOVED — numeric joint_feedback, not prose, is the actual per-object correction mechanism; keep overall_verdict
to ONE whole-scene passage, do not produce per-object prose breakdowns, to keep response length manageable
given joint_feedback below already requires one entry per ARAP object).

Respond with ONLY a JSON object, no other text, in this exact format:
{{
  "plausibility_score": <1-5>,
  "plausibility_notes": "<one sentence, for logging/reasoning only>",
  "faithfulness_score": <1-5 or "N/A">,
  "faithfulness_notes": "<one sentence, for logging/reasoning only>",
  "quality_score": <1-5>,
  "quality_notes": "<one sentence — if score is low due to unintended distortion of a supposedly-stable part (not just messy lines), say specifically what part and which keyframe(s)>",
  "overall_verdict": "<ONE whole-scene passage, for logging/summary only>"{joint_feedback_schema}
}}
"""


def parse_judge_response(response_text):
    match = re.search(r"\{.*\}", response_text, re.DOTALL)
    if not match:
        raise ValueError("No JSON object found in response:\n" + response_text)
    return json.loads(match.group(0))


def validate_judge_response(parsed, valid_arap_objects=None):
    """
    valid_arap_objects: the actual list/set of ARAP object names that were
    given a joint legend in this call's prompt (e.g. list(objects_info.keys())).
    When given, any joint_feedback entry naming an object OUTSIDE this set
    is a hallucination — the object either doesn't exist in this scene or
    has no joints (e.g. a TRAJ_ONLY object) — and is flagged as a
    validation problem rather than silently trusted. This is the code-side
    backstop for the prompt's "only reference these ARAP objects" instruction:
    the prompt asks the model not to do this, but doesn't structurally
    prevent it, so this check exists for when it does anyway.
    """
    required = ["plausibility_score", "faithfulness_score", "quality_score", "overall_verdict"]
    problems = [f"missing key '{k}'" for k in required if k not in parsed]
    for score_key in ["plausibility_score", "quality_score"]:
        val = parsed.get(score_key)
        if val is not None and not (isinstance(val, (int, float)) and 1 <= val <= 5):
            problems.append(f"'{score_key}' = {val!r}, expected a number 1-5")

    # joint_feedback is OPTIONAL — only present when the prompt included
    # coordinate context (objects_info etc. were passed to build_judge_prompt).
    # Absence is not itself a problem; malformed entries are.
    if "joint_feedback" in parsed:
        jf = parsed["joint_feedback"]
        if not isinstance(jf, list):
            problems.append(f"'joint_feedback' = {jf!r}, expected a list")
        else:
            required_entry_keys = ["object", "joint", "keyframe", "issue", "target_coords", "anchored_to"]
            for i, entry in enumerate(jf):
                if not isinstance(entry, dict):
                    problems.append(f"joint_feedback[{i}] is not an object: {entry!r}")
                    continue
                missing = [k for k in required_entry_keys if k not in entry]
                if missing:
                    problems.append(f"joint_feedback[{i}] missing keys: {missing}")
                # hallucination check: does this entry's object actually have a joint legend in
                # this call's prompt? A TRAJ_ONLY or nonexistent object name here means the model
                # invented a joint for something that structurally cannot have one — this is
                # flagged, not silently trusted, since it wastes the judge's feedback on an
                # object that can never use it and can crowd out real feedback for objects that
                # actually needed it.
                if valid_arap_objects is not None:
                    obj = entry.get("object")
                    if obj is not None and obj not in valid_arap_objects:
                        problems.append(
                            f"joint_feedback[{i}]['object'] = {obj!r} is not an ARAP object with a joint "
                            f"legend in this prompt (valid: {sorted(valid_arap_objects)}) — likely "
                            f"hallucinated, e.g. a TRAJ_ONLY object mistakenly given joint feedback"
                        )
                if "joint" in entry and not isinstance(entry["joint"], int):
                    problems.append(f"joint_feedback[{i}]['joint'] = {entry['joint']!r}, expected an int")
                if "keyframe" in entry and not isinstance(entry["keyframe"], int):
                    problems.append(f"joint_feedback[{i}]['keyframe'] = {entry['keyframe']!r}, expected an int")
                tc = entry.get("target_coords")
                if tc is not None:
                    if not isinstance(tc, dict) or "x" not in tc or "y" not in tc:
                        problems.append(f"joint_feedback[{i}]['target_coords'] = {tc!r}, expected {{'x':num,'y':num}} or null")
                    else:
                        for axis in ("x", "y"):
                            v = tc.get(axis)
                            if v is not None and not isinstance(v, (int, float)):
                                problems.append(f"joint_feedback[{i}]['target_coords']['{axis}'] = {v!r}, expected a number or null")
                        # target_coords was given but anchored_to doesn't say WHICH block (per-keyframe
                        # deformed vs rest-pose) it came from — both blocks reuse the same "stroke_N"
                        # numbering, so a bare "stroke_2" can't be trusted to be the right one. This is a
                        # validation problem because a wrong-block anchor silently produces a wrong
                        # but confident-looking delta_px downstream, with no other signal catching it.
                        anchor = entry.get("anchored_to", "")
                        has_x = tc.get("x") is not None
                        if has_x and isinstance(anchor, str) and anchor.strip():
                            disambiguated = any(term in anchor.lower() for term in
                                                 ("kf", "keyframe", "deformed", "rest-pose", "rest pose"))
                            if not disambiguated:
                                problems.append(
                                    f"joint_feedback[{i}]['anchored_to'] = {anchor!r} does not specify which "
                                    f"geometry block (per-keyframe deformed vs rest-pose) the coordinate came "
                                    f"from — ambiguous, both blocks reuse the same stroke numbering"
                                )
    return problems


def load_keyframe_images(frames_dir, n_keyframes=5):
    from PIL import Image
    images = []
    for kf in range(n_keyframes):
        path = os.path.join(frames_dir, f"kf{kf}.png")
        if not os.path.exists(path):
            raise SystemExit(f"Missing expected keyframe image: {path}")
        images.append(Image.open(path).convert("RGB"))
    return images


def run_judge(model, processor, images, prompt):
    import torch
    content = [{"type": "image", "image": img} for img in images]
    content.append({"type": "text", "text": prompt})
    messages = [{"role": "user", "content": content}]

    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[text], images=images, return_tensors="pt").to(model.device)

    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=1200, temperature=0.2, do_sample=True)

    generated = output_ids[:, inputs["input_ids"].shape[1]:]
    response = processor.batch_decode(generated, skip_special_tokens=True)[0]
    return response


def load_vlm(model_path, quantize=None):
    """
    quantize: "8bit", "4bit", or None (full precision) — currently
    defaulting to None (full precision) per explicit choice to accept
    the risk of the CUDA OOM this GPU has repeatedly hit before, over
    the ~8-bit slowdown (dequantize overhead per layer means 8bit is
    typically SLOWER per-token than full precision, not faster — it
    only exists as a memory-fit tradeoff, never was a speed win).

    CONFIRMED on real hardware: full-precision loading consistently uses
    ~13.8 GiB on a GPU with ~14.6 GiB actually usable (15.59 GiB nominal
    minus this machine's desktop/display overhead), leaving under 100 MiB
    free — repeatedly triggered CUDA OOM across multiple different
    sketches (dog9, ice2). If this recurs, pass quantize="8bit" or "4bit".

    Model class is auto-selected from model_path, since Qwen3-VL and
    Qwen2.5-VL are NOT interchangeable — different class, not just a
    different string. Qwen2.5-VL requires transformers>=4.49.0 (confirmed
    fine on this environment's transformers==5.14.1); it was NOT viable
    back when this environment was still on Python 3.8.
    """
    if "qwen2.5-vl" in model_path.lower() or "qwen2_5_vl" in model_path.lower():
        from transformers import Qwen2_5_VLForConditionalGeneration as ModelClass, AutoProcessor
    elif "qwen2-vl" in model_path.lower():
        from transformers import Qwen2VLForConditionalGeneration as ModelClass, AutoProcessor
    else:
        from transformers import Qwen3VLForConditionalGeneration as ModelClass, AutoProcessor
    print(f"loading VLM from {model_path} (class={ModelClass.__name__}, quantize={quantize}) ...")

    kwargs = {"device_map": "auto"}
    if quantize == "8bit":
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
    elif quantize == "4bit":
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_compute_dtype="float16", bnb_4bit_quant_type="nf4")
    else:
        kwargs["dtype"] = "auto"

    model = ModelClass.from_pretrained(model_path, **kwargs)
    processor = AutoProcessor.from_pretrained(model_path)
    print("VLM loaded.")
    return model, processor


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("sketch", type=str,
                     help="sketch name — used to auto-derive the frames folder and look up the caption")
    ap.add_argument("--frames-dir", type=str, default=None,
                     help="directory containing kf0.png ... kf4.png (default: frames/{sketch}, "
                          "matching pipeline.py's default output layout)")
    ap.add_argument("--model", type=str, default="Qwen/Qwen3-VL-4B-Instruct")
    ap.add_argument("--caption", type=str, default=None,
                     help="override the caption directly instead of looking it up")
    ap.add_argument("--caption-file", type=str, default=CAPTION_FILE_DEFAULT)
    ap.add_argument("--no-caption", action="store_true",
                     help="skip caption lookup entirely (faithfulness will be judged as N/A)")
    ap.add_argument("--out", type=str, default=None,
                     help="where to write the verdict JSON (default: json/{sketch}/judge_verdict.json)")
    args = ap.parse_args()

    frames_dir = args.frames_dir or os.path.join("frames", args.sketch)

    caption = None
    if args.no_caption:
        print("--no-caption given, skipping caption lookup")
    elif args.caption:
        caption = args.caption
    else:
        from lib import get_caption
        caption = get_caption(args.caption_file, args.sketch)
        print(f"looked up caption for '{args.sketch}': {caption!r}")

    images = load_keyframe_images(frames_dir)
    print(f"loaded {len(images)} keyframe images from {frames_dir}")

    prompt = build_judge_prompt(args.sketch, caption)
    print("\n--- JUDGE PROMPT ---")
    print(prompt)

    model, processor = load_vlm(args.model)
    response = run_judge(model, processor, images, prompt)
    print("\n--- RAW RESPONSE ---")
    print(response)

    parsed = parse_judge_response(response)
    problems = validate_judge_response(parsed)
    print("\n--- PARSED VERDICT ---")
    print(json.dumps(parsed, indent=2))
    if problems:
        print("--- VALIDATION PROBLEMS ---")
        for p in problems:
            print(f"  - {p}")

    out_path = args.out
    if not out_path:
        out_path = os.path.join(os.path.dirname(os.path.dirname(frames_dir.rstrip("/"))),
                                 "json", args.sketch, "judge_verdict.json")
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    with open(out_path, "w") as f:
        json.dump(parsed, f, indent=2)
    print(f"\nwrote {out_path}")


if __name__ == "__main__":
    main()