"""
Mass Iteration Studio — one image in, N variants out (PNG only).

Built for ZeroGPU: work is split into short GPU calls so a large run never
exceeds the per-call duration limit. Alpha channels survive the round trip.
"""

import os
import csv
import math
import random
import tempfile
import zipfile
from datetime import datetime

import gradio as gr
import numpy as np
import spaces
import torch
from PIL import Image, ImageOps
from diffusers import AutoPipelineForImage2Image, EulerDiscreteScheduler
from huggingface_hub import hf_hub_download

# --------------------------------------------------------------------------
# Models
# --------------------------------------------------------------------------
MODELS = {
    "SDXL · 20 steps · strongest prompt": {
        "base": "stabilityai/stable-diffusion-xl-base-1.0",
        "lora": None,
        "steps": 20, "guidance": 7.0, "size": 1024, "trailing": False,
        "sec_per_image": 2.5,
    },
    "SDXL-Lightning · 4 steps · fast": {
        "base": "stabilityai/stable-diffusion-xl-base-1.0",
        "lora": ("ByteDance/SDXL-Lightning", "sdxl_lightning_4step_lora.safetensors"),
        "steps": 4, "guidance": 1.0, "size": 1024, "trailing": True,
        "sec_per_image": 0.8,
    },
    "SDXL-Turbo · 2 steps · 768px": {
        "base": "stabilityai/sdxl-turbo", "lora": None,
        "steps": 2, "guidance": 0.0, "size": 768, "trailing": False,
        "sec_per_image": 0.45,
    },
    "SD-Turbo · 2 steps · 512px": {
        "base": "stabilityai/sd-turbo", "lora": None,
        "steps": 2, "guidance": 0.0, "size": 512, "trailing": False,
        "sec_per_image": 0.2,
    },
}
DEFAULT_MODEL = "SDXL · 20 steps · strongest prompt"

DTYPE = torch.float16
_pipes = {}


def get_pipe(name: str):
    """Load once, keep in memory. First call downloads several GB."""
    if name in _pipes:
        return _pipes[name]

    cfg = MODELS[name]
    pipe = AutoPipelineForImage2Image.from_pretrained(
        cfg["base"], torch_dtype=DTYPE, variant="fp16", use_safetensors=True
    )
    if cfg["lora"]:
        repo, ckpt = cfg["lora"]
        pipe.load_lora_weights(hf_hub_download(repo, ckpt))
        pipe.fuse_lora()
    if cfg["trailing"]:
        pipe.scheduler = EulerDiscreteScheduler.from_config(
            pipe.scheduler.config, timestep_spacing="trailing"
        )
    pipe.set_progress_bar_config(disable=True)
    pipe.to("cuda")
    _pipes[name] = pipe
    return pipe


# --------------------------------------------------------------------------
# Color handling
# --------------------------------------------------------------------------
COLOR_NAMES = {
    "#000000": "black", "#1a1a1a": "near black", "#3d3d3d": "charcoal",
    "#6b6b6b": "grey", "#a8a8a8": "light grey", "#e8e4dc": "bone white",
    "#ffffff": "white", "#7f1d1d": "oxblood", "#dc2626": "bright red",
    "#f87171": "soft coral", "#c2410c": "burnt orange", "#e8794a": "terracotta",
    "#fb923c": "warm orange", "#f59e0b": "amber", "#facc15": "yellow",
    "#eab308": "brass gold", "#a3a635": "olive", "#65a30d": "leaf green",
    "#16a34a": "emerald green", "#22c55e": "bright green", "#a3e635": "acid green",
    "#0d9488": "deep teal", "#2dd4bf": "turquoise", "#06b6d4": "cyan",
    "#0284c7": "azure blue", "#1e3a8a": "navy blue", "#4338ca": "indigo",
    "#6366f1": "periwinkle", "#7c3aed": "violet", "#a78bfa": "lilac",
    "#c026d3": "magenta", "#ec4899": "hot pink", "#f9a8d4": "pastel pink",
    "#78350f": "chocolate brown", "#b45309": "clay", "#d6c6a8": "sand",
    "#94a3b8": "slate", "#334155": "graphite", "#f5f5dc": "cream",
}
_NAME_RGB = {h: tuple(int(h[i:i + 2], 16) for i in (1, 3, 5)) for h in COLOR_NAMES}


