#!/usr/bin/env python3
"""v4 Phase 0 calibration: C1 redux mode, C2 resolution, C3 CN end, C4 bfl spec. Seed 777."""
import json, os, shutil, time, urllib.request
from run_sweeps import post, HOST, SEED

PAIRS = {
 "r1xt2": ("r1.jpg", "t2.jpg",
   "A photograph of a woman with long dark wavy hair seated on a pebble beach in a sheer black "
   "mini dress with gold necklaces, a turquoise cove and seaside buildings behind her, looking at "
   "the camera. Warm golden-hour sunlight, glistening sun-kissed skin, deep blue sea and dramatic "
   "clouds, crisp editorial color. Sharp focus, natural proportions."),
 "r2xt5": ("r2.jpg", "t5.jpg",
   "A photograph of a woman with long hair seated on the edge of a luxury pool deck at night, "
   "leaning back on one arm, wearing a white bikini, closed umbrellas and a cabana fading into "
   "darkness behind her. Hard direct camera flash, glossy high-contrast flash photography, "
   "saturated skin tones. Sharp focus, natural proportions."),
 "r3xt6": ("r3.jpg", "t6.jpg",
   "A photograph of a woman with long dark hair standing in a bright hotel bedroom beside a tall "
   "window, wearing a delicate lace bikini set, a bed with flowers and a nightstand behind her. "
   "Soft diffuse natural daylight, oiled bronze skin, warm earthy organic palette. Sharp focus, "
   "natural proportions."),
}

