# app.py (v4, CPU-optimized + brede Gradio-compat)
# Neurale netwerken als tekenaars — Stable Diffusion + ControlNet (scribble)
# - Werkt op CPU (langzamer dan GPU, maar demowaardig)
# - Gebruikt gr.Sketchpad (compatibel met oudere Gradio-versies)
# - Geen queue-parameters (demo.queue() zonder args)

import os, random, numpy as np
from typing import Optional
from PIL import Image, ImageDraw
import gradio as gr
import torch
from diffusers import (
    ControlNetModel,
    StableDiffusionControlNetPipeline,
    DPMSolverMultistepScheduler,
)

# Modellen (kun je via Space Secrets/Opsional env-vars overriden)
MODEL_BASE = os.environ.get("SD_BASE_MODEL", "runwayml/stable-diffusion-v1-5")
MODEL_CN   = os.environ.get("CN_SCRIBBLE_MODEL", "lllyasviel/sd-controlnet-scribble")

DEVICE = "cpu"              # geforceerd CPU
DTYPE  = torch.float32
TARGET_SIZE = 512           # lager = sneller op CPU
DEFAULT_STEPS = 12          # CPU-vriendelijk

pipe: Optional[StableDiffusionControlNetPipeline] = None

def _lazy_load_pipeline():
    global pipe
    if pipe is not None:
        return pipe
    controlnet = ControlNetModel.from_pretrained(MODEL_CN, torch_dtype=DTYPE)
    pipe = StableDiffusionControlNetPipeline.from_pretrained(
        MODEL_BASE, controlnet=controlnet, torch_dtype=DTYPE, safety_checker=None
    )
    # Snellere scheduler → minder stappen nodig
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to(DEVICE)
    pipe.enable_attention_slicing()
    return pipe

def to_rgba(img):
    """Accepteer numpy of PIL en geef een PIL RGBA image terug."""
    if img is None:
        return None
    if isinstance(img, np.ndarray):
        if img.dtype != np.uint8:
            img = img.astype(np.uint8)
        return Image.fromarray(img).convert("RGBA")
    if isinstance(img, Image.Image):
        return img.convert("RGBA")
    try:
        return Image.open(img).convert("RGBA")
    except Exception:
        return None

