pipe_1.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
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 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 format_joint_feedback_for_object(object_name, joint_feedback):
    """
    Filters the judge's structured joint_feedback list down to entries for
    THIS object and formats them as explicit, actionable lines — e.g.
    "joint_2 at kf2, kf3: reaching too little toward the frisbee -> move
    further up-left." This replaces feeding the whole undifferentiated
    critique string into every object's section: each object now only
    sees the feedback that's actually about it, in a form that names the
    exact joint and keyframe(s) rather than describing the problem in prose.
    Returns "" if there's no feedback for this object (nothing to add).
    """
    if not joint_feedback:
        return ""
    relevant = [f for f in joint_feedback if f.get("object") == object_name]
    if not relevant:
        return ""
    lines = []
    for f in relevant:
        kfs = ", ".join(f"kf{k}" for k in f.get("keyframes", []))
        lines.append(
            f"    - joint_{f.get('joint')} at {kfs or 'unspecified keyframe(s)'}: {f.get('issue', '')} "
            f"-> {f.get('suggested_direction', 'adjust as needed')}"
        )
    return "\n".join(lines)


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, joint_feedback=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 the
    attached image(s) and feedback, same as before.

    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.

    joint_feedback: the FULL joint_feedback list from the judge's verdict
    (all objects) — filtered down to this object's entries and formatted
    as precise per-joint lines. Falls back to nothing (not an error) if
    the judge didn't return joint_feedback (e.g. coordinate context wasn't
    given to build_judge_prompt) — the object still gets the old flat
    `feedback` string via previous_block/feedback_line below.
    """
    cap = round(bbox_size * 0.25, 1)
    joint_feedback_text = format_joint_feedback_for_object(object_name, joint_feedback)
    joint_feedback_block = (
        f"\n  Specific per-joint corrections from the judge (apply these precisely, this is not general "
        f"guidance):\n{joint_feedback_text}\n"
    ) if joint_feedback_text else ""
    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"
        )

    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  The attached images show exactly what the previous attempt's target positions " \
                         "actually looked like when rendered — use them to see specifically what needs to change numerically."
        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}
{joint_feedback_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("The attached images show exactly what that previous attempt actually looked like when "
                      "rendered. Look at them, understand what specifically was wrong, and revise BOTH the "
                      "narrative and the target positions to fix it — don't just reword the narrative "
                      "superficially while leaving the same underlying problem.")
        previous_block = "\n  " + "\n  ".join(parts) + "\n"

    return f"""Object: "{object_name}"
  Joints:
{joint_lines}
{geometry_block}
{previous_block}
{joint_feedback_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, joint_feedback=None):
    sections, example_parts = [], []
    for name, info in objects_info.items():
        prev_narrative_for_obj = (previous_narratives or {}).get(name)
        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"), joint_feedback=joint_feedback,
        ))
        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 = ("The FIRST image attached is the object's original rest pose (undeformed). "
                          "The remaining images are the actual rendered result of your PREVIOUS attempt, "
                          "one per keyframe, in order.")
    else:
        image_context = ("The attached image shows the object's original rest pose (undeformed) — use this "
                          "to understand what strokes actually 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 — reused
    here since both are Qwen3-VL calls with a list of images + one 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=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 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, "P_1", 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)

        # preprocessed bbox trajectory data — fixed ground truth, given to the
        # judge as spatial context for EVERY object (ARAP and TRAJ_ONLY alike),
        # not something the judge critiques or the generator controls
        real_trajectories = load_trajectories(traj_path)
        all_object_names = list(real_trajectories.keys())

        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
        joint_feedback = None           # judge's structured per-joint corrections from the last attempt
        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)

            # image list: always the rest pose; from attempt 2+, ALSO the
            # previous attempt's actual rendered keyframes, so Qwen sees
            # exactly what its last attempt looked like, not just a text
            # description of it
            from PIL import Image
            images = [Image.open(rest_pose_image_path).convert("RGB")]
            is_retry = attempt > 1
            if is_retry:
                images.extend(vlm_judge.load_keyframe_images(previous_temp_dir))

            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, joint_feedback=joint_feedback)

            print(f"\n---------- STEP 2+3: NARRATE+DEFORM (attempt {attempt}, "
                  f"{len(images)} image{'s' if len(images) != 1 else ''}"
                  f"{', narrative FROZEN' if freeze_narrative else ''}) ----------")
            vlm_model, vlm_processor = vlm_judge.load_vlm("Qwen/Qwen2.5-VL-3B-Instruct")
            response = run_combined_narrate_deform(vlm_model, vlm_processor, images, 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)
            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}, images + joint/bbox coordinates) ----------")
            # reuse the SAME already-loaded VLM for judging — no reload
            # needed, since narrate+deform and judge are both Qwen3-VL calls now.
            # Judge now sees: rest pose + this attempt's 5 rendered keyframes
            # (images), PLUS rest-pose stroke text + joint legend + this
            # attempt's joint targets + every object's fixed bbox trajectory
            # (text) — NOT raw per-stroke deformed coordinates, which gave no
            # joint identity to anchor feedback to. This is what makes
            # joint_feedback (precise "move joint_2 up at kf3" critique)
            # possible instead of only a vague overall_verdict.
            judge_images = [Image.open(rest_pose_image_path).convert("RGB")]
            judge_images.extend(vlm_judge.load_keyframe_images(attempt_temp_dir))

            judge_prompt = vlm_judge.build_judge_prompt(
                name, caption=caption, dino_stagnation=stagnation_result,
                dino_temporal=temporal_result, clip_scores=clip_result,
                objects_info=objects_info, deform_outputs=deform_outputs_this_attempt,
                real_trajectories=real_trajectories, all_object_names=all_object_names)
            judge_response = vlm_judge.run_judge(vlm_model, vlm_processor, judge_images, 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
                joint_feedback = None  # verdict didn't parse, so any joint_feedback in it is unusable/unknown — don't carry stale feedback forward
                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
            # carry the judge's structured per-joint corrections into the NEXT
            # attempt's prompt — empty list is valid (judge found nothing
            # specific to fix), missing key means coordinate context wasn't
            # given to build_judge_prompt at all; either way, default to None
            # so format_joint_feedback_for_object just adds nothing
            joint_feedback = verdict.get("joint_feedback")
            # 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}")

                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 P_1/{name}/ locations")
            print(f"all {len(all_attempts_summary)} attempts preserved under "
                  f"json/{name}/attempts/ and P_1/{name}/attempts/ for comparison")
    else:
        print("\nNo ARAP objects — skipping narrate/deform/judge entirely.")
        temp_dir = os.path.join(args.out_dir, "P_1", name)
        print("\n########## RENDER (no Qwen) ##########")
        run_render(json_dir, svg_path, semantic_path, traj_path, frames_dir=temp_dir)


if __name__ == "__main__":
    main()