svg.py

"""
mosketch_pipeline.py — the full pipeline as one script with subcommands.

Pipeline:
    1. Identify objects        -> from the semantic file (no Qwen)
    2. Classify ARAP vs. not   -> `classify` subcommand (one Qwen call, all objects)
    3. Narrate + deform        -> `narrate` + `deform` subcommands, ONLY for
                                   objects marked ARAP in step 2
    4. Render                  -> `render` subcommand; every object gets real
                                   trajectory translation; ARAP objects
                                   additionally get deformation on top

Run steps individually for debugging, or use `full` to run everything for
one sketch in one process — this loads the Qwen model ONCE and reuses it
across classify/narrate/deform, instead of loading it 3 separate times.

Examples:
    # step by step
    python mosketch_pipeline.py classify --model M --caption-file C --sketch-name S --semantic SEM --out deformation.json
    python mosketch_pipeline.py narrate --model M --caption-file C --sketch-name S --objects dog --out narratives.json
    python mosketch_pipeline.py deform --model M --svg S.svg --semantic SEM --traj T --narratives narratives.json --deformation deformation.json --out-dir .
    python mosketch_pipeline.py render --svg S.svg --semantic SEM --traj T --handles-dir .

    # everything at once, one model load
    python mosketch_pipeline.py full --model M --caption-file C --sketch-name S --svg S.svg --semantic SEM --traj T --out-dir .
"""

import argparse
import json
import os
import re
import sys

import numpy as np

from lib import (
    load_strokes_from_svg, load_semantic_assignments, filter_strokes, flatten_strokes,
    load_object, deduplicate_points, build_mesh, nearest_mesh_vertex,
    auto_select_handles_deduped, object_bbox_size, arap_deform,
    load_trajectories, bbox_deltas, get_caption, build_stroke_geometry_text,
)

N_KEYFRAMES = 5
MAX_RETRIES = 5
LOW_UTILIZATION_THRESHOLD = 0.30  # mean fraction of the movement cap actually used —
                                   # confirmed on real hardware (horsecar5) that Qwen can
                                   # stay well within a generous cap without ever being told
PLAUSIBILITY_THRESHOLD = 4
FAITHFULNESS_THRESHOLD = 4  # both must pass to stop — faithfulness previously only
                            # affected feedback text, never actually gated success
QUALITY_THRESHOLD = 4       # same upgrade applied to the new quality criterion —
                            # scored but not gating would repeat the same mistake


def faithfulness_passed(score):
    """
    faithfulness_score can be a number 1-5, the string "N/A" (no caption
    was given to compare against, so there's nothing to fail), or missing
    entirely (treated as NOT passed — can't confirm it's good, so err
    toward regenerating the narrative rather than assuming it's fine).
    """
    if score is None:
        return False
    if isinstance(score, str):
        return score.strip().upper() == "N/A"
    if isinstance(score, (int, float)):
        return score >= FAITHFULNESS_THRESHOLD
    return False


def unload_model(model):
    """Frees GPU memory before loading a different model. Necessary because
    the narrate/deform steps use a text Qwen model and the judge step uses
    a separate vision-language model (Qwen3-VL) — on hardware with limited
    VRAM (this project's RTX A4000, 16GB, already documented as a tight
    fit for a single model), loading both at once risks the same OOM issue
    that blocked Wan2.2 integration earlier. Load/unload sequentially
    instead of assuming both fit simultaneously.

    CONFIRMED BUG (found on real hardware, invisible to all mocked testing
    since no real GPU was available to catch it): `del model` here only
    clears THIS function's own local reference — it does nothing to the
    caller's variable, which stays alive and keeps the whole model
    resident in VRAM. torch.cuda.empty_cache() then has nothing to
    actually free, because the refcount never reaches zero. Fixed by
    returning None — callers MUST reassign their variable to this return
    value (e.g. `model = unload_model(model)`), or the bug reappears."""
    import gc
    del model
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except ImportError:
        pass
    return None

# Dog3's real, human-reviewed narratives — used BOTH as the few-shot example
# in `narrate` and as the fallback default if --narratives is omitted in
# `deform`. One constant, one source of truth (previously duplicated across
# two separate files under two different names with identical content).
DOG3_CAPTION = ("The person throws a frisbee through the air, and the dog sits poised, "
                "ready to sprint forward and catch it with its mouth in a swift motion.")
DOG3_NARRATIVES = {
    "dog": [
        "the dog is sitting alert, watching the frisbee as it is thrown",
        "the dog is beginning to rise, weight shifting forward, head reaching toward the frisbee",
        "the dog is mid-leap, body extended, reaching far forward and up toward the frisbee",
        "the dog is at the peak of its jump, reaching as far as possible toward the frisbee",
        "the dog is landing after catching the frisbee, body compacting back down",
    ],
    "frisbee": [
        "the frisbee has just left the thrower's hand, angled slightly upward",
        "the frisbee is gliding through the air, tilting slightly as it arcs",
        "the frisbee is near the peak of its arc, angled toward the dog",
        "the frisbee is descending toward the dog, tilting down slightly",
        "the frisbee is at the dog's mouth, being caught",
    ],
}

SVG_PATH_DEFAULT = "/mnt/user-data/uploads/dog3.svg"
SEMANTIC_PATH_DEFAULT = "dog3_semantic.txt"
DEFAULT_COLOR = "#444444"
DEFAULT_LINEWIDTH = 1.1
OBJECT_COLORS = {"dog": "black", "person": "#3F4C57", "frisbee": "#B0463C"}
OBJECT_LINEWIDTH = {"dog": 1.1, "person": 1.1, "frisbee": 1.4}


# =============================================================================
# shared: Qwen call + response parsing
# =============================================================================

def query_qwen(model, tokenizer, prompt, device, max_new_tokens=500, temperature=0.1):
    import torch
    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer([text], return_tensors="pt").to(device)
    input_token_count = inputs["input_ids"].shape[1]
    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens,
                                     temperature=temperature, do_sample=True)
    generated = output_ids[0][inputs["input_ids"].shape[1]:]
    output_token_count = generated.shape[0]
    response_text = tokenizer.decode(generated, skip_special_tokens=True)
    return response_text, input_token_count, output_token_count


def load_qwen_model(model_path):
    from transformers import AutoModelForCausalLM, AutoTokenizer
    import torch
    print(f"loading model from {model_path} ...")
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, device_map="auto")
    device = next(model.parameters()).device
    print("model loaded.")
    return model, tokenizer, device


def parse_json_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))


# =============================================================================
# STEP 2: classify — ARAP vs TRAJ_ONLY, one call, all objects
# =============================================================================

def build_deformation_prompt(caption, object_names):
    objects_str = ", ".join(f'"{o}"' for o in object_names)
    lines = "\n".join(
        f'  "{o}": "ARAP" or "TRAJ_ONLY"' + ("," if i < len(object_names) - 1 else "")
        for i, o in enumerate(object_names)
    )
    return f"""Scene: "{caption}"

Objects in this scene: {objects_str}

For each object, decide whether representing it correctly needs NON-RIGID DEFORMATION (its body/shape changes — e.g. limbs moving, a neck reaching, a body crouching or leaning) or whether simple RIGID TRANSLATION (the object moves/rotates as a whole, unchanged in shape, or doesn't move at all) is enough.

Answer "ARAP" if the object's shape or body configuration changes at any point in the action, even if its overall position doesn't change. Answer "TRAJ_ONLY" if the object is rigid (a vehicle, tool, projectile, furniture, background element) or is simply carried by its own movement without changing shape.

Respond with ONLY a JSON object, no other text, in this exact format:
{{
{lines}
}}
"""