def to_rgb(value, fallback=(255, 255, 255)):
    """Accept #rrggbb or rgba(r, g, b, a) — Gradio's picker returns both."""
    if not value:
        return fallback
    v = str(value).strip()
    if v.startswith("#") and len(v) >= 7:
        return tuple(int(v[i:i + 2], 16) for i in (1, 3, 5))
    if v.startswith("rgb"):
        nums = [float(x) for x in v[v.find("(") + 1:v.find(")")].split(",")[:3]]
        return tuple(int(round(n)) for n in nums)
    return fallback


def name_color(value) -> str:
    """Nearest plain-language name — models respond to words, not hex codes."""
    r, g, b = to_rgb(value)
    best = min(_NAME_RGB.items(),
               key=lambda kv: (kv[1][0] - r) ** 2 + (kv[1][1] - g) ** 2 + (kv[1][2] - b) ** 2)
    return COLOR_NAMES[best[0]]


def palette_lut(colors):
    """256×3 lookup mapping luminance onto the picked colors, dark to light."""
    rgbs = sorted((to_rgb(c) for c in colors), key=lambda c: 0.299 * c[0] + 0.587 * c[1] + 0.114 * c[2])
    stops = np.linspace(0, 255, len(rgbs))
    xs = np.arange(256)
    lut = np.stack([np.interp(xs, stops, [c[i] for c in rgbs]) for i in range(3)], axis=1)
    return lut.astype(np.uint8)


def force_palette(img: Image.Image, colors, mix: float) -> Image.Image:
    """Gradient-map the result onto the exact picked colors."""
    alpha = img.split()[3] if img.mode == "RGBA" else None
    lum = np.asarray(img.convert("L"))
    mapped = palette_lut(colors)[lum]
    if mix < 1.0:
        orig = np.asarray(img.convert("RGB")).astype(np.float32)
        mapped = (mapped * mix + orig * (1 - mix)).astype(np.uint8)
    out = Image.fromarray(mapped, "RGB")
    if alpha is not None:
        out = out.convert("RGBA")
        out.putalpha(alpha)
    return out


def keyed_alpha(img: Image.Image, matte, tol: int = 38, soft: int = 26) -> Image.Image:
    """Rebuild transparency from the matte color, so newly invented shapes
    keep their own silhouette instead of being clipped to the old one."""
    rgb = np.asarray(img.convert("RGB")).astype(np.int16)
    m = np.array(to_rgb(matte), dtype=np.int16)
    dist = np.sqrt(((rgb - m) ** 2).sum(axis=-1))
    a = np.clip((dist - tol) / max(soft, 1), 0.0, 1.0)
    out = img.convert("RGBA")
    out.putalpha(Image.fromarray((a * 255).astype(np.uint8), "L"))
    return out


# --------------------------------------------------------------------------
# Prompt pools
# --------------------------------------------------------------------------
POOL_COLOR = """warm terracotta and bone white
cold teal and graphite
monochrome charcoal on paper
acid green on deep black
dusty lilac and sand
indigo with brass gold
faded coral and sea foam
oxblood red and cream
electric cyan and magenta
muted olive and clay"""

POOL_STYLE = """flat vector illustration
risograph print with grain
thick ink outlines, cel shaded
soft airbrush gradients
woodcut engraving
matte gouache painting
chrome and glass render
halftone comic print
minimal bauhaus poster
chalk on blackboard"""

POOL_SHAPE = """rounded organic shapes
sharp angular geometry
elongated slender proportions
chunky bold silhouettes
fragmented and shattered forms
symmetrical and centered
loose hand drawn contours
tightly packed dense composition"""

