"""Generate P7-5.1 style-reference candidates with Qwen Image.

Outputs are generation records. This script does not assign approval states.
Use ``--scene`` to regenerate one named existing or extension scene.
The default run covers all twenty contract rows. Outputs are candidates only;
human observation records describe each image's reference role and limitations.
Each saved PNG receives a neighboring ``-result.json`` generation record.
"""

import argparse
import json
import subprocess
import sys
import threading
import time
from pathlib import Path

ASSET_DIR = Path(__file__).resolve().parent
MODEL_ID = "Qwen/Qwen-Image"
TRANSFORMER_REPOSITORY = "nunchaku-tech/nunchaku-qwen-image"
TRANSFORMER_FILENAME = "svdq-fp4_r128-qwen-image.safetensors"
TRANSFORMER_ID = f"{TRANSFORMER_REPOSITORY}/{TRANSFORMER_FILENAME}"
HF_HUB_CACHE = ASSET_DIR.parents[4] / ".tmp" / "download" / "huggingface" / "hub"
# All ordinary style-reference runs use a square 1024 canvas.  A comparison
# experiment may override this explicitly with --width/--height.
SIZE = (1024, 1024)
# Default quality/throughput operating point for subsequent style candidates.
# The 4-step screen remains a recorded performance probe, not a style-master
# setting; each run prints any explicit override in its terminal summary.
STEPS = 30
TRUE_CFG_SCALE = 4.0
STYLE_PROMPT_PATH = ASSET_DIR / "p7-5-1-style-prompt-contract.json"
COMMON_CONTRACT = json.loads(STYLE_PROMPT_PATH.read_text(encoding="utf-8"))["common_contract"]
SCENES = [
    {
        "id": "atrium-dawn-high-angle",
        "generate_by_default": False,
        "seed": 420713,
        "prompt": "Vertical empty indoor atrium at early dawn, steep high-angle from an upper landing. One broad straight concrete stair descends diagonally beside a solid wall to an open tiled floor with one bench and two plants. No exposed railings, balusters, or floating metal lines. Cool off-white tiles, blue-gray daylight, teal wall shadows, and one muted-apricot reflection. ",
    },
    {
        "id": "courtyard-early-morning-high-angle",
        "generate_by_default": False,
        "seed": 420702,
        "prompt": "Vertical empty Seoul residential courtyard in clear early morning, steep high-angle from a balcony. Diagonal paving, one tree, benches, planters, and low roofs cross the edges. Cool off-white paving, blue-teal shadows, leaf green, pale-blue sky reflection, and one tiny warm window glint. ",
    },
    {
        "id": "downtown-clear-day-wide",
        "generate_by_default": False,
        "seed": 420703,
        "prompt": "Vertical empty Seoul business intersection at clear midday, side view from a shaded sidewalk corner. Diagonal glass towers and street trees, cool off-white pavement; avoid a centered road corridor. Pale-blue glass reflections, teal building shadows, green foliage shadows, cool-gray highlights. ",
    },
    {
        "id": "residential-sunset-low-angle",
        "generate_by_default": False,
        "seed": 420704,
        "prompt": "Vertical empty residential street at sunset, low angle from curb height toward bicycle rack, house facades, branches, and narrow sky. Strong foreground-to-sky scale; avoid eye level and a centered corridor. Teal pavement shadow, olive foliage, blue-gray walls, narrow muted-apricot sky rim. ",
    },
    {
        "id": "night-lit-reading-room-oblique",
        "generate_by_default": False,
        "seed": 420705,
        "prompt": "Vertical empty reading room at deep night, oblique view of one square wood desk beside a tall window. One small shaded lamp lights only the desktop and a short patch of floor; the rest is a simple blue-gray wall and dark wood floor. Indigo night outside, compact amber reflection; avoid long corridor lines and extra furniture. ",
    },
    {
        "id": "rooftop-rainy-night-overhead",
        "generate_by_default": False,
        "seed": 420706,
        "prompt": "Vertical empty Seoul rooftop plaza after rain at late night, steep overhead from a high terrace. Wet paving, two planters, and one shallow puddle form diagonal planes. Deep indigo shadow, navy pavement, cyan puddle reflections, and a few small tungsten reflections; no rain streaks or warm sky. ",
    },
    {
        "id": "venice-sunset-oblique",
        "generate_by_default": False,
        "seed": 420707,
        "prompt": "Vertical empty Venice canal at sunset, oblique view from a stone bridge edge. A diagonal canal bends between pale ochre facades; avoid a centered canal. Medium teal water, small indigo water shadows, pale stone reflection, and a narrow muted-apricot sky opening. ",
    },
    {
        "id": "park-clear-day-eye-level",
        "generate_by_default": False,
        "seed": 420708,
        "prompt": "Vertical empty city park pond at clear midday, calm eye-level diagonal from a cool off-white path. Diagonal water edge, teal pond reflection, leaf green, pale-blue sky reflection, and blue-green tree shade; avoid a centered path corridor. ",
    },
    {
        "id": "train-platform-rainy-night-oblique",
        "generate_by_default": False,
        "seed": 420709,
        "prompt": "Vertical empty open-air Seoul train platform after rain at late night, oblique view under a simple canopy. Platform edge, columns, blank benches, and two rail lines recede diagonally. Indigo wet pavement, navy shadows, cool-white canopy light, small tungsten pools, and cyan puddle reflections. ",
    },
    {
        "id": "gallery-midday-oblique", "generate_by_default": True, "seed": 420810,
        "prompt": "Vertical empty contemporary gallery at clear midday, oblique view from a near corner. Off-white walls, cool-gray floor, blank plinths, and ceiling tracks form diagonal planes crossing the edges. Pale-blue reflection, blue-gray shadow, and one muted-apricot accent; avoid artworks, visitors, labels, and a centered corridor. ",
    },
    {
        "id": "library-stairwell-day-high-angle", "generate_by_default": True, "seed": 420811,
        "prompt": "Daylit public-library interior with a tall high-ceiling lobby, high oblique view. Tall warm-wood bookcase against one wall, frosted window, cool off-white floors, pale-cyan daylight reflection, blue-teal shadow. ",
    },
    {
        "id": "harbor-plaza-sunrise-high", "generate_by_default": True, "seed": 420812,
        "prompt": "Vertical empty harbor terrace at sunrise, high oblique view across one broad cool-off-white paved foreground to a low straight seawall and open water. Use only three dark mooring bollards along the seawall; no plants or boats. Blue-teal water, indigo paving shadow, and a narrow muted-apricot horizon; avoid people, signs, broad orange light, and a centered waterfront. ",
    },
    {
        "id": "underpass-rainy-twilight", "generate_by_default": True, "seed": 420813,
        "prompt": "Vertical empty pedestrian underpass just after rain at blue twilight, diagonal view from its entrance into a gently bending passage. Concrete walls, wet tile, blank columns, and a narrow cool-sky opening cross the edges. Indigo wet shadows, blue-gray concrete, cyan puddles, and small warm safety lights; avoid people, graffiti, signs, trains, and a centered tunnel. ",
    },
    {
        "id": "hillside-alley-late-afternoon", "generate_by_default": True, "seed": 420814,
        "prompt": "Vertical empty hillside alley in late afternoon, eye-level view along a diagonal climbing path. Retaining walls, blank small-house facades, unmarked poles, steps, and foliage overlap at the edges. Teal pavement shadows, leaf-green foliage, blue-gray walls, and a narrow muted-apricot rim; avoid people, vehicles, signs, broad sunset orange, and a centered corridor. ",
    },
    {
        "id": "market-arcade-overcast", "generate_by_default": True, "seed": 420815,
        "prompt": "Vertical empty covered market arcade under soft overcast daylight, oblique view across shuttered blank stalls. Canopy ribs, damp cool-gray floor, unmarked crates, and side openings create diagonal depth at the edges. Cool off-white skylight, blue-teal shadow, muted olive, and faint cyan reflections; avoid shoppers, products, signs, logos, and a centered corridor. ",
    },
    {
        "id": "riverside-terrace-night", "generate_by_default": True, "seed": 420816,
        "prompt": "Vertical empty riverside terrace at night, oblique view beside a low stone planter. Broad promenade, river edge, blank benches, distant bridge silhouette, and sparse trees cross the edges. Navy pavement, indigo water, cyan-blue reflections, cool-white lights, and two muted tungsten pools; avoid people, boats, signs, neon, and a centered corridor. ",
    },
    {
        "id": "greenhouse-blue-hour", "generate_by_default": True, "seed": 420817,
        "prompt": "Vertical empty greenhouse conservatory at blue hour, quiet eye-level diagonal. Glass roof ribs, damp stone path, leafy plants, benches, and a distant glass door cross the edges. Pale-blue exterior light, blue-teal glass shadow, leaf green, cool off-white highlights, and tiny warm glints; avoid people, labels, signs, animals, and a centered aisle. ",
    },
    {
        "id": "ferry-deck-morning", "generate_by_default": True, "seed": 420818,
        "prompt": "Vertical empty open ferry deck in clear morning light, oblique view beside a blank bench toward railings and distant water. Deck planks, simple rail posts, unmarked life-ring housing, and horizon cross the edges. Cool off-white deck light, teal sea, pale-blue reflections, navy rail shadow, muted-apricot accent; avoid people, boats, text, logos, and a centered corridor. ",
    },
    {
        "id": "cinema-foyer-night", "generate_by_default": True, "seed": 420819,
        "prompt": "Vertical empty neighborhood cinema foyer at night, eye-level view from a side corner. Dark-indigo tiles, blank ticket counter planes, unlettered poster frames, ceiling lights, and glass reflections create an oblique composition at the edges. Deep navy shadow, cool-white light, cyan reflections, restrained amber pools; avoid people, film images, signs, logos, and a centered hallway. ",
    },
    {
        "id": "ceramics-studio-afternoon", "generate_by_default": True, "seed": 420820,
        "prompt": "Vertical empty ceramics studio in quiet afternoon light, diagonal view across a worktable. Pottery wheel, blank shelves, unmarked clay forms, tall windows, and cool concrete floor cross the edges. Pale-blue window light, blue-gray shadow, muted clay beige, leaf-green reflection, and a small apricot highlight; avoid people, lettering, logos, and a centered aisle. ",
    },
]