def validate_deformation(parsed, object_names):
    problems = []
    for obj in object_names:
        if obj not in parsed:
            problems.append(f"'{obj}' missing from response")
            continue
        val = str(parsed[obj]).strip().upper()
        if val not in ("ARAP", "TRAJ_ONLY"):
            problems.append(f"'{obj}' has invalid value {parsed[obj]!r}, expected ARAP or TRAJ_ONLY")
    return problems


def run_classify(model, tokenizer, device, caption, semantic_path, out_path):
    assignments = load_semantic_assignments(semantic_path)
    object_names = list(assignments.keys())
    print(f"objects found in {semantic_path}: {object_names}")

    prompt = build_deformation_prompt(caption, object_names)
    print("\n--- CLASSIFY PROMPT ---")
    print(prompt)

    response, in_tok, out_tok = query_qwen(model, tokenizer, prompt, device,
                                            max_new_tokens=250, temperature=0.1)
    print(f"\ntokens: {in_tok} in / {out_tok} out")
    print("--- RAW RESPONSE ---")
    print(response)

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

    arap_objs = [o for o in object_names if str(parsed.get(o, "")).strip().upper() == "ARAP"]
    traj_only_objs = [o for o in object_names if o not in arap_objs]
    print(f"\nARAP: {arap_objs}")
    print(f"TRAJ_ONLY: {traj_only_objs}")

    with open(out_path, "w") as f:
        json.dump(parsed, f, indent=2)
    print(f"wrote {out_path}")
    return parsed, arap_objs


def build_objects_info(svg_path, semantic_path, arap_objects):
    """
    Standalone version of the mesh/joint setup previously embedded inside
    run_deform — factored out so the unified narrate+deform+judge retry
    loop can build this ONCE before the loop (mesh/joints never change
    between attempts) instead of recomputing it every attempt.
    """
    objects_info = {}
    for obj_name in arap_objects:
        points, slices = load_object(obj_name, svg_path, semantic_path)
        bbox_size = object_bbox_size(points)
        unique_points, p2u = deduplicate_points(points, tol=0.35)
        tri, edges = build_mesh(unique_points)
        strokes = filter_strokes(load_strokes_from_svg(svg_path),
                                  load_semantic_assignments(semantic_path)[obj_name])
        joints, anchor_idx, handle_idxs, joint_mesh_indices = auto_select_handles_deduped(
            strokes, unique_points, k=4)

        if len(handle_idxs) == 0:
            print(f"WARNING: '{obj_name}' has no independent handles after dedup, skipping")
            continue

        print(f"'{obj_name}': {len(joints)} joints, anchor={anchor_idx}, handles={handle_idxs}, "
              f"mesh_indices={joint_mesh_indices}, bbox_size={bbox_size:.1f}")

        objects_info[obj_name] = {
            "joints": joints, "anchor_idx": anchor_idx, "handle_idxs": handle_idxs,
            "joint_mesh_indices": joint_mesh_indices, "bbox_size": bbox_size, "strokes": strokes,
            "points": points, "slices": slices, "unique_points": unique_points,
            "p2u": p2u, "edges": edges,
        }
    return objects_info


def build_deformed_stroke_geometry_text(object_info, kf_targets, n_points=2):
    """
    Reconstructs what a previous attempt's ACTUAL DEFORMED shape looked
    like, as compact text — used to replace the "attached image of the
    previous attempt" memory in svg.py, which has no images at all.

    kf_targets: {joint_i_str: [x, y]} for ONE keyframe, in the same
    "joint_i" key format apply_deform_clip_and_write saves (i.e. one
    entry of deform_outputs[obj_name][kf_key]).

    Runs the SAME ARAP solve used for real rendering, then formats the
    resulting deformed stroke points the same way build_stroke_geometry_text
    formats the rest pose — so a retry sees the previous attempt's actual
    resulting SHAPE in text, not just the isolated handle-joint numbers.
    """
    import re as _re
    unique_points = object_info["unique_points"]
    edges = object_info["edges"]
    p2u = object_info["p2u"]
    points = object_info["points"]
    slices = object_info["slices"]
    anchor_idx = object_info["anchor_idx"]
    joints = object_info["joints"]
    joint_mesh_indices = object_info["joint_mesh_indices"]

    handle_mesh_indices, handle_targets = [], []
    for key, target in kf_targets.items():
        m = _re.match(r"joint_(\d+)", key)
        if not m:
            continue
        joint_i = int(m.group(1))
        if joint_i >= len(joint_mesh_indices):
            continue
        handle_mesh_indices.append(joint_mesh_indices[joint_i])
        handle_targets.append(target)
    # anchor is always held at rest, same convention as rendering
    handle_mesh_indices.append(joint_mesh_indices[anchor_idx])
    handle_targets.append(joints[anchor_idx].tolist())

    if not handle_mesh_indices:
        return None

    deformed_unique = arap_deform(unique_points, edges, handle_mesh_indices,
                                   np.array(handle_targets), iterations=10)
    deformed_points = deformed_unique[p2u]

    lines = []
    for i, (start, end) in enumerate(slices):
        stroke_pts = deformed_points[start:end]
        idxs = np.linspace(0, len(stroke_pts) - 1, n_points).astype(int)
        pts = stroke_pts[idxs]
        pts_str = " -> ".join(f"({x:.0f},{y:.0f})" for x, y in pts)
        lines.append(f"  stroke_{i}: {pts_str}")
    return "\n".join(lines)


def render_rest_pose_multi(object_names, svg_path, semantic_path, out_path):
    """
    Renders multiple objects' ORIGINAL strokes together, at their real
    positions in the source SVG (no deformation, no trajectory translation)
    — this is the "what does the sketch actually look like" image every
    attempt is grounded in, so Qwen can see what's actually drawable
    (e.g. whether the dog's back legs even exist as strokes) instead of
    only reasoning from the caption's text description.
    """
    import matplotlib.pyplot as plt
    fig, ax = plt.subplots(figsize=(6, 6))
    for name in object_names:
        points, slices = load_object(name, svg_path, semantic_path)
        for start, end in slices:
            seg = points[start:end]
            ax.plot(seg[:, 0], seg[:, 1], color="black", linewidth=1.2)
    ax.invert_yaxis()
    ax.set_aspect("equal")
    ax.set_title("rest pose")
    fig.savefig(out_path, dpi=150, bbox_inches="tight")
    plt.close(fig)
    return out_path