NEGATIVE = "blurry, low quality, watermark, jpeg artifacts, deformed, garbled text"


def as_list(text: str):
    return [line.strip() for line in text.splitlines() if line.strip()]


def build_recipes(n, instruction, color_phrase, colors, styles, shapes,
                  strat, seed, smin, smax):
    """One recipe per variant: prompt, strength, seed."""
    rng = random.Random(seed)
    pools = [p for p in (colors, styles, shapes) if p]
    combos = []

    if strat.startswith("Grid") and pools:
        total = 1
        for p in pools:
            total *= len(p)
        for i in range(n):
            idx, parts = i % total, []
            for p in reversed(pools):
                parts.append(p[idx % len(p)])
                idx //= len(p)
            combos.append(list(reversed(parts)))
    else:
        for _ in range(n):
            combos.append([rng.choice(p) for p in pools])

    # snap strength onto a few levels so same-level variants batch together
    levels = 6
    step = (smax - smin) / (levels - 1) if smax > smin else 0.0

    recipes = []
    for i, parts in enumerate(combos):
        bits = []
        if instruction.strip():
            bits.append(instruction.strip())
        if color_phrase:
            bits.append(color_phrase)
        bits += parts
        strength = smin + round((rng.uniform(smin, smax) - smin) / step) * step if step else smin
        recipes.append({
            "index": i + 1,
            "prompt": ", ".join(bits) or "graphic design variation",
            "strength": round(strength, 2),
            "seed": seed + i,
        })
    return recipes