def gpu_memory_mib() -> int:
    result = subprocess.run(
        ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
        check=True,
        capture_output=True,
        text=True,
    )
    return int(result.stdout.splitlines()[0])


def parse_args(argv=None):
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    scene_ids = [scene["id"] for scene in SCENES]
    parser.add_argument("--scene", choices=scene_ids, help="Generate one named scene.")
    parser.add_argument("--exclude", nargs="+", choices=scene_ids, default=[],
                        help="Scene IDs to exclude, separated by spaces.")
    parser.add_argument("--include-existing", action=argparse.BooleanOptionalAction,
                        default=True, help="Include existing rows when no --scene is given.")
    parser.add_argument("--run-label", default="v1", help="Label used in output filenames.")
    parser.add_argument("--steps", type=int, default=STEPS, help="Diffusion steps.")
    parser.add_argument("--width", type=int, default=SIZE[0], help="Output width in pixels.")
    parser.add_argument("--height", type=int, default=SIZE[1], help="Output height in pixels.")
    args = parser.parse_args(argv)
    if args.steps < 1:
        parser.error("--steps must be positive")
    if any(value <= 0 or value % 16 for value in (args.width, args.height)):
        parser.error("--width and --height must be positive multiples of 16")
    if not args.run_label or any(c not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for c in args.run_label):
        parser.error("--run-label must contain only letters, digits, underscores or hyphens")
    args.selected_scenes = [scene for scene in SCENES if
        (scene["id"] == args.scene if args.scene else
         args.include_existing or scene["generate_by_default"])
        and scene["id"] not in args.exclude]
    if not args.selected_scenes:
        parser.error("No scenes selected for generation")
    return args