def build_combined_object_section(object_name, joints, anchor_idx, handle_idxs, bbox_size,
                                   previous_narrative=None, feedback=None, n_keyframes=N_KEYFRAMES,
                                   freeze_narrative=False, strokes=None, previous_shape_text=None):
    """
    freeze_narrative: if True, previous_narrative is used as a FIXED target
    pose description (the object's narrative already passed faithfulness —
    only the numeric deformation needs to improve, not the story). If
    False (default), the narrative is regenerated fresh, informed by
    feedback (and, in the image-based pipeline variants, the attached
    image(s) — this variant, svg.py, has NO images at all, relying
    entirely on text: stroke geometry + previous numeric targets).

    strokes: if given, the object's actual stroke points are included as
    TEXT (not just the rendered image) — added specifically because
    multimodal LLMs can under-attend to image content relative to text;
    this gives the same geometric information in a text-native form the
    model is more likely to actually use. Kept deliberately sparse (2
    points per stroke, start+end only) since a complex object can have
    60+ strokes — measured on real dog3 data: 2 points/stroke costs
    ~570 tokens for a 64-stroke object vs ~2650 for the full 12
    points/stroke used internally for the ARAP mesh.

    previous_shape_text: pre-formatted text (from build_deformed_stroke_geometry_text,
    computed one level up where the full mesh/ARAP data is available) showing
    the previous attempt's ACTUAL DEFORMED STROKE positions at 2 representative
    keyframes — not raw joint numbers, the reconstructed resulting SHAPE, in the
    same stroke-point text format as the rest pose. This is svg.py's replacement
    for showing Qwen the previous attempt's rendered images — since there are no
    images here at all, without this the model has zero way to know what its own
    previous output actually looked like on a retry.
    """
    cap = round(bbox_size * 0.25, 1)
    joint_lines = "\n".join(
        f'    - joint_{i}: rest position (x={joints[i][0]:.1f}, y={joints[i][1]:.1f})'
        + ("  <-- ANCHOR, must stay at or near this position in EVERY keyframe" if i == anchor_idx else "")
        for i in range(len(joints))
    )
    handle_list_str = ', joint_'.join(str(i) for i in handle_idxs)

    geometry_block = ""
    if strokes:
        geometry_text = build_stroke_geometry_text(strokes, n_points=2)
        geometry_block = (
            f"\n  This object's ACTUAL drawn strokes (start -> end point of each stroke, same coordinate "
            f"space as the joints above) — use this to know exactly what is and isn't actually drawn, don't "
            f"invent motion for parts that have no strokes here:\n{geometry_text}\n"
        )

    previous_targets_block = ""
    if previous_shape_text:
        previous_targets_block = (
            f"\n  Your previous attempt's ACTUAL DEFORMED SHAPE at two representative keyframes "
            f"(reconstructed from your own previous target positions — same stroke-point format as "
            f"the rest pose above, so you can directly compare what changed):\n{previous_shape_text}\n"
        )

    if freeze_narrative and previous_narrative:
        pose_lines = "\n".join(f"    kf{i}: {desc}" for i, desc in enumerate(previous_narrative))
        feedback_line = f'\n  This pose story already matches the intended action — it is FIXED, do not change it. ' \
                         f'Only the numeric target positions need to improve.' \
                         + (f' Previous attempt was judged: "{feedback}"' if feedback else "") + \
                         "\n  Adjust specifically from your previous target positions shown above where the " \
                         "critique applies — this is a revision of known values, not a fresh guess."
        return f"""Object: "{object_name}"
  Joints:
{joint_lines}
{geometry_block}
  Target pose across all {n_keyframes} keyframes (FIXED, already correct — do not rewrite):
{pose_lines}
{feedback_line}
{previous_targets_block}
  For non-anchor joints (joint_{handle_list_str}), do not move more than {cap} pixels from REST in any keyframe. IMPORTANT: these {n_keyframes} keyframes are SPARSE anchor points spanning the ENTIRE action, NOT consecutive video frames — a large, dramatic difference between consecutive keyframes is NORMAL and EXPECTED, not an error; the actual in-between motion will be generated separately later by a different model. Positions should progress in a DIRECTIONALLY COHERENT way (don't make real progress toward the action and then have a LATER keyframe randomly revert backward without the narrative describing a reason to — e.g. only "landing"/"settling" should move back toward rest). Small, timid, barely-different positions between keyframes are themselves a mistake, not a safe choice."""

    previous_block = ""
    if previous_narrative or feedback:
        parts = []
        if previous_narrative:
            parts.append(f"Your previous narrative attempt was:\n{json.dumps(previous_narrative, indent=2)}")
        if feedback:
            parts.append(f'That attempt was judged and received this critique: "{feedback}"')
        parts.append("Revise BOTH the narrative and the target positions to address the critique — don't just "
                      "reword the narrative superficially while leaving the same underlying numeric problem.")
        previous_block = "\n  " + "\n  ".join(parts) + "\n"
        previous_block += previous_targets_block

    return f"""Object: "{object_name}"
  Joints:
{joint_lines}
{geometry_block}
{previous_block}
  For non-anchor joints (joint_{handle_list_str}), do not move more than {cap} pixels from REST in any keyframe. IMPORTANT: these {n_keyframes} keyframes are SPARSE anchor points spanning the ENTIRE action, NOT consecutive video frames — a large, dramatic difference between consecutive keyframes is NORMAL and EXPECTED, not an error; the actual in-between motion will be generated separately later by a different model. Positions should progress in a DIRECTIONALLY COHERENT way (don't make real progress toward the action and then have a LATER keyframe randomly revert backward without the narrative describing a reason to — e.g. only "landing"/"settling" should move back toward rest). Small, timid, barely-different positions between keyframes are themselves a mistake, not a safe choice."""


DOG3_COMBINED_FEWSHOT_EXAMPLE = {
    "dog": {
        "narrative": [
            "the dog is sitting alert, watching the frisbee as it is thrown",
            "the dog is beginning to rise, weight shifting forward, head reaching toward the frisbee",
            "the dog is mid-leap, body extended, reaching far forward and up toward the frisbee",
            "the dog is at the peak of its jump, reaching as far as possible toward the frisbee",
            "the dog is landing after catching the frisbee, body compacting back down",
        ],
        # real dog3 joint rest positions: joint_1=(178.28,136.85) head, joint_2=(228.78,194.94)
        # tail, joint_3=(189.29,151.35) neck — every value below verified to stay within a
        # 27px cap of rest. Notice the progression BUILDS UP through kf0->kf3 (increasing
        # displacement, matching "rising -> leaping -> peak reach") and only SETTLES BACK at
        # kf4 ("landing") — this is the exact monotonic-then-settle shape that was missing
        # when a real run produced a kf2 spike with kf3/kf4 reverting toward rest with no
        # narrative reason to.
        "targets": {
            "kf0": {"joint_1": [178.3, 136.8], "joint_2": [228.8, 194.9], "joint_3": [189.3, 151.3]},
            "kf1": {"joint_1": [168.0, 127.0], "joint_2": [232.0, 191.0], "joint_3": [184.0, 144.0]},
            "kf2": {"joint_1": [160.0, 120.0], "joint_2": [237.0, 186.0], "joint_3": [177.0, 137.0]},
            "kf3": {"joint_1": [159.0, 119.0], "joint_2": [240.0, 183.0], "joint_3": [174.0, 134.0]},
            "kf4": {"joint_1": [168.0, 128.0], "joint_2": [231.0, 192.0], "joint_3": [185.0, 146.0]},
        },
    }
}


