"""
Minimal script to verify PR #13461 does not change ZImagePipeline outputs.

Compares latent outputs between the current branch (hitchhiker3010-main, with PR changes)
and the main branch (without PR changes) using a fixed seed.

The PR makes two changes:
  1. Precompute cfg_truncation t_norms outside the loop (avoids DtoH sync)
  2. Use torch.where instead of boolean mask indexing in the transformer

Both are pure optimizations — outputs should be identical.

Usage:
    # On the current branch (with PR changes):
    python test_zimage_pr13461.py --save current_branch.pt

    # On the main branch (without PR changes):
    git checkout main
    python test_zimage_pr13461.py --save main_branch.pt

    # Compare:
    python test_zimage_pr13461.py --compare current_branch.pt main_branch.pt
"""

import argparse
import torch


def run_pipeline():
    from diffusers import ZImagePipeline

    pipe = ZImagePipeline.from_pretrained(
        "Tongyi-MAI/Z-Image-Turbo",
        torch_dtype=torch.bfloat16,
    )
    pipe.to("cuda")

    generator = torch.Generator(device="cuda").manual_seed(0)

    # Use guidance_scale > 1 and cfg_truncation to exercise both changed code paths.
    # Small resolution + few steps + latent output for speed.
    output = pipe(
        prompt="a cat",
        height=256,
        width=256,
        num_inference_steps=2,
        guidance_scale=3.5,
        cfg_truncation=0.5,
        output_type="latent",
        generator=generator,
    )
    return output.images


def save(latents, path):
    torch.save(latents.cpu(), path)
    print(f"Saved latents with shape {latents.shape} and dtype {latents.dtype} to {path}")


def compare(path_a, path_b):
    a = torch.load(path_a, weights_only=True)
    b = torch.load(path_b, weights_only=True)

    print(f"Tensor A: shape={a.shape}, dtype={a.dtype}")
    print(f"Tensor B: shape={b.shape}, dtype={b.dtype}")

    if a.shape != b.shape:
        print("FAIL: shapes differ")
        return

    exact_match = torch.equal(a, b)
    max_diff = (a.float() - b.float()).abs().max().item()
    print(f"Exact match: {exact_match}")
    print(f"Max absolute difference: {max_diff}")

    if exact_match:
        print("PASS: outputs are identical")
    elif max_diff < 1e-3:
        print(f"PASS: outputs differ by at most {max_diff} (within tolerance)")
    else:
        print(f"FAIL: outputs differ by {max_diff}")


def main():
    parser = argparse.ArgumentParser(description="ZImage PR #13461 output comparison")
    parser.add_argument("--save", type=str, help="Run pipeline and save latents to this path")
    parser.add_argument("--compare", nargs=2, metavar=("A", "B"), help="Compare two saved latent files")
    args = parser.parse_args()

    if args.save:
        latents = run_pipeline()
        save(latents, args.save)
    elif args.compare:
        compare(args.compare[0], args.compare[1])
    else:
        parser.print_help()


if __name__ == "__main__":
    main()