import os
import gradio as gr
import numpy as np
from PIL import Image, ImageOps
import torch
from diffusers import StableDiffusionImg2ImgPipeline, EulerAncestralDiscreteScheduler

# -----------------------------
# CPU performance knobs
# -----------------------------
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["MKL_NUM_THREADS"] = "2"
try:
    torch.set_num_threads(2)
except Exception:
    pass

torch.set_grad_enabled(False)

device = "cpu"
dtype = torch.float32

MODEL_ID = "runwayml/stable-diffusion-v1-5"

pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=dtype,
    safety_checker=None,
    requires_safety_checker=False,
)

pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to(device)

# Make diffusers quieter / slightly faster
pipe.set_progress_bar_config(disable=True)

# Helpful on CPU
try:
    pipe.enable_attention_slicing()
except Exception:
    pass

# VAE optimizations (often helps)
try:
    pipe.enable_vae_slicing()
except Exception:
    pass
try:
    pipe.enable_vae_tiling()
except Exception:
    pass

# -----------------------------
# Image utils
# -----------------------------
def _to_pil(img_np):
    if img_np is None:
        return None
    if isinstance(img_np, Image.Image):
        return img_np.convert("RGB")
    arr = np.asarray(img_np)
    if arr.dtype != np.uint8:
        arr = np.clip(arr, 0, 255).astype(np.uint8)
    if arr.ndim == 3 and arr.shape[2] == 4:
        arr = arr[:, :, :3]
    return Image.fromarray(arr).convert("RGB")

def _center_square(pil_img: Image.Image, out_size=512):
    w, h = pil_img.size
    s = min(w, h)
    left = (w - s) // 2
    top = (h - s) // 2
    pil_img = pil_img.crop((left, top, left + s, top + s))
    pil_img = pil_img.resize((out_size, out_size), Image.LANCZOS)
    return pil_img

def _blend_parents(dad_pil, mom_pil, out_size=512):
    dad = _center_square(dad_pil, out_size)
    mom = _center_square(mom_pil, out_size)
    blended = Image.blend(dad, mom, alpha=0.5)
    blended = ImageOps.autocontrast(blended, cutoff=1)
    return blended

# -----------------------------
# Generation
# -----------------------------
def generate_cartoon_kid(
    dad_np,
    mom_np,
    age,
    style_strength,
    steps,
    guidance_scale,
    seed,
    extra_prompt,
    negative_prompt
):
    dad = _to_pil(dad_np)
    mom = _to_pil(mom_np)
    if dad is None or mom is None:
        raise gr.Error("ارفع صورتين: الأب + الأم.")

    init = _blend_parents(dad, mom, out_size=512)

    # ✅ Faster on CPU: 256 instead of 384
    init_small = init.resize((256, 256), Image.LANCZOS)

    base_prompt = (
        "cute 3d animated child portrait, pixar-like style, "
        "soft studio lighting, smooth skin, big friendly eyes, "
        "high quality, centered face, clean background"
    )

    age = int(age)
    if age <= 7:
        age_prompt = f"a {age}-year-old kid, very cute, youthful"
    elif age <= 12:
        age_prompt = f"a {age}-year-old pre-teen kid, youthful"
    else:
        age_prompt = f"a {age}-year-old teen, youthful"

    prompt = f"{base_prompt}, {age_prompt}"
    if extra_prompt and extra_prompt.strip():
        prompt = f"{prompt}, {extra_prompt.strip()}"

    if seed is None or int(seed) < 0:
        seed = np.random.randint(0, 2**31 - 1)

    gen = torch.Generator(device=device).manual_seed(int(seed))
    strength = float(style_strength)

    # ✅ Inference mode saves overhead
    with torch.inference_mode():
        out = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=init_small,
            strength=strength,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance_scale),
            generator=gen,
        ).images[0]

    out = out.resize((512, 512), Image.LANCZOS)
    return init, out, int(seed), prompt

# -----------------------------
# UI (same)
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 👶 Kid Generator — Cartoon (CPU)\nFaster CPU settings (256px + fewer steps).")

    with gr.Row():
        dad_in = gr.Image(type="numpy", label="Father")
        mom_in = gr.Image(type="numpy", label="Mother")

    with gr.Row():
        age = gr.Slider(4, 15, value=11, step=1, label="Age")
        style_strength = gr.Slider(0.45, 0.85, value=0.65, step=0.01, label="Cartoon Strength")

    with gr.Row():
        # ✅ recommend smaller steps default
        steps = gr.Slider(6, 20, value=10, step=1, label="Steps (CPU faster)")
        guidance_scale = gr.Slider(3, 8, value=5.5, step=0.1, label="Guidance Scale")
        seed = gr.Number(value=-1, precision=0, label="Seed (-1 random)")

    extra_prompt = gr.Textbox(label="Extra Prompt (optional)", placeholder="smiling, curly hair, freckles")
    negative_prompt = gr.Textbox(
        label="Negative Prompt",
        value="realistic photo, ugly, deformed, blurry, low quality, watermark, text, scary, old, wrinkles"
    )

    btn = gr.Button("Generate Cartoon Kid", variant="primary")

    with gr.Row():
        init_out = gr.Image(label="Blended Parents (Init)")
        out_img = gr.Image(label="Cartoon Kid Result")

    used_seed = gr.Number(label="Used Seed", precision=0)
    used_prompt = gr.Textbox(label="Used Prompt")

    btn.click(
        fn=generate_cartoon_kid,
        inputs=[dad_in, mom_in, age, style_strength, steps, guidance_scale, seed, extra_prompt, negative_prompt],
        outputs=[init_out, out_img, used_seed, used_prompt],
    )

demo.launch()