def build_combined_narrate_deform_prompt(objects_info, caption, previous_narratives=None, feedback=None,
                                          n_keyframes=N_KEYFRAMES, is_retry=False, few_shot=True,
                                          freeze_narrative=False, previous_deform_outputs=None):
    # keyframes shown for the "previous attempt's actual shape" memory —
    # deliberately 2, not all 5: measured real cost on dog3's 64-stroke dog
    # object at 5/5 keyframes was ~2,866 tokens for ONE object alone, which
    # multiplies fast for multi-object scenes. kf2 (mid-action) + kf4 (final)
    # gives a genuine "what happened" signal at 60% less cost.
    SHAPE_MEMORY_KEYFRAMES = [2, 4]

    sections, example_parts = [], []
    for name, info in objects_info.items():
        prev_narrative_for_obj = (previous_narratives or {}).get(name)

        previous_shape_text = None
        prev_deform_for_obj = (previous_deform_outputs or {}).get(name)
        if prev_deform_for_obj:
            shape_parts = []
            for kf in SHAPE_MEMORY_KEYFRAMES:
                kf_key = f"kf{kf}"
                if kf_key not in prev_deform_for_obj:
                    continue
                shape_text = build_deformed_stroke_geometry_text(info, prev_deform_for_obj[kf_key], n_points=2)
                if shape_text:
                    shape_parts.append(f"  -- {kf_key} --\n{shape_text}")
            if shape_parts:
                previous_shape_text = "\n".join(shape_parts)

        sections.append(build_combined_object_section(
            name, info["joints"], info["anchor_idx"], info["handle_idxs"], info["bbox_size"],
            previous_narrative=prev_narrative_for_obj, feedback=feedback, n_keyframes=n_keyframes,
            freeze_narrative=freeze_narrative, strokes=info.get("strokes"),
            previous_shape_text=previous_shape_text,
        ))
        kf_examples = ",\n".join(
            "        \"kf%d\": {%s}" % (kf, ", ".join(f'"joint_{i}": [x, y]' for i in info["handle_idxs"]))
            for kf in range(n_keyframes)
        )
        if freeze_narrative:
            example_parts.append(
                f'  "{name}": {{\n'
                f'    "targets": {{\n{kf_examples}\n    }}\n'
                f'  }}'
            )
        else:
            example_parts.append(
                f'  "{name}": {{\n'
                f'    "narrative": [<{n_keyframes} short pose description strings, one per keyframe>],\n'
                f'    "targets": {{\n{kf_examples}\n    }}\n'
                f'  }}'
            )

    all_sections = "\n\n".join(sections)
    example_json = "{\n" + ",\n".join(example_parts) + "\n}"

    image_context = ""
    if is_retry:
        image_context = ("This is a RETRY. There are no images — everything you need is given as TEXT below: "
                          "each object's rest-pose stroke geometry, and (where available) your previous "
                          "attempt's actual deformed shape at two representative keyframes.")
    else:
        image_context = ("There are no images in this prompt — the object's rest-pose stroke geometry is "
                          "given below as TEXT. Use it to understand exactly what strokes exist and are "
                          "available to move; do not invent motion for body parts that aren't actually drawn.")

    fewshot_block = ""
    if few_shot:
        fewshot_json = json.dumps(DOG3_COMBINED_FEWSHOT_EXAMPLE, indent=2)
        fewshot_block = f"""Example — for the scene "{DOG3_CAPTION}", a good answer looks like:
{fewshot_json}

Notice: each narrative keyframe reads as a distinct, substantially different stage of the action — not a near-duplicate of its neighbor, and not a small incremental change from it. The joint targets BUILD UP smoothly (kf0 -> kf1 -> kf2 -> kf3 each moving further than the last) and only settle back toward rest at the FINAL keyframe, matching the narrative's "landing" moment — no keyframe overshoots and then has a later keyframe revert back toward rest without a narrative reason to. Match this style and this kind of numeric consistency for the new scene below.

"""

    if freeze_narrative:
        output_instruction = (
            'For EACH object above, the narrative/pose story is already fixed (shown above) — '
            'produce ONLY:\n'
            '  "targets": target (x, y) positions for its non-anchor joints, at every keyframe, '
            'consistent with the fixed pose story above.'
        )
    else:
        output_instruction = (
            "For EACH object above, produce BOTH:\n"
            '  1. "narrative": a plain-English pose description for each keyframe. REMEMBER: these are '
            "SPARSE keyframes spanning the WHOLE action, not consecutive video frames — each description "
            "should be a meaningfully, substantially different stage of the action from its neighbors, not "
            "a small incremental change. Write these like 5 distinct captions for 5 different moments spread "
            "across an entire action, not like 5 near-duplicate snapshots a split-second apart. Under 20 "
            "words each.\n"
            '  2. "targets": target (x, y) positions for its non-anchor joints, at every keyframe, '
            "consistent with your own narrative."
        )

    return f"""{fewshot_block}You are directing a {n_keyframes}-keyframe animated sequence for a hand-drawn sketch, viewed from the side. Coordinate system: x increases rightward, y increases DOWNWARD.

IMPORTANT: these {n_keyframes} keyframes are SPARSE anchor points sampled across the ENTIRE action from start to finish — NOT consecutive video frames. Think of them like 5 widely-spaced snapshots of a whole motion, not neighboring frames a fraction of a second apart. Large, dramatic pose changes between consecutive keyframes are normal and expected; a separate model will generate the actual in-between motion frames later. Do not treat these like near-continuous animation frames.

Scene: "{caption}"

{image_context}

{all_sections}

{output_instruction}

Consider objects together (e.g. a dog reaching toward a frisbee should be spatially consistent with the frisbee's own position) and consider each object's OWN sequence together — these {n_keyframes} keyframes are SPARSE anchor points spanning the WHOLE action, not consecutive video frames, so large differences between consecutive keyframes are expected and correct, not something to avoid. The many actual in-between motion frames will be generated separately later. Only avoid a keyframe making real progress and then a LATER keyframe randomly reverting backward without the narrative describing why.

Respond with ONLY one JSON object, no other text, in this exact format:
{example_json}
"""


def run_combined_narrate_deform(model, processor, images, prompt):
    """
    Same multi-image calling pattern as vlm_judge.run_judge, but svg.py
    calls this with images=[] (or None) on every call — no images at
    all, text-only. UNCONFIRMED against a real model: whether
    processor(images=[]) behaves identically to omitting the images kwarg
    entirely. Handled defensively here rather than assumed, since this
    can't be tested without real hardware.
    """
    import torch
    has_images = bool(images)
    content = [{"type": "image", "image": img} for img in images] if has_images else []
    content.append({"type": "text", "text": prompt})
    messages = [{"role": "user", "content": content}]

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

    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=500 * N_KEYFRAMES, temperature=0.1, do_sample=True)

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


def build_svg_judge_prompt(sketch_name, all_keyframes_geometry_text, caption=None,
                            dino_stagnation=None, dino_temporal=None, clip_scores=None,
                            cap_utilization=None, n_keyframes=N_KEYFRAMES):
    """
    Text-only judge prompt — reconstructs each keyframe's ACTUAL deformed
    stroke geometry (same function used for retry memory) and gives ALL
    5 keyframes to the judge as text, instead of showing it the rendered
    images. Kept separate from vlm_judge.build_judge_prompt (which stays
    image-based) so pipe.py/pipeline1.py are never affected by this.

    HONEST CAVEAT, not glossed over: the QUALITY criterion below ("does
    this look like a clean line drawing vs garbled/noisy") is inherently
    a pixel-level, visual question. Assessing it from coordinate text
    alone is a fundamentally harder, more indirect task than looking at
    the actual rendering — this is the part of the experiment most
    likely to perform worse than the image-based judge, not a solved
    problem.
    """
    caption_block = ""
    if caption:
        caption_block = f'\nThe sketch is supposed to depict: "{caption}"\n'

    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's RENDERED images, per keyframe "
                              f"(1.0 = identical): {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' RENDERED images: {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: {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 (from the RENDERED images): "
                              f"{clip_str} (mean {clip_scores['mean_clip_score']:.3f})")
    metrics_block = ""
    if metrics_lines:
        metrics_block = ("\nObjective measurements computed from the actual rendered images (these ARE "
                          "image-derived even though you are not shown the images directly — use them to "
                          "inform your reasoning):\n" + "\n".join(metrics_lines) + "\n")

    return f"""You are judging a sequence of {n_keyframes} keyframes generated for an animated sketch named "{sketch_name}". You are NOT shown images — instead, each keyframe's actual deformed stroke geometry is given below as TEXT (start -> end point of each stroke, in pixel coordinates, y increases DOWNWARD).

IMPORTANT: these are SPARSE keyframes sampled across the ENTIRE action from start to finish — NOT consecutive video frames. Large, dramatic differences between consecutive keyframes are normal and correct.
{caption_block}{metrics_block}
{all_keyframes_geometry_text}

Evaluate the sequence on these three criteria, reasoning from the coordinate data above:

1. PLAUSIBILITY (1-5): Does the sequence of stroke positions, taken as sparse waypoints across the whole action, describe a coherent and physically believable progression? 5 = each keyframe is a sensible, meaningfully different stage of the action, in a believable order. 1 = static (coordinates barely change between keyframes) or the positions describe something physically impossible/nonsensical.

2. FAITHFULNESS (1-5, or "N/A" if no caption was given above): Does the described motion match what the caption says should be happening? 5 = clearly matches, 1 = unrelated to the description.

3. QUALITY (1-5): Judging ONLY from the coordinate data (you cannot see the actual rendered image) — do the stroke positions look geometrically coherent, e.g. no wildly self-intersecting or degenerate configurations that would likely render as visual noise? 5 = coordinates look clean and coherent, 1 = coordinates suggest severe geometric distortion. NOTE: this criterion is fundamentally harder to assess without seeing the actual rendering — be appropriately uncertain rather than overconfident.

Respond with ONLY a JSON object, no other text, in this exact format:
{{
  "plausibility_score": <1-5>,
  "plausibility_notes": "<one sentence>",
  "faithfulness_score": <1-5 or "N/A">,
  "faithfulness_notes": "<one sentence>",
  "quality_score": <1-5>,
  "quality_notes": "<one sentence>",
  "overall_verdict": "<one or two sentence summary>"
}}
"""