def main() -> None:
    args = parse_args()
    requested_scene = args.scene
    run_label = args.run_label
    include_existing = args.include_existing
    steps = args.steps
    size = (args.width, args.height)
    excluded_scenes = set(args.exclude)
    scenes = args.selected_scenes

    import torch
    from diffusers import QwenImagePipeline
    from huggingface_hub import snapshot_download
    from nunchaku import NunchakuQwenImageTransformer2DModel
    # The shared naming helper remains in the chapter asset directory.
    sys.path.insert(0, str(ASSET_DIR.parent))
    from p7_5_image_output_naming import candidate_stem

    before = gpu_memory_mib()
    peak = before
    stop = threading.Event()

    def observe_peak() -> None:
        nonlocal peak
        while not stop.is_set():
            peak = max(peak, gpu_memory_mib())
            time.sleep(0.2)

    observer = threading.Thread(target=observe_peak, daemon=True)
    observer.start()
    started = time.monotonic()
    runs = []
    try:
        transformer_path = Path(
            snapshot_download(TRANSFORMER_REPOSITORY, cache_dir=HF_HUB_CACHE, local_files_only=True)
        ) / TRANSFORMER_FILENAME
        model_path = Path(snapshot_download(MODEL_ID, cache_dir=HF_HUB_CACHE, local_files_only=True))
        transformer = NunchakuQwenImageTransformer2DModel.from_pretrained(transformer_path)
        transformer.set_offload(True, use_pin_memory=False, num_blocks_on_gpu=1)
        pipe = QwenImagePipeline.from_pretrained(
            model_path, transformer=transformer, torch_dtype=torch.bfloat16, local_files_only=True
        )
        pipe._exclude_from_cpu_offload.append("transformer")
        pipe.enable_sequential_cpu_offload()
        for scene in scenes:
            scene_started = time.monotonic()
            prompt = scene["prompt"] + COMMON_CONTRACT
            run_contract = {
                "model": MODEL_ID,
                "transformer": TRANSFORMER_ID,
                "prompt": prompt,
                "size": size,
                "true_cfg_scale": TRUE_CFG_SCALE,
            }
            image = pipe(
                prompt=prompt,
                width=size[0],
                height=size[1],
                num_inference_steps=steps,
                true_cfg_scale=TRUE_CFG_SCALE,
                negative_prompt=" ",
                generator=torch.Generator(device="cpu").manual_seed(scene["seed"]),
            ).images[0]
            output_stem = candidate_stem(
                f"p7-5-1-style-{scene['id']}-qwen-image-{run_label}",
                seed=scene["seed"],
                steps=steps,
                contract=run_contract,
            )
            image_name = f"{output_stem}.png"
            result_name = f"{output_stem}-result.json"
            image.save(ASSET_DIR / image_name)
            elapsed_seconds = round(time.monotonic() - scene_started, 1)
            result = {
                "status": "generated",
                "record_kind": "generation_result",
                "scene_id": scene["id"],
                "model": MODEL_ID,
                "transformer": TRANSFORMER_ID,
                "runtime": "local GPU via Diffusers QwenImagePipeline with Nunchaku FP4 r128 sequential CPU offload",
                "prompt": prompt,
                "seed": scene["seed"],
                "size": list(size),
                "steps": steps,
                "true_cfg_scale": TRUE_CFG_SCALE,
                "output": {"path": image_name},
                "elapsed_seconds": elapsed_seconds,
            }
            (ASSET_DIR / result_name).write_text(
                json.dumps(result, ensure_ascii=False, indent=2) + "\n",
                encoding="utf-8",
            )
            runs.append(
                {
                    "id": scene["id"],
                    "seed": scene["seed"],
                    "prompt": prompt,
                    "prompt_word_count": len(prompt.split()),
                    "asset": image_name,
                    "result": result_name,
                    "elapsed_seconds": elapsed_seconds,
                    "record_kind": "generation_result",
                }
            )
            torch.cuda.empty_cache()
    finally:
        stop.set()
        observer.join(timeout=2)

    record = {
        "record_kind": "generation_record",
        "model_id": MODEL_ID,
        "transformer_id": TRANSFORMER_ID,
        "runtime": "local GPU via Diffusers QwenImagePipeline with Nunchaku FP4 r128 sequential CPU offload",
        "size": list(size),
        "steps": steps,
        "true_cfg_scale": TRUE_CFG_SCALE,
        "elapsed_seconds": round(time.monotonic() - started, 1),
        "gpu_memory_before_mib": before,
        "gpu_memory_peak_mib": peak,
        "requested_scene": requested_scene,
        "include_existing": include_existing,
        "excluded_scenes": sorted(excluded_scenes),
        "runs": runs,
    }
    print(json.dumps(record, indent=2))


if __name__ == "__main__":
    main()