"""
Standalone seed-sensitivity test for narrate+deform.

WHY THIS EXISTS
----------------
On eat2, one full pipeline run (3 attempts, each reloading the VLM from
scratch and getting different feedback) produced BYTE-IDENTICAL joint
targets across all 3 attempts, at temperature=0.6 with do_sample=True.
That could mean one of two very different things:

  (a) The model's output distribution for THIS prompt/image is genuinely
      peaked enough that independent sampling calls land on the same
      answer most of the time — a real property of the model+input, not
      a bug. If so, no code fix here can change it; the fix (if any) is
      changing the model, the prompt, or the images.

  (b) Something in the pipeline (most likely: reloading the model fresh
      before every attempt) is putting the RNG into the same state each
      time, so "random" sampling isn't actually random across attempts.
      If so, this IS fixable in code (e.g. load the model once per run,
      not once per attempt).

This script isolates the test: load the model ONCE, then generate from
the SAME prompt+images TWICE in the same process — once letting the RNG
run naturally, once after an explicit, different manual_seed. No model
reload happens between the two calls. If the two outputs still match,
that's evidence for (a). If they differ, that's evidence AGAINST (a) for
this specific in-process case, which strengthens (b) as the thing to fix
in the real pipeline (since the real pipeline's difference from this
script is exactly the "reload before every attempt" step).

This does NOT run the whole pipeline, ARAP, or rendering — narrate+deform
only, since that's the specific call under suspicion.

USAGE
-----
Grab the exact prompt text for one attempt straight from what the
pipeline already saves to disk, e.g.:
    ./json/eat2__<run_label>/attempts/attempt_1/eat2_person_narrate_deform_prompt.txt
(exact filename pattern may vary slightly by version — look inside the
attempt_1 folder for whichever file has "_prompt" in the name for the
object you care about).

Then:
    python test_seed_sensitivity.py \\
        --model-path /user/HS400/rk01499/my_scratch/models/qwen2.5-vl-3b/ \\
        --prompt-file /path/to/that/saved/prompt.txt \\
        --image /path/to/eat2_rest_pose.png \\
        --temperature 0.6 \\
        --seed-b 999

If you don't have the rest-pose image path handy, --image can be omitted
if the saved prompt was TEXT-ONLY for that attempt (check the file — if
it references image content the model needs to see, you need --image).
"""
import argparse
import sys


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model-path", required=True,
                     help="path to the Qwen2.5-VL checkpoint, same value you'd pass elsewhere in the pipeline")
    ap.add_argument("--prompt-file", required=True,
                     help="path to a .txt file containing the EXACT prompt text to reuse — copy this "
                          "straight from a saved attempt_N/*_prompt.txt in a real run's json output dir")
    ap.add_argument("--image", action="append", default=[],
                     help="path to an image the prompt references (repeat --image for multiple views). "
                          "Omit only if the saved prompt was genuinely text-only for this attempt.")
    ap.add_argument("--temperature", type=float, default=0.6,
                     help="must match whatever temperature produced the identical-repeat you're "
                          "investigating — default 0.6 matches this pipeline's current default")
    ap.add_argument("--seed-a", type=int, default=None,
                     help="seed for call A. Default: none (let RNG run naturally, whatever state it's "
                          "already in — this is the more realistic 'first call in a fresh process' case)")
    ap.add_argument("--seed-b", type=int, default=999,
                     help="seed for call B — deliberately different from call A, so if the two outputs "
                          "still match, that's not explained by 'same seed was used both times'")
    ap.add_argument("--max-new-tokens", type=int, default=2500,
                     help="matches pipe11.py's run_combined_narrate_deform: 500 * N_KEYFRAMES, "
                          "N_KEYFRAMES=5 by default")
    args = ap.parse_args()

    import torch
    from PIL import Image
    from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor

    with open(args.prompt_file) as f:
        prompt_text = f.read()

    images = [Image.open(p).convert("RGB") for p in args.image]

    print(f"loading model from {args.model_path} ...")
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(args.model_path, device_map="auto", dtype="auto")
    processor = AutoProcessor.from_pretrained(args.model_path)
    print("model loaded — will NOT be reloaded between call A and call B.\n")

    def generate_once(seed):
        if seed is not None:
            torch.manual_seed(seed)
        content = [{"type": "image", "image": img} for img in images]
        content.append({"type": "text", "text": prompt_text})
        messages = [{"role": "user", "content": content}]
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = processor(text=[text], images=images if images else None, return_tensors="pt").to(model.device)
        with torch.no_grad():
            output_ids = model.generate(**inputs, max_new_tokens=args.max_new_tokens,
                                         temperature=args.temperature, do_sample=True)
        generated = output_ids[:, inputs["input_ids"].shape[1]:]
        return processor.batch_decode(generated, skip_special_tokens=True)[0]

    print(f"--- CALL A (seed={args.seed_a}) ---")
    response_a = generate_once(args.seed_a)
    print(response_a)

    print(f"\n--- CALL B (seed={args.seed_b}) ---")
    response_b = generate_once(args.seed_b)
    print(response_b)

    print("\n" + "=" * 60)
    if response_a == response_b:
        print("RESULT: IDENTICAL — same output despite different seeds, same process, no reload.")
        print("This means the model's distribution for THIS prompt/image is genuinely peaked —")
        print("not a pipeline bug. A code fix (e.g. 'load model once per run') would NOT help.")
        print("Next step would be changing the prompt, the images, or the model, not this pipeline's")
        print("reload/retry logic.")
    else:
        print("RESULT: DIFFERENT — outputs changed once the seed was actually varied.")
        print("This means sampling DOES work normally in-process. The eat2 identical-repeat across")
        print("real pipeline attempts is then more likely caused by something specific to the real")
        print("pipeline's per-attempt flow (most suspect: reloading the model fresh before every")
        print("attempt) rather than a fundamental property of the model+prompt. Worth checking next:")
        print("does load_vlm() / from_pretrained() put the RNG into a repeatable state, and does")
        print("loading the model ONCE per run instead of once per attempt fix the real pipeline?")
        sys.exit(0)


if __name__ == "__main__":
    main()