def run_svg_judge(model, processor, prompt):
    """Text-only judge call — same no-images pattern as run_combined_narrate_deform."""
    import torch
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = processor(text=[text], return_tensors="pt").to(model.device)

    with torch.no_grad():
        output_ids = model.generate(**inputs, max_new_tokens=400, 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 build_all_keyframes_geometry_text(objects_info, deform_outputs, n_keyframes=N_KEYFRAMES):
    """
    For EVERY object and EVERY keyframe of the CURRENT attempt, reconstruct
    the actual deformed stroke geometry as text — unlike the 2-keyframe
    retry memory, judging needs the FULL sequence to assess the whole
    story, not just a sample. Real, measured cost: ~2,866 tokens for ONE
    64-stroke object across all 5 keyframes — this multiplies per object,
    so multi-object scenes get expensive fast. No trimming applied here;
    if this needs to be cut down for a specific scene, that's the next
    thing to adjust.
    """
    sections = []
    for obj_name, info in objects_info.items():
        obj_targets = deform_outputs.get(obj_name)
        if not obj_targets:
            continue
        kf_parts = []
        for kf in range(n_keyframes):
            kf_key = f"kf{kf}"
            if kf_key not in obj_targets:
                continue
            shape_text = build_deformed_stroke_geometry_text(info, obj_targets[kf_key], n_points=2)
            if shape_text:
                kf_parts.append(f"  -- {kf_key} --\n{shape_text}")
        if kf_parts:
            sections.append(f'Object "{obj_name}":\n' + "\n".join(kf_parts))
    return "\n\n".join(sections)


def apply_deform_clip_and_write(parsed, objects_info, out_dir, sketch_name, n_keyframes=N_KEYFRAMES,
                                 frozen_narratives=None):
    """
    Shared post-processing for the combined call's "targets" section:
    same hard-clip logic as the old run_deform, applied here instead.
    Returns (narratives_dict, deform_outputs_dict, cap_utilization_dict).

    frozen_narratives: {obj_name: [...]} — used as a fallback when the
    response doesn't include a "narrative" key for an object, which
    happens when freeze_narrative=True was used in the prompt (the model
    was never asked to produce one, so its absence is expected, not an
    error — carry the frozen one forward instead of losing it).

    cap_utilization: {obj_name: mean_fraction_of_cap_used} — CONFIRMED on
    real hardware (horsecar5's person) that Qwen can propose displacement
    well within the movement cap without ever being told it did so — the
    handle selection and clipping were both working correctly, but the
    actual output was too timid to be visible (e.g. a leg moving only 18%
    of its allowed range). The hard clip only ever catches OVER the cap;
    nothing previously caught UNDER-using it. This surfaces that as an
    explicit number so it can be fed back to Qwen directly.
    """
    narratives_out = {}
    deform_outputs = {}
    utilization_by_obj = {}  # {obj_name: [fraction, fraction, ...]} across all joints/keyframes

    for obj_name, info in objects_info.items():
        if obj_name not in parsed:
            print(f"  WARNING: '{obj_name}' missing from response entirely, skipping")
            continue
        obj_result = parsed[obj_name]
        utilization_by_obj[obj_name] = []

        if "narrative" in obj_result:
            narratives_out[obj_name] = obj_result["narrative"]
        elif frozen_narratives and obj_name in frozen_narratives:
            narratives_out[obj_name] = frozen_narratives[obj_name]
        else:
            print(f"  WARNING: '{obj_name}' has no narrative in response and no frozen narrative "
                  f"to fall back to — narratives.json will be missing this object")

        deform_outputs[obj_name] = {}
        targets = obj_result.get("targets", {})
        for kf in range(n_keyframes):
            kf_key = f"kf{kf}"
            if kf_key not in targets:
                print(f"  WARNING: '{obj_name}' missing {kf_key} targets, skipping this frame")
                continue
            kf_result = targets[kf_key]
            out = {}
            joint_targets_this_kf = {}
            for name, target in kf_result.items():
                m = re.match(r"joint_(\d+)", name)
                if not m:
                    print(f"  WARNING: unexpected key '{name}' for '{obj_name}' {kf_key}, skipping")
                    continue
                joint_i = int(m.group(1))
                if joint_i >= len(info["joint_mesh_indices"]):
                    print(f"  WARNING: '{obj_name}' joint_{joint_i} out of range, skipping")
                    continue
                mesh_idx = info["joint_mesh_indices"][joint_i]

                rest = np.array(info["joints"][joint_i])
                cap = round(info["bbox_size"] * 0.25, 1)
                target_arr = np.array(target, dtype=float)
                disp = target_arr - rest
                dist = np.linalg.norm(disp)
                utilization_by_obj[obj_name].append(min(dist / cap, 1.0) if cap > 0 else 0.0)
                if dist > cap:
                    clipped = rest + disp / dist * cap
                    print(f"  CLIPPED '{obj_name}' {kf_key} joint_{joint_i}: requested {dist:.1f}px "
                          f"(cap {cap}px) -> clipped to {cap}px, direction preserved")
                    target = clipped.tolist()

                out[str(mesh_idx)] = target
                joint_targets_this_kf[f"joint_{joint_i}"] = target
            anchor_mesh_idx = info["joint_mesh_indices"][info["anchor_idx"]]
            out[str(anchor_mesh_idx)] = info["joints"][info["anchor_idx"]].tolist()
            deform_outputs[obj_name][kf_key] = joint_targets_this_kf

            out_path = os.path.join(out_dir, f"qwen_{sketch_name}_{obj_name}_kf{kf}.json")
            with open(out_path, "w") as f:
                json.dump(out, f, indent=2)
            print(f"  wrote {out_path}")

    cap_utilization = {}
    for obj_name, fractions in utilization_by_obj.items():
        if fractions:
            mean_frac = sum(fractions) / len(fractions)
            cap_utilization[obj_name] = mean_frac
            print(f"  '{obj_name}': mean cap utilization = {mean_frac*100:.0f}% "
                  f"(across {len(fractions)} joint-keyframe pairs)")

    return narratives_out, deform_outputs, cap_utilization
# =============================================================================
# STEP 4: render — compose the full scene (no Qwen call at all)
# =============================================================================

def run_render(handles_dir, svg_path, semantic_path, traj_path,
                out_path=None, frames_dir=None):
    """
    out_path: if given, ALSO saves the combined strip image (all 5 keyframes
              side by side) here, outside frames_dir. Optional — pass None
              to keep output confined to frames_dir only.
    frames_dir: if given, saves each keyframe as its own individual PNG
                (kf0.png ... kf4.png) plus the combined strip, named after
                the sketch itself ({sketch_name}.png), all inside this one
                folder.
    """
    import matplotlib.pyplot as plt

    sketch_name = os.path.splitext(os.path.basename(svg_path))[0]
    real_trajectories = load_trajectories(traj_path)
    object_names = list(real_trajectories.keys())

    object_data = {}
    for name in object_names:
        points, slices = load_object(name, svg_path, semantic_path)
        dx_vals, dy_vals = bbox_deltas(real_trajectories[name])
        unique_points, p2u = deduplicate_points(points, tol=0.35)
        tri, edges = build_mesh(unique_points)
        object_data[name] = {
            "points": points, "slices": slices, "dx": dx_vals, "dy": dy_vals,
            "unique_points": unique_points, "p2u": p2u, "edges": edges,
        }

    if frames_dir:
        os.makedirs(frames_dir, exist_ok=True)

    xmin, xmax, ymin, ymax = 0, 260, 60, 230

    # compute each keyframe's drawing data once, reused for both the combined
    # strip and the individual per-keyframe images
    keyframe_lines = []  # list of {name: [(seg_x, seg_y), ...]} per keyframe
    for kf in range(N_KEYFRAMES):
        lines_this_kf = {}
        for name in object_names:
            od = object_data[name]
            handles_path = os.path.join(handles_dir, f"qwen_{sketch_name}_{name}_kf{kf}.json")

            if os.path.exists(handles_path):
                with open(handles_path) as f:
                    spec = json.load(f)
                handle_indices = [int(k) for k in spec.keys()]
                handle_targets = np.array([spec[k] for k in spec.keys()])
                n_verts = len(od["unique_points"])
                bad = [i for i in handle_indices if i >= n_verts]
                if bad:
                    print(f"  ERROR: {handles_path} has out-of-bounds indices {bad} for '{name}' "
                          f"({n_verts} mesh vertices) — likely from a DIFFERENT sketch's mesh. "
                          f"Falling back to translation-only.")
                    deformed_points = od["points"]
                    mode = "translation-only (handles file failed validation)"
                else:
                    deformed_unique = arap_deform(od["unique_points"], od["edges"],
                                                   handle_indices, handle_targets, iterations=10)
                    deformed_points = deformed_unique[od["p2u"]]
                    mode = "ARAP"
            else:
                deformed_points = od["points"]
                mode = "translation-only (no handles file found)"

            moved = deformed_points + np.array([od["dx"][kf], od["dy"][kf]])
            lines_this_kf[name] = [moved[start:end] for start, end in od["slices"]]

            print(f"kf{kf} '{name}': {mode}, points_after_move_range="
                  f"x[{moved[:,0].min():.1f},{moved[:,0].max():.1f}] "
                  f"y[{moved[:,1].min():.1f},{moved[:,1].max():.1f}]")

        keyframe_lines.append(lines_this_kf)

        if frames_dir:
            fig_i, ax_i = plt.subplots(figsize=(6, 5.5))
            for name, segs in lines_this_kf.items():
                for seg in segs:
                    ax_i.plot(seg[:, 0], seg[:, 1],
                              color=OBJECT_COLORS.get(name, DEFAULT_COLOR),
                              linewidth=OBJECT_LINEWIDTH.get(name, DEFAULT_LINEWIDTH))
            ax_i.set_xlim(xmin, xmax)
            ax_i.set_ylim(ymax, ymin)
            ax_i.set_aspect("equal")
            ax_i.set_title(f"{sketch_name} — kf{kf}", fontsize=12, fontweight="bold")
            frame_path = os.path.join(frames_dir, f"kf{kf}.png")
            fig_i.savefig(frame_path, dpi=140, bbox_inches="tight")
            plt.close(fig_i)
            print(f"  wrote {frame_path}")

    # combined strip, same as before
    fig, axes = plt.subplots(1, N_KEYFRAMES, figsize=(24, 5))
    for kf in range(N_KEYFRAMES):
        ax = axes[kf]
        for name, segs in keyframe_lines[kf].items():
            for seg in segs:
                ax.plot(seg[:, 0], seg[:, 1],
                        color=OBJECT_COLORS.get(name, DEFAULT_COLOR),
                        linewidth=OBJECT_LINEWIDTH.get(name, DEFAULT_LINEWIDTH))
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(ymax, ymin)
        ax.set_aspect("equal")
        ax.set_title(f"kf{kf}", fontsize=13, fontweight="bold")

    plt.tight_layout()

    if out_path:
        plt.savefig(out_path, dpi=140, bbox_inches="tight")
        print(f"wrote {out_path}")

    if frames_dir:
        combined_frame_path = os.path.join(frames_dir, f"{sketch_name}.png")
        plt.savefig(combined_frame_path, dpi=140, bbox_inches="tight")
        print(f"wrote {combined_frame_path}")

    if not out_path and not frames_dir:
        print("WARNING: neither out_path nor frames_dir given, combined strip image not saved anywhere")

    plt.close(fig)


# =============================================================================
# CLI — single input: a sketch name or an SVG path. Everything else is
# derived automatically from the directory conventions used throughout
# this dataset. Override flags exist for the rare case a path doesn't
# match convention, but nothing is required beyond the sketch itself.
# =============================================================================

# Confirmed real paths from this dataset, used as defaults so nothing else
# needs to be typed per run. If your layout differs, override with the
# corresponding --*-dir / --*-file flag below.
SVG_DIR_DEFAULT = "/user/HS400/rk01499/my_scratch/sketch/data/raw/60sketches/svg"
PROCESSED_DIR_DEFAULT = "/user/HS400/rk01499/my_scratch/sketch/data/processed"
CAPTION_FILE_DEFAULT = "/user/HS400/rk01499/my_scratch/sketch/data/raw/60sketches/caption.txt"
MODEL_PATH_DEFAULT = "/user/HS400/rk01499/my_scratch/models/qwen2.5-7b/"


def resolve_sketch_paths(sketch, svg_dir, processed_dir, caption_file):
    """
    sketch: either a bare sketch name ("dog9") or a path to its SVG
            ("/path/to/dog9.svg") — either way, everything else (semantic,
            traj, caption) is derived from the same naming convention used
            across this dataset: {name}.svg, {name}/{name}_semantic.txt,
            {name}/{name}_traj.txt, and a lookup in one shared caption.txt.
    """
    name = os.path.splitext(os.path.basename(sketch))[0]
    svg_path = sketch if sketch.endswith(".svg") else os.path.join(svg_dir, f"{name}.svg")
    semantic_path = os.path.join(processed_dir, name, f"{name}_semantic.txt")
    traj_path = os.path.join(processed_dir, name, f"{name}_traj.txt")

    missing = [p for p in [svg_path, semantic_path, traj_path, caption_file] if not os.path.exists(p)]
    if missing:
        raise SystemExit(
            f"Could not find these expected files for sketch '{name}':\n  " +
            "\n  ".join(missing) +
            "\n\nIf your directory layout differs from the default, pass --svg-dir / "
            "--processed-dir / --caption-file explicitly."
        )

    caption = get_caption(caption_file, name)
    return name, svg_path, semantic_path, traj_path, caption


def main():
    ap = argparse.ArgumentParser(
        description="Run the full sketch deformation pipeline for one image. "
                     "The only required input is the sketch — everything else "
                     "(semantic assignments, trajectory, caption) is looked up "
                     "automatically from the standard dataset layout.")
    ap.add_argument("sketch", type=str,
                     help="sketch name (e.g. 'dog9') or path to its .svg file")
    ap.add_argument("--model", type=str, default=MODEL_PATH_DEFAULT)
    ap.add_argument("--svg-dir", type=str, default=SVG_DIR_DEFAULT)
    ap.add_argument("--processed-dir", type=str, default=PROCESSED_DIR_DEFAULT)
    ap.add_argument("--caption-file", type=str, default=CAPTION_FILE_DEFAULT)
    ap.add_argument("--out-dir", type=str, default=".")
    ap.add_argument("--no-fewshot", action="store_true",
                     help="disable the dog3 few-shot example in narrate (for A/B comparison)")
    ap.add_argument("--deform-only", action="store_true",
                     help="run classify + ONE narrate+deform attempt + render, then STOP — no judge, "
                          "no retries. Prints cap utilization directly and saves the render, so you can "
                          "inspect raw generation quality without the judge's assessment as a confound.")
    args = ap.parse_args()

    name, svg_path, semantic_path, traj_path, caption = resolve_sketch_paths(
        args.sketch, args.svg_dir, args.processed_dir, args.caption_file)
    print(f"sketch: {name}")
    print(f"  svg:      {svg_path}")
    print(f"  semantic: {semantic_path}")
    print(f"  traj:     {traj_path}")
    print(f"  caption:  {caption!r}")

    os.makedirs(args.out_dir, exist_ok=True)
    json_dir = os.path.join(args.out_dir, "json", name)
    os.makedirs(json_dir, exist_ok=True)
    print(f"  json output dir: {json_dir}")

    print("\n########## STEP 1: CLASSIFY ##########")
    classify_model, classify_tokenizer, classify_device = load_qwen_model(args.model)
    deformation, arap_objects = run_classify(
        classify_model, classify_tokenizer, classify_device, caption, semantic_path,
        os.path.join(json_dir, f"{name}_deformation.json"))
    classify_model = unload_model(classify_model)

    if arap_objects:
        import vlm_judge

        temp_dir = os.path.join(args.out_dir, "temp", name)
        objects_info = build_objects_info(svg_path, semantic_path, arap_objects)
        if not objects_info:
            print("No objects with valid handles found after mesh setup. Nothing to do.")
            objects_info = None

        rest_pose_image_path = os.path.join(json_dir, f"{name}_rest_pose.png")
        render_rest_pose_multi(arap_objects, svg_path, semantic_path, rest_pose_image_path)
        # NOTE: in svg.py this image is generated but never fed to the model —
        # narrate+deform is text-only here. Kept on disk purely as a human-
        # readable reference for comparing against the stroke-geometry text.

        import dino_similarity
        dino_model, dino_processor = dino_similarity.load_dino_model()

        import clip_score
        clip_model, clip_processor = clip_score.load_clip_model()

        feedback = None
        previous_narratives = None      # {obj_name: [5 descriptions]} from the last attempt
        previous_temp_dir = None      # where the last attempt's rendered kf0..kf4 images live
        freeze_narrative = False        # only frozen once faithfulness has already passed once
        consecutive_stagnant = 0        # early-stop if DINOv2 confirms no real change 2 attempts in a row
        previous_deform_outputs = None  # {obj_name: {kf_key: {joint_i: [x,y]}}} — svg.py's text-only
                                          # replacement for "show it the previous rendered images"
        final_verdict = None
        winning_attempt = None
        all_attempts_summary = []

        for attempt in range(1, MAX_RETRIES + 1):
            if not objects_info:
                break
            print(f"\n########## ATTEMPT {attempt}/{MAX_RETRIES} ##########")

            attempt_json_dir = os.path.join(json_dir, "attempts", f"attempt_{attempt}")
            attempt_temp_dir = os.path.join(temp_dir, "attempts", f"attempt_{attempt}")
            os.makedirs(attempt_json_dir, exist_ok=True)

            # NO IMAGES AT ALL in svg.py — narrate+deform relies entirely on
            # text: rest-pose stroke geometry (always) and, on retries, the
            # previous attempt's reconstructed deformed shape (also text,
            # via previous_deform_outputs -> build_deformed_stroke_geometry_text
            # inside build_combined_narrate_deform_prompt).
            is_retry = attempt > 1

            prompt = build_combined_narrate_deform_prompt(
                objects_info, caption, previous_narratives=previous_narratives,
                feedback=feedback, is_retry=is_retry, few_shot=not args.no_fewshot,
                freeze_narrative=freeze_narrative, previous_deform_outputs=previous_deform_outputs)

            print(f"\n---------- STEP 2+3: NARRATE+DEFORM (attempt {attempt}, "
                  f"NO IMAGES, text-only"
                  f"{', narrative FROZEN' if freeze_narrative else ''}) ----------")
            vlm_model, vlm_processor = vlm_judge.load_vlm("Qwen/Qwen3-VL-4B-Instruct")
            response = run_combined_narrate_deform(vlm_model, vlm_processor, [], prompt)
            try:
                parsed = parse_json_response(response)
            except (ValueError, json.JSONDecodeError) as e:
                print(f"FAILED TO PARSE: {e}\nraw: {response}")
                vlm_model = unload_model(vlm_model)
                feedback = "the previous attempt's output could not be parsed; produce valid JSON in the exact requested format"
                all_attempts_summary.append({"attempt": attempt, "plausibility_score": None, "note": "narrate+deform parse failed"})
                continue

            narratives_this_attempt, deform_outputs_this_attempt, cap_utilization_this_attempt = apply_deform_clip_and_write(
                parsed, objects_info, attempt_json_dir, name,
                frozen_narratives=previous_narratives if freeze_narrative else None)
            previous_deform_outputs = deform_outputs_this_attempt
            with open(os.path.join(attempt_json_dir, f"{name}_narratives.json"), "w") as f:
                json.dump(narratives_this_attempt, f, indent=2)

            print(f"\n---------- STEP 4: RENDER (attempt {attempt}, no Qwen) ----------")
            run_render(attempt_json_dir, svg_path, semantic_path, traj_path, frames_dir=attempt_temp_dir)

            if args.deform_only:
                print(f"\n########## --deform-only: STOPPING after attempt 1, no judge ##########")
                print(f"cap_utilization (raw, unfiltered by any threshold):")
                for obj, frac in (cap_utilization_this_attempt or {}).items():
                    print(f"  {obj}: {frac*100:.1f}% of allowed movement used")
                print(f"\nInspect the actual render directly at: {attempt_temp_dir}")
                print(f"(kf0.png ... kf4.png, plus the combined strip)")
                sys.exit(0)

            stagnation_result = None
            if previous_temp_dir:
                print(f"\n---------- STAGNATION CHECK (attempt {attempt} vs attempt {attempt - 1}) ----------")
                prev_dino_images = vlm_judge.load_keyframe_images(previous_temp_dir)
                curr_dino_images = vlm_judge.load_keyframe_images(attempt_temp_dir)
                stagnation_result = dino_similarity.stagnation_score(
                    dino_model, dino_processor, prev_dino_images, curr_dino_images)
                stagnant = dino_similarity.is_stagnant(stagnation_result)
                print(f"mean attempt-to-attempt similarity: {stagnation_result['mean_similarity']:.4f} "
                      f"({'STAGNANT' if stagnant else 'changed'})")
                consecutive_stagnant = consecutive_stagnant + 1 if stagnant else 0

            print(f"\n---------- TEMPORAL CONSISTENCY (attempt {attempt}, diagnostic only) ----------")
            temporal_images = vlm_judge.load_keyframe_images(attempt_temp_dir)
            temporal_result = dino_similarity.temporal_consistency(dino_model, dino_processor, temporal_images)

            clip_result = None
            if caption:
                print(f"\n---------- CLIP SCORE (attempt {attempt}) ----------")
                clip_images = vlm_judge.load_keyframe_images(attempt_temp_dir)
                clip_result = clip_score.compute_sequence_clip_scores(clip_model, clip_processor, clip_images, caption)

            print(f"\n---------- STEP 5: JUDGE (attempt {attempt}, SVG-based, no images) ----------")
            # reuse the SAME already-loaded VLM for judging — no reload
            # needed, since narrate+deform and judge are both Qwen3-VL calls now
            all_kf_geometry_text = build_all_keyframes_geometry_text(objects_info, deform_outputs_this_attempt)
            judge_prompt = build_svg_judge_prompt(
                name, all_kf_geometry_text, caption, dino_stagnation=stagnation_result,
                dino_temporal=temporal_result, clip_scores=clip_result,
                cap_utilization=cap_utilization_this_attempt)
            judge_response = run_svg_judge(vlm_model, vlm_processor, judge_prompt)
            vlm_model = unload_model(vlm_model)

            try:
                verdict = vlm_judge.parse_judge_response(judge_response)
            except (ValueError, json.JSONDecodeError) as e:
                print(f"JUDGE FAILED TO PARSE: {e}\nraw: {judge_response}")
                print("Treating as a failed attempt, retrying without specific feedback.")
                feedback = "the previous attempt's evaluation could not be parsed; try a clearer, more varied pose progression"
                previous_narratives = narratives_this_attempt
                previous_temp_dir = attempt_temp_dir
                freeze_narrative = False  # unknown state — safest to regenerate rather than assume faithfulness held
                all_attempts_summary.append({"attempt": attempt, "plausibility_score": None, "note": "judge parse failed"})
                continue

            problems = vlm_judge.validate_judge_response(verdict)
            print("\n--- JUDGE VERDICT ---")
            print(json.dumps(verdict, indent=2))
            if problems:
                for p in problems:
                    print(f"  VALIDATION PROBLEM: {p}")

            with open(os.path.join(attempt_json_dir, f"{name}_judge_verdict.json"), "w") as f:
                json.dump(verdict, f, indent=2)

            final_verdict = verdict
            winning_attempt = attempt
            score = verdict.get("plausibility_score")
            faith_score = verdict.get("faithfulness_score")
            quality_score = verdict.get("quality_score")
            all_attempts_summary.append({"attempt": attempt, "plausibility_score": score,
                                          "plausibility_notes": verdict.get("plausibility_notes"),
                                          "faithfulness_score": faith_score,
                                          "faithfulness_notes": verdict.get("faithfulness_notes"),
                                          "quality_score": quality_score,
                                          "quality_notes": verdict.get("quality_notes"),
                                          "dino_stagnant": dino_similarity.is_stagnant(stagnation_result) if stagnation_result else None})

            plausibility_ok = isinstance(score, (int, float)) and score >= PLAUSIBILITY_THRESHOLD
            faithfulness_ok = faithfulness_passed(faith_score)
            quality_ok = isinstance(quality_score, (int, float)) and quality_score >= QUALITY_THRESHOLD
            print(f"\nplausibility_score = {score} (threshold = {PLAUSIBILITY_THRESHOLD}, "
                  f"{'PASS' if plausibility_ok else 'FAIL'})")
            print(f"faithfulness_score = {faith_score} (threshold = {FAITHFULNESS_THRESHOLD}, "
                  f"{'PASS' if faithfulness_ok else 'FAIL'})")
            print(f"quality_score = {quality_score} (threshold = {QUALITY_THRESHOLD}, "
                  f"{'PASS' if quality_ok else 'FAIL'})")

            if plausibility_ok and faithfulness_ok and quality_ok:
                print(f"All three thresholds met on attempt {attempt} — stopping.")
                break

            if consecutive_stagnant >= 2:
                print(f"\nDINOv2 confirmed NO real change across {consecutive_stagnant} consecutive attempts "
                      f"(attempt {attempt} vs {attempt-1}, and {attempt-1} vs {attempt-2}) — further retries "
                      f"are very unlikely to help. Stopping early and keeping this attempt's result rather "
                      f"than burning through the remaining {MAX_RETRIES - attempt} attempts.")
                break

            previous_narratives = narratives_this_attempt
            previous_temp_dir = attempt_temp_dir
            # freeze the narrative on the NEXT attempt only if faithfulness
            # already passed THIS attempt — no reason to keep regenerating
            # a story that's already correct, only the numbers need work
            freeze_narrative = faithfulness_ok

            if attempt == MAX_RETRIES:
                print("Below threshold on the final attempt — no retry left, skipping feedback construction.")
            else:
                plausibility_note = verdict.get("plausibility_notes", "the pose progression needs to look more plausible")
                faithfulness_note = verdict.get("faithfulness_notes")
                quality_note = verdict.get("quality_notes")
                notes = [f"plausibility: {plausibility_note}"]
                if not faithfulness_ok and faithfulness_note and faithfulness_note != "N/A":
                    notes.append(f"faithfulness to the intended action: {faithfulness_note}")
                if not quality_ok and quality_note:
                    notes.append(f"rendering quality: {quality_note}")
                feedback = " | ".join(notes)

                # DINOv2 objective override: if the images barely changed from the
                # previous attempt, say so explicitly and forcefully — this is the
                # exact failure mode confirmed on real hardware (cannon1 ran all
                # MAX_RETRIES with no real change), where the judge's own text
                # critique was never specific enough for Qwen to act on. An
                # objective embedding-distance number doesn't have that problem.
                if stagnation_result and dino_similarity.is_stagnant(stagnation_result):
                    feedback = (f"CRITICAL: your last attempt was measured as nearly IDENTICAL to the one "
                                f"before it (DINOv2 similarity {stagnation_result['mean_similarity']:.3f}) — "
                                f"you are NOT making real changes. You MUST produce substantially different "
                                f"target positions this time, not a superficial rewording. Original feedback: {feedback}")
                elif cap_utilization_this_attempt:
                    # objective override: CONFIRMED on real hardware (horsecar5's walking
                    # person) that Qwen can propose displacement well within the allowed
                    # cap without any prior signal catching it — technically nonzero
                    # movement that's visually negligible at render scale. The hard clip
                    # only ever caught OVER-cap requests; this catches UNDER-using it.
                    low_objs = {obj: frac for obj, frac in cap_utilization_this_attempt.items()
                                if frac < LOW_UTILIZATION_THRESHOLD}
                    if low_objs:
                        low_str = ", ".join(f"{obj} used only {frac*100:.0f}%" for obj, frac in low_objs.items())
                        feedback = (f"CRITICAL: {low_str} of their allowed movement range — this is too timid "
                                    f"to be visible. You MUST commit to larger, more decisive displacement "
                                    f"(up to the stated cap) for these objects. Original feedback: {feedback}")

                print(f"Below threshold — retrying with feedback: {feedback!r} "
                      f"(narrative will be {'FROZEN' if freeze_narrative else 'regenerated'})")
        else:
            print(f"\nReached MAX_RETRIES ({MAX_RETRIES}) without meeting both thresholds. "
                  f"Using the last attempt's result.")

        # print a compact table so the score progression across attempts is
        # visible in one place, not just scattered through the full log
        print("\n########## ATTEMPT SUMMARY ##########")
        for a in all_attempts_summary:
            print(f"  attempt {a['attempt']}: plausibility={a.get('plausibility_score')} "
                  f"faithfulness={a.get('faithfulness_score')} quality={a.get('quality_score')} "
                  f"dino_stagnant={a.get('dino_stagnant')} "
                  f"— {a.get('plausibility_notes') or a.get('note', '')}")

        if winning_attempt:
            import shutil
            winning_json = os.path.join(json_dir, "attempts", f"attempt_{winning_attempt}")
            winning_temp = os.path.join(temp_dir, "attempts", f"attempt_{winning_attempt}")
            for f in os.listdir(winning_json):
                shutil.copy2(os.path.join(winning_json, f), os.path.join(json_dir, f))
            for f in os.listdir(winning_temp):
                src = os.path.join(winning_temp, f)
                if os.path.isfile(src):
                    shutil.copy2(src, os.path.join(temp_dir, f))
            print(f"\ncopied winning attempt ({winning_attempt}) to the top-level "
                  f"json/{name}/ and temp/{name}/ locations")
            print(f"all {len(all_attempts_summary)} attempts preserved under "
                  f"json/{name}/attempts/ and temp/{name}/attempts/ for comparison")
    else:
        print("\nNo ARAP objects — skipping narrate/deform/judge entirely.")
        temp_dir = os.path.join(args.out_dir, "temp", name)
        print("\n########## RENDER (no Qwen) ##########")
        run_render(json_dir, svg_path, semantic_path, traj_path, frames_dir=temp_dir)


if __name__ == "__main__":
    main()