#!/usr/bin/env python3
"""Sweep ControlNet strength for the Z-Image FPV depth pipeline and montage it.

With the rev-3 depth maps (full per-frame normalization, commit 018909b) the
control signal is now high-contrast, so the right strength almost certainly
shifted. This renders ONE open frame + ONE closed frame across a strength matrix
at a fixed seed/prompt/sampler, then stacks them next to their depth map so you
can read adherence (does the passage land where the black recess is? does a
closed wall stay a wall?) at a glance.

  python3 tools/comfy-spike/sweep_strength.py http://192.168.1.26:8188 \
      --seed-dir tools/out/atlas/7 --open r8_N --closed r0_N \
      --strengths 0.55,0.7,0.85,1.0

Output: tools/out/compare/sweep_<run>.png  (rows = frames, cols = depth + each strength)
Re-run after re-baking depth maps to compare pipeline revisions head to head.
"""
import argparse
import glob
import os
import subprocess
import sys

REPO = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
COMFY = os.path.join(REPO, "tools", "comfy-spike", "comfy.py")
VENV_PY = os.path.join(REPO, "tools", "pixelart", ".venv", "bin", "python")

# The final montage needs Pillow, which only lives in the pixelart venv; the
# renders themselves shell out to `python3 comfy.py` (stdlib only). If we were
# launched under a Python without PIL, re-exec under the venv so one invocation
# both renders and montages. (Without this the whole sweep runs, then dies at the
# last step — wasting every render.)
if __name__ == "__main__":
    try:
        import PIL  # noqa: F401
    except ImportError:
        # The venv's bin/python is a symlink back to the base interpreter, so a
        # realpath() comparison wrongly reports "already in the venv" and skips
        # the hop — what matters is sys.path (venv site-packages), not the binary.
        # Guard the re-exec with an env sentinel so a still-missing PIL raises
        # instead of looping forever.
        if os.path.exists(VENV_PY) and not os.environ.get("_SWEEP_REEXEC"):
            os.environ["_SWEEP_REEXEC"] = "1"
            os.execv(VENV_PY, [VENV_PY, os.path.abspath(__file__), *sys.argv[1:]])
        raise

# Mirror bake_atlas.py's proven prompt skeleton so the sweep reflects production.
STYLE = ("dark fantasy, ancient torchlit stone, warm orange torchlight, deep black shadows, "
         "wet stone, volumetric haze, architectural photography of an empty ruin, "
         "empty, deserted, no people, uninhabited, atmospheric, highly detailed, ")
OPEN_BODY = "grand pillared hall, rows of carved stone columns, a dark arched passage leads onward into shadow ahead"
WALL_BODY = ("dungeon gatehouse of carved stone, the way ahead ends at a blank unbroken wall of stone blocks, "
             "sealed solid masonry with no door and no opening, a dead end")
NEGATIVE = ("person, people, human, figure, character, hands, fingers, hand, arm, arms, "
            "holding weapon, sword, feet, legs, body, silhouette, UI, HUD, text, watermark, modern")


def render(url, depth, prompt, strength, seed, prefix, wait):
    ok = subprocess.run(["python3", COMFY, "zrender", url, "--depth", depth,
                         "--width", "1280", "--height", "720", "--strength", str(strength),
                         "--seed", str(seed), "--prompt", prompt, "--negative", NEGATIVE,
                         "--prefix", prefix, "--wait", str(wait)],
                        capture_output=True, text=True)
    if ok.returncode != 0:
        sys.stderr.write(ok.stdout + ok.stderr + "\n")
        return None
    hits = glob.glob(os.path.join(os.path.dirname(depth), f"out_{prefix}_*.png"))
    return hits[0] if hits else None


def montage(rows, strengths, out):
    from PIL import Image, ImageDraw
    cw, ch, pad = 360, 203, 6
    ncol = 1 + len(strengths)
    sheet = Image.new("RGB", (cw * ncol + pad * (ncol + 1), (ch + 20) * len(rows) + pad), (16, 16, 16))
    dr = ImageDraw.Draw(sheet)
    y = pad
    for label, depth, imgs in rows:
        cells = [("depth", depth)] + [(f"str {s}", imgs.get(s)) for s in strengths]
        x = pad
        for cap, path in cells:
            if path and os.path.exists(path):
                sheet.paste(Image.open(path).convert("RGB").resize((cw, ch)), (x, y + 18))
            dr.text((x, y + 4), f"{label} · {cap}", fill=(0, 230, 120))
            x += cw + pad
        y += ch + 20
    sheet.save(out)
    return out


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("url")
    ap.add_argument("--seed-dir", default=os.path.join(REPO, "tools", "out", "atlas", "7"))
    ap.add_argument("--alt-work", default="", help="optional 2nd depth dir to A/B (e.g. tools/out/compare/rev2_depths "
                    "= the old washed-out maps); adds an OLD row per frame so old-vs-new adherence is conclusive")
    ap.add_argument("--open", default="r8_N", help="frame id for the open-passage test")
    ap.add_argument("--closed", default="r0_N", help="frame id for the closed-wall test")
    ap.add_argument("--strengths", default="0.55,0.7,0.85,1.0")
    ap.add_argument("--seed", type=int, default=7, help="fixed diffusion seed (style lock)")
    ap.add_argument("--wait", type=int, default=300)
    a = ap.parse_args()

    work = os.path.join(a.seed_dir, "_work")
    strengths = [float(s) for s in a.strengths.split(",") if s.strip()]
    run = os.getpid()
    cases = [(a.open, OPEN_BODY), (a.closed, WALL_BODY)]
    # (source-tag, depth-dir): NEW always; OLD too when --alt-work is given.
    sources = [("NEW", work)] + ([("OLD", a.alt_work)] if a.alt_work else [])

    rows = []
    for frame, body in cases:
        prompt = STYLE + body
        for tag, srcdir in sources:
            depth = os.path.join(srcdir, f"{frame}_depth.png")
            if not os.path.exists(depth):
                sys.exit(f"!! missing depth map {depth} — bake the atlas first")
            imgs = {}
            for s in strengths:
                prefix = f"sweep_{frame}_{tag}_{int(s*100)}_{run}"
                print(f"[{frame}·{tag} @ strength {s}] rendering ...", flush=True)
                out = render(a.url, depth, prompt, s, a.seed, prefix, a.wait)
                if out:
                    imgs[s] = out
            rows.append((f"{frame}·{tag}", depth, imgs))

    out_dir = os.path.join(REPO, "tools", "out", "compare")
    os.makedirs(out_dir, exist_ok=True)
    out = montage(rows, strengths, os.path.join(out_dir, f"sweep_{run}.png"))
    # tidy the per-strength intermediates; the montage is the artifact
    for _, _, imgs in rows:
        for p in imgs.values():
            try:
                os.remove(p)
            except OSError:
                pass
    print(f"\nmontage → {os.path.relpath(out, REPO)}")


if __name__ == "__main__":
    main()