def rgba_to_scribble(img) -> Image.Image:
    """Zet RGBA schets om naar zwart-op-wit, vierkant en resized (TARGET_SIZE)."""
    img = to_rgba(img)
    if img is None:
        return None
    # transparant → wit
    bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
    img = Image.alpha_composite(bg, img)
    # grijs + drempel
    arr = np.array(img.convert("L"))
    thr = max(25, int(arr.mean() * 0.9))
    lines = arr < thr
    out = np.full((*arr.shape, 3), 255, dtype=np.uint8)
    out[lines] = (0, 0, 0)
    scribble = Image.fromarray(out, "RGB")
    # vierkant + resize
    m = max(scribble.width, scribble.height)
    sq = Image.new("RGB", (m, m), (255, 255, 255))
    sq.paste(scribble, ((m - scribble.width) // 2, (m - scribble.height) // 2))
    return sq.resize((TARGET_SIZE, TARGET_SIZE), Image.BICUBIC)

def run(drawing, prompt, negative_prompt,
        guidance_scale=7.0, controlnet_conditioning_scale=1.0,
        num_inference_steps=DEFAULT_STEPS, seed=-1):
    if drawing is None:
        raise gr.Error("Teken eerst iets of gebruik de voorbeeld-schets.")
    cn_image = rgba_to_scribble(drawing)
    if cn_image is None:
        raise gr.Error("Schets niet leesbaar. Probeer opnieuw of upload PNG/JPG.")

    if not prompt or not str(prompt).strip():
        prompt = "clean pencil sketch, coherent completion, subtle shading, high quality"
    if seed is None or int(seed) < 0:
        seed = random.randint(0, 2**31 - 1)
    gen = torch.Generator(device=DEVICE).manual_seed(int(seed))

    p = _lazy_load_pipeline()
    full_prompt = (
        "black pencil sketch, clean lines, coherent shape completion, subtle shading, "
        "consistent with the input scribble, " + prompt
    )
    result = p(
        prompt=full_prompt,
        negative_prompt=negative_prompt or "text, watermark, extra limbs, low quality, distorted face, nsfw",
        image=cn_image,
        width=cn_image.width, height=cn_image.height,
        generator=gen, num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale),
        controlnet_conditioning_scale=float(controlnet_conditioning_scale),
    )
    image = result.images[0]
    pad = 12
    combo = Image.new("RGB", (cn_image.width * 2 + pad, cn_image.height), (245, 245, 245))
    combo.paste(cn_image, (0, 0))
    combo.paste(image, (cn_image.width + pad, 0))
    return image, cn_image, combo, f"Seed: {seed}, steps: {num_inference_steps}"

def make_example_scribble() -> Image.Image:
    """Kleine voorbeeldschets (appelvorm) voor snelle test."""
    w, h = 360, 360
    img = Image.new("RGBA", (w, h), (0, 0, 0, 0))
    d = ImageDraw.Draw(img)
    d.line([(100,300),(90,230),(110,170),(160,140),(200,140),
            (260,170),(280,230),(260,300),(180,320),(100,300)],
           fill=(0,0,0,255), width=12)
    d.line([(200,140),(210,110)], fill=(0,0,0,255), width=10)
    d.line([(210,110),(250,120),(265,150)], fill=(0,0,0,255), width=9)
    return img

CSS = ".gradio-container{max-width:1100px} #combo img{border-radius:10px}"

with gr.Blocks(title="🖌️ Neurale netwerken als tekenaars (CPU v4)", css=CSS) as demo:
    gr.Markdown("## 🖌️ Neurale netwerken als tekenaars (CPU)\nTeken links of klik **Voorbeeld-schets**; AI vult aan met Stable Diffusion + ControlNet (scribble).")

    with gr.Row():
        with gr.Column():
            canvas = gr.Sketchpad(label="Schets (teken in zwart op transparant/wit)", height=420)
            prompt = gr.Textbox(label="Prompt", value="realistische appel met blad, potloodschets")
            negative_prompt = gr.Textbox(label="Negative prompt (optioneel)", value="tekst, watermerk, wazig")
            with gr.Accordion("Geavanceerd", open=False):
                guidance_scale = gr.Slider(5.0, 10.0, value=7.0, step=0.5, label="Guidance scale")
                cn_strength   = gr.Slider(0.6, 1.6, value=1.0, step=0.05, label="ControlNet conditioning scale")
                steps         = gr.Slider(6, 25, value=DEFAULT_STEPS, step=1, label="Aantal diffusion-steps")
                seed          = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
            with gr.Row():
                run_btn = gr.Button("Vervolledig ✨", variant="primary")
                example_btn = gr.Button("Voorbeeld-schets", variant="secondary")

        with gr.Column():
            out_img = gr.Image(label="Vervolledigde tekening", interactive=False)
            cn_prev = gr.Image(label="Scribble (input voor AI)", interactive=False)
            combo   = gr.Image(label="Vergelijking (links: input, rechts: AI)", elem_id="combo")
            meta    = gr.Markdown()

    def fill_and_run(p, n, gs, cs, st, sd):
        img = make_example_scribble()
        a, b, c, d = run(img, p, n, gs, cs, st, sd)
        return img, a, b, c, d

    run_btn.click(run, inputs=[canvas, prompt, negative_prompt, guidance_scale, cn_strength, steps, seed],
                  outputs=[out_img, cn_prev, combo, meta])
    example_btn.click(fill_and_run, inputs=[prompt, negative_prompt, guidance_scale, cn_strength, steps, seed],
                      outputs=[canvas, out_img, cn_prev, combo, meta])

demo.queue()  # zonder parameters: compatibel met oudere Gradio-versies

if __name__ == "__main__":
    demo.launch()