def graph(pair, tag, *, mode="mult", redux=0.30, w=896, h=1152, cn_str=0.7, cn_end=0.8,
          bfl=False, guidance=None, lora=1.0):
    style, comp, prompt = PAIRS[pair]
    g = guidance if guidance is not None else (10.0 if bfl else 3.0)
    p = {}
    p["u"] = {"class_type": "UNETLoader", "inputs": {"unet_name": "flux1-dev.safetensors", "weight_dtype": "default"}}
    p["c"] = {"class_type": "DualCLIPLoader", "inputs": {"clip_name1": "t5xxl_fp16.safetensors",
        "clip_name2": "clip_l.safetensors", "type": "flux", "device": "default"}}
    p["v"] = {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}}
    p["txt"] = {"class_type": "CLIPTextEncode", "inputs": {"clip": ["c", 0], "text": prompt}}
    p["guid"] = {"class_type": "FluxGuidance", "inputs": {"conditioning": ["txt", 0], "guidance": g}}
    p["neg"] = {"class_type": "ConditioningZeroOut", "inputs": {"conditioning": ["guid", 0]}}
    p["si"] = {"class_type": "LoadImage", "inputs": {"image": style}}
    p["cvl"] = {"class_type": "CLIPVisionLoader", "inputs": {"clip_name": "sigclip_vision_patch14_384.safetensors"}}
    p["sml"] = {"class_type": "StyleModelLoader", "inputs": {"style_model_name": "flux1-redux-dev.safetensors"}}
    if mode == "radv":
        p["sma"] = {"class_type": "ReduxAdvanced", "inputs": {"conditioning": ["guid", 0],
            "style_model": ["sml", 0], "clip_vision": ["cvl", 0], "image": ["si", 0],
            "downsampling_factor": 3, "downsampling_function": "area",
            "mode": "center crop (square)", "weight": 1.0, "autocrop_margin": 0.1}}
    else:
        p["cve"] = {"class_type": "CLIPVisionEncode", "inputs": {"clip_vision": ["cvl", 0], "image": ["si", 0], "crop": "center"}}
        p["sma"] = {"class_type": "StyleModelApply", "inputs": {"conditioning": ["guid", 0],
            "style_model": ["sml", 0], "clip_vision_output": ["cve", 0], "strength": redux, "strength_type": "multiply"}}
    p["ci"] = {"class_type": "LoadImage", "inputs": {"image": comp}}
    p["rs"] = {"class_type": "ImageResize+", "inputs": {"image": ["ci", 0], "width": w, "height": h,
        "interpolation": "lanczos", "method": "fill / crop", "condition": "always", "multiple_of": 0}}
    p["depth"] = {"class_type": "DepthAnythingV2Preprocessor", "inputs": {
        "image": ["rs", 0], "ckpt_name": "depth_anything_v2_vitl.pth", "resolution": w}}
    if bfl:
        p["lora"] = {"class_type": "LoraLoaderModelOnly", "inputs": {"model": ["u", 0],
            "lora_name": "flux1-depth-dev-lora.safetensors", "strength_model": lora}}
        p["ip2p"] = {"class_type": "InstructPixToPixConditioning", "inputs": {"positive": ["sma", 0],
            "negative": ["neg", 0], "vae": ["v", 0], "pixels": ["depth", 0]}}
        model_in, pos, negs, lat = ["lora", 0], ["ip2p", 0], ["ip2p", 1], ["ip2p", 2]
    else:
        p["cnl"] = {"class_type": "ControlNetLoader", "inputs": {"control_net_name": "FLUX.1-dev-ControlNet-Union-Pro-2.0.safetensors"}}
        p["cn"] = {"class_type": "ControlNetApplySD3", "inputs": {"positive": ["sma", 0], "negative": ["neg", 0],
            "control_net": ["cnl", 0], "vae": ["v", 0], "image": ["depth", 0],
            "strength": cn_str, "start_percent": 0.0, "end_percent": cn_end}}
        p["lat0"] = {"class_type": "EmptySD3LatentImage", "inputs": {"width": w, "height": h, "batch_size": 1}}
        model_in, pos, negs, lat = ["u", 0], ["cn", 0], ["cn", 1], ["lat0", 0]
    p["msf"] = {"class_type": "ModelSamplingFlux", "inputs": {"model": model_in,
        "max_shift": 1.15, "base_shift": 0.5, "width": w, "height": h}}
    p["ks"] = {"class_type": "KSampler", "inputs": {"model": ["msf", 0], "positive": pos, "negative": negs,
        "latent_image": lat, "seed": SEED, "steps": 32, "cfg": 1.0,
        "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0}}
    p["dec"] = {"class_type": "VAEDecode", "inputs": {"samples": ["ks", 0], "vae": ["v", 0]}}
    p["save"] = {"class_type": "SaveImage", "inputs": {"images": ["dec", 0], "filename_prefix": f"calib/{tag}"}}
    return p

q = {}
for pair in PAIRS:
    q[f"C1-mult25-{pair}"] = graph(pair, f"C1-mult25-{pair}", redux=0.25)
    q[f"C1-mult35-{pair}"] = graph(pair, f"C1-mult35-{pair}", redux=0.35)
    q[f"C1-radv-{pair}"]  = graph(pair, f"C1-radv-{pair}", mode="radv")
    q[f"C2-896-{pair}"]   = graph(pair, f"C2-896-{pair}")                       # mult .30 @896x1152 (also C3 end .8 baseline)
    q[f"C2-768-{pair}"]   = graph(pair, f"C2-768-{pair}", w=768, h=1024)
    q[f"C3-end4-{pair}"]  = graph(pair, f"C3-end4-{pair}", cn_end=0.4)
    q[f"C3-end6-{pair}"]  = graph(pair, f"C3-end6-{pair}", cn_end=0.6)
    q[f"C3-s8e6-{pair}"]  = graph(pair, f"C3-s8e6-{pair}", cn_str=0.8, cn_end=0.6)
    q[f"C4-spec-{pair}"]  = graph(pair, f"C4-spec-{pair}", bfl=True, guidance=10.0, lora=0.85)
    q[f"C4-alt-{pair}"]   = graph(pair, f"C4-alt-{pair}", bfl=True, guidance=4.0, lora=1.0)

ids = {post(p): tag for tag, p in q.items()}
print("queued", len(ids), flush=True)
pending = set(ids)
errors = []
while pending:
    time.sleep(12)
    for pid in list(pending):
        try:
            with urllib.request.urlopen(f"{HOST}/history/{pid}") as r:
                h = json.loads(r.read())
        except Exception: continue
        if pid not in h: continue
        st = h[pid].get("status", {})
        if st.get("completed"):
            pending.discard(pid); print(f"done {ids[pid]} ({len(ids)-len(pending)}/{len(ids)})", flush=True)
        elif st.get("status_str") == "error":
            pending.discard(pid); errors.append(ids[pid])
            msgs = [m for m in st.get("messages", []) if m[0]=="execution_error"]
            print("ERROR", ids[pid], (msgs[-1][1].get("exception_message","?") if msgs else "?")[:200], flush=True)
os.makedirs("/workspace/outputs_v4/calib", exist_ok=True)
src = "/workspace/ComfyUI/output/calib"
for f in sorted(os.listdir(src)):
    shutil.copy(os.path.join(src, f), f"/workspace/outputs_v4/calib/{f.split('_0')[0]}.png")
print("ERRORS:", errors, flush=True)
print("CALIB COMPLETE", flush=True)