# --------------------------------------------------------------------------
# Image preparation — alpha in, alpha out
# --------------------------------------------------------------------------
def prepare(img: Image.Image, target: int, matte):
    """Composite onto a matte for the model, keep the mask for the way back."""
    img = ImageOps.exif_transpose(img)
    rgba = img.convert("RGBA")
    w, h = rgba.size
    s = target / max(w, h)
    size = (max(64, int(w * s) // 8 * 8), max(64, int(h * s) // 8 * 8))
    rgba = rgba.resize(size, Image.LANCZOS)

    alpha = rgba.split()[3]
    bg = Image.new("RGBA", size, to_rgb(matte) + (255,))
    flat = Image.alpha_composite(bg, rgba).convert("RGB")
    transparent = alpha.getextrema()[0] < 250
    return flat, alpha, transparent


# --------------------------------------------------------------------------
# GPU work
# --------------------------------------------------------------------------
@spaces.GPU(duration=90)
def render_batch(model_name, image, recipes, negative, guidance):
    cfg = MODELS[model_name]
    pipe = get_pipe(model_name)

    prompts = [r["prompt"] for r in recipes]
    strength = float(recipes[0]["strength"])
    guidance = float(guidance)
    # img2img only runs steps*strength of them, so scale up to keep the
    # model's native step count — otherwise low strength means no denoising
    steps = min(cfg["steps"] * 2,
                max(cfg["steps"], math.ceil(cfg["steps"] / max(strength, 0.15))))
    gens = [torch.Generator("cuda").manual_seed(r["seed"]) for r in recipes]

    out = pipe(
        prompt=prompts,
        negative_prompt=[negative] * len(prompts) if guidance > 1.0 else None,
        image=[image] * len(prompts),
        strength=strength,
        num_inference_steps=steps,
        guidance_scale=guidance,
        generator=gens,
    )
    return out.images


# --------------------------------------------------------------------------
# Orchestration
# --------------------------------------------------------------------------
def generate(image, model_name, count, instruction, guidance, use_picked, c1, c2, c3,
             apply_mode, palette_mix, alpha_mode, matte,
             colors_raw, styles_raw, shapes_raw, strat, smin, smax,
             seed, randomize, batch_size, negative, progress=gr.Progress()):

    if image is None:
        raise gr.Error("Upload an image first.")
    if smax < smin:
        smin, smax = smax, smin

    count = int(count)
    seed = random.randint(0, 2**31 - 1) if randomize else int(seed)
    cfg = MODELS[model_name]

    src, alpha, had_alpha = prepare(image, cfg["size"], matte)
    picked = [c1, c2, c3]

    color_phrase, pool_colors = "", as_list(colors_raw)
    if use_picked:
        names = [name_color(c) for c in picked]
        color_phrase = f"color palette of {names[0]}, {names[1]} and {names[2]}"
        if apply_mode != "Prompt + exact remap":
            pool_colors = []          # picked colors replace the pool
    recolor = use_picked and apply_mode in ("Exact remap only", "Prompt + exact remap")

    recipes = build_recipes(count, instruction, color_phrase, pool_colors,
                            as_list(styles_raw), as_list(shapes_raw),
                            strat, seed, float(smin), float(smax))
    recipes.sort(key=lambda r: r["strength"])

    run_dir = os.path.join(tempfile.gettempdir(), f"run_{datetime.now():%H%M%S}_{seed}")
    os.makedirs(run_dir, exist_ok=True)

    # group into batches that share one denoising schedule
    batches, i, bs = [], 0, int(batch_size)
    while i < len(recipes):
        chunk = [recipes[i]]
        i += 1
        while i < len(recipes) and len(chunk) < bs and \
                recipes[i]["strength"] == chunk[0]["strength"]:
            chunk.append(recipes[i])
            i += 1
        batches.append(chunk)

    gallery, done = [], 0
    for chunk in batches:
        images = render_batch(model_name, src, chunk, negative, guidance)
        for r, img in zip(chunk, images):
            if alpha_mode.startswith("Reuse") and had_alpha:
                img = img.convert("RGBA")
                img.putalpha(alpha)
            elif alpha_mode.startswith("Key"):
                img = keyed_alpha(img, matte)
            if recolor:
                img = force_palette(img, picked, float(palette_mix))
            path = os.path.join(run_dir, f"variant_{r['index']:03d}.png")
            img.save(path, "PNG")
            r["file"] = os.path.basename(path)
            gallery.append((path, f"#{r['index']} · str {r['strength']} · {r['prompt'][:70]}"))
        done += len(chunk)
        progress(done / count, desc=f"{done} / {count} rendered")
        yield gallery, None, f"Rendering… {done} / {count}"

    manifest = os.path.join(run_dir, "manifest.csv")
    with open(manifest, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=["index", "file", "prompt", "strength", "seed"])
        w.writeheader()
        for r in sorted(recipes, key=lambda x: x["index"]):
            w.writerow({k: r.get(k, "") for k in w.fieldnames})

    bundle = os.path.join(run_dir, f"variants_{seed}.zip")
    with zipfile.ZipFile(bundle, "w", zipfile.ZIP_DEFLATED) as z:
        for r in recipes:
            z.write(os.path.join(run_dir, r["file"]), r["file"])
        z.write(manifest, "manifest.csv")

    gpu_s = count * cfg["sec_per_image"]
    if alpha_mode.startswith("Reuse"):
        note = "source mask reused"
    elif alpha_mode.startswith("Key"):
        note = "background keyed out"
    else:
        note = "opaque output"
    yield (gallery, bundle,
           f"**{count} variants** · seed `{seed}` · {note} · "
           f"~{gpu_s:.0f} s GPU (~{1500 / max(gpu_s, 1):.0f} runs/day on Pro quota)")


def estimate(model_name, count):
    s = MODELS[model_name]["sec_per_image"] * int(count)
    return (f"≈ {s:.0f} s GPU · {s / 15:.0f} % of a Pro day · "
            f"≈ {1500 // max(s, 1):.0f} runs before the quota resets")


def preview_palette(use_picked, c1, c2, c3):
    if not use_picked:
        return "<span class='mi-note'>Colors come from the pool below.</span>"
    chips = "".join(
        f"<span class='mi-chip' style='background:rgb{to_rgb(c)}'></span>"
        f"<span class='mi-note'>{name_color(c)}</span>"
        for c in (c1, c2, c3)
    )
    return f"<div class='mi-row'>{chips}</div>"


# --------------------------------------------------------------------------
# Interface
# --------------------------------------------------------------------------
CSS = """
.mi-note { font-size: 0.85rem; opacity: 0.78; line-height: 1.5; }
.mi-row { display: flex; align-items: center; gap: 8px; flex-wrap: wrap; }
.mi-chip { width: 20px; height: 20px; border-radius: 5px; display: inline-block;
           border: 1px solid rgba(128,128,128,.4); }
footer { display: none !important; }
"""
_MAJOR = int(gr.__version__.split(".")[0])
_STYLE = {"theme": gr.themes.Soft(), "css": CSS}
_BLOCKS_KW = {} if _MAJOR >= 6 else _STYLE
_LAUNCH_KW = _STYLE if _MAJOR >= 6 else {}

# img2img denoises toward a DESCRIPTION of the finished image — never phrase
# these as an instruction ("turn X into Y"), the model has no notion of "turn".
EXAMPLES = [
    "a minimal flat logo mark of a bicycle, thick even strokes, plain background",
    "a simple geometric mountain range logo, bold shapes, plain background",
    "a camping tent icon, flat vector logo, plain background",
    "a wordmark logo in heavy rounded sans serif, plain background",
]

POOL_SUBJECT_HINT = ("Put one target per line here — a bicycle logo, a tent "
                     "logo, a mountain logo — and every variant gets exactly "
                     "one of them instead of all three at once.")

with gr.Blocks(title="Mass Iteration Studio", **_BLOCKS_KW) as demo:
    gr.Markdown(
        "## Mass Iteration Studio\n"
        "One image in, N PNG variants out. Say what should change, pick your "
        "colors, set how far each variant may drift."
    )

    with gr.Row():
        with gr.Column(scale=4):
            image = gr.Image(label="Source image", type="pil", height=240,
                             image_mode="RGBA", sources=["upload", "clipboard"])

            instruction = gr.Textbox(
                label="Describe the finished image",
                placeholder="a minimal flat logo mark of a bicycle, thick even "
                            "strokes, plain background",
                lines=3,
            )
            gr.Markdown(
                "Describe **what you want to see**, not what to change. This "
                "model has no concept of *turn the sun into a bike* — it "
                "denoises toward your description. Naming several objects at "
                "once averages them; list them one per line under *Form* instead.",
                elem_classes="mi-note",
            )
            gr.Examples(EXAMPLES, inputs=instruction, label="Starting points")

            model_name = gr.Dropdown(list(MODELS), value=DEFAULT_MODEL, label="Model")
            guidance = gr.Slider(1.0, 12.0, MODELS[DEFAULT_MODEL]["guidance"],
                                 step=0.5, label="Prompt strength (guidance)")
            gr.Markdown(
                "How hard the model is pushed toward your text. 1.0 means the "
                "prompt barely registers — that is the price of the 2- and "
                "4-step models. 6–8 is where descriptions actually take hold.",
                elem_classes="mi-note",
            )
            count = gr.Slider(4, 200, 20, step=2, label="Number of variants")
            budget = gr.Markdown(estimate(DEFAULT_MODEL, 20), elem_classes="mi-note")
            run = gr.Button("Generate variants", variant="primary")

            with gr.Accordion("Colors", open=True):
                use_picked = gr.Checkbox(True, label="Use my colors")
                with gr.Row():
                    c1 = gr.ColorPicker("#e8794a", label="Primary")
                    c2 = gr.ColorPicker("#1a1a1a", label="Secondary")
                    c3 = gr.ColorPicker("#e8e4dc", label="Accent")
                swatches = gr.HTML(preview_palette(True, "#e8794a", "#1a1a1a", "#e8e4dc"))
                apply_mode = gr.Radio(
                    ["Prompt only", "Prompt + exact remap", "Exact remap only"],
                    value="Prompt only", label="How to apply them",
                )
                palette_mix = gr.Slider(0.2, 1.0, 0.85, step=0.05,
                                        label="Remap intensity")
                gr.Markdown(
                    "*Prompt only* lets the model interpret the palette — natural "
                    "results, approximate colors. *Exact remap* maps brightness onto "
                    "your exact values afterwards — precise colors, flatter look.",
                    elem_classes="mi-note",
                )

            with gr.Accordion("Transparency", open=True):
                alpha_mode = gr.Radio(
                    ["Key out the background color",
                     "Reuse the source mask",
                     "Opaque output"],
                    value="Key out the background color", label="Alpha channel",
                )
                matte = gr.ColorPicker("#ffffff", label="Background color")
                gr.Markdown(
                    "The model always needs an opaque image, so the source is "
                    "composited onto this color first.\n\n"
                    "**Key out** removes that color again afterwards — use this "
                    "whenever the shape may change, since a new subject needs its "
                    "own silhouette.\n\n"
                    "**Reuse the source mask** cuts every result to the original "
                    "outline. Correct below 0.4, destructive above it.",
                    elem_classes="mi-note",
                )

            with gr.Accordion("How far it may drift", open=False):
                smin = gr.Slider(0.15, 0.95, 0.55, step=0.05, label="Minimum reinvention")
                smax = gr.Slider(0.15, 0.95, 0.85, step=0.05, label="Maximum reinvention")
                gr.Markdown(
                    "**0.2–0.4** recolors and relights, the subject survives — this "
                    "is the range where a prompt looks ignored.\n\n"
                    "**0.5–0.65** forms shift, lettering starts garbling.\n\n"
                    "**0.75–0.9** a new subject can appear. Needed for *sun becomes "
                    "bicycle*; any wordmark will be destroyed.",
                    elem_classes="mi-note",
                )

            with gr.Accordion("Variation pools — one option per line", open=False):
                colors_raw = gr.Textbox(POOL_COLOR, label="Color (ignored when using your colors)", lines=5)
                styles_raw = gr.Textbox(POOL_STYLE, label="Style", lines=6)
                shapes_raw = gr.Textbox(POOL_SHAPE, label="Form / subject", lines=5)
                gr.Markdown(POOL_SUBJECT_HINT, elem_classes="mi-note")
                strat = gr.Radio(
                    ["Random mix", "Grid sweep (every combination in order)"],
                    value="Random mix", label="Sampling",
                )

            with gr.Accordion("Advanced", open=False):
                seed = gr.Number(1234, label="Base seed", precision=0)
                randomize = gr.Checkbox(True, label="New random seed each run")
                batch_size = gr.Slider(1, 8, 4, step=1, label="Batch size per GPU call")
                negative = gr.Textbox(NEGATIVE, label="Negative prompt", lines=2)

        with gr.Column(scale=6):
            gallery = gr.Gallery(label="Variants", columns=4, height=640,
                                 object_fit="contain", preview=True)
            status = gr.Markdown("")
            bundle = gr.File(label="Download all as ZIP (+ manifest.csv)", height=90)

    for c in (model_name, count):
        c.change(estimate, [model_name, count], budget)
    model_name.change(lambda m: gr.update(value=MODELS[m]["guidance"]),
                      model_name, guidance)
    for c in (use_picked, c1, c2, c3):
        c.change(preview_palette, [use_picked, c1, c2, c3], swatches)

    run.click(
        generate,
        [image, model_name, count, instruction, guidance, use_picked, c1, c2, c3,
         apply_mode, palette_mix, alpha_mode, matte,
         colors_raw, styles_raw, shapes_raw, strat, smin, smax,
         seed, randomize, batch_size, negative],
        [gallery, bundle, status],
    )

if __name__ == "__main__":
    demo.queue(max_size=12).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        ssr_mode=False,
        **_LAUNCH_KW,
    )