import os
import io
import zipfile
import random
from dataclasses import dataclass, asdict
from typing import List, Tuple

import gradio as gr
import numpy as np
import pandas as pd
from PIL import Image, ImageDraw, ImageFilter
import plotly.express as px
import plotly.graph_objects as go


# =====================================================
# Data structures & App State
# =====================================================
@dataclass
class Sample:
    image: Image.Image
    age: str
    gender: str
    ethnicity: str
    expression: str
    lighting: str
    background: str
    glasses: bool
    facial_hair: bool
    freckles: bool
    blush: bool
    detail: int
    hair_style: str


class AppState:
    def __init__(self):
        self.samples: List[Sample] = []
        self.label_target = "expression"


STATE = AppState()


# =====================================================
# Utility helpers
# =====================================================
def lerp(a, b, t):
    return int(a + (b - a) * t)

def lerp_color(c1, c2, t):
    return tuple(lerp(c1[i], c2[i], t) for i in range(3))

def linear_gradient(size: Tuple[int,int], c_top: Tuple[int,int,int], c_bottom: Tuple[int,int,int]):
    w, h = size
    img = Image.new("RGB", (w, h), c_top)
    draw = ImageDraw.Draw(img)
    for y in range(h):
        t = y / max(1, h - 1)
        draw.line([(0, y), (w, y)], fill=lerp_color(c_top, c_bottom, t))
    return img

def radial_highlight(img: Image.Image, center: Tuple[int,int], radius: int, strength: float=0.18):
    overlay = Image.new("RGBA", img.size, (0,0,0,0))
    d = ImageDraw.Draw(overlay)
    for r in range(radius, 0, -6):
        alpha = int(255 * strength * (r / radius) ** 2)
        d.ellipse([center[0]-r, center[1]-r, center[0]+r, center[1]+r], fill=(255,255,255, alpha))
    return Image.alpha_composite(img.convert("RGBA"), overlay).convert("RGB")

def seeded_rng(*args):
    s = "_".join(map(str, args))
    rnd = random.Random()
    rnd.seed(s)
    return rnd


# =====================================================
# Palettes (neutraal en breed)
# =====================================================
def skin_palette():
    return [
        (235,205,180), (226,188,156), (210,170,140), (190,150,120),
        (170,130,100), (150,110, 85), (130, 95, 70), (110, 80, 60)
    ]

def hair_palette():
    return [
        (30,30,30), (60,45,35), (90,70,55), (140,105,70),
        (200,170,110), (220,190,140), (50,50,70), (20,20,20)
    ]

def eye_palette():
    return [(80,110,170), (70,140,120), (120,90,60), (50,50,50)]

def clothing_palette():
    return [(60,70,110), (40,100,120), (120,60,80), (90,90,90), (30,120,60)]

def bg_sets():
    return {
        "soft":   ((245,235,235), (220,220,240)),
        "dramatic": ((50,60,80), (20,25,40)),
        "bright": ((250,245,210), (230,225,170)),
        "moody": ((45,70,85), (25,45,65)),
    }


# =====================================================
# Drawing primitives
# =====================================================
def draw_face_base(draw: ImageDraw.ImageDraw, cx, cy, w, h, color, outline=(50,40,40), jaw_ratio=0.14):
    jaw = int(h * jaw_ratio)
    left, right = cx - w//2, cx + w//2
    top, bottom = cy - h//2, cy + h//2
    draw.ellipse([left, top, right, bottom - jaw], fill=color, outline=outline, width=2)
    jaw_w = int(w * 0.86)
    draw.polygon([
        (cx - jaw_w//2, bottom - jaw),
        (cx + jaw_w//2, bottom - jaw),
        (cx + int(jaw_w*0.45), bottom),
        (cx - int(jaw_w*0.45), bottom),
    ], fill=color, outline=outline)

def draw_ears(draw, cx, cy, w, h, skin, outline):
    ear_w = int(w * 0.14)
    ear_h = int(h * 0.22)
    y = cy - int(h * 0.08)
    left = [cx - w//2 - ear_w//2, y - ear_h//2, cx - w//2 + ear_w//2, y + ear_h//2]
    right = [cx + w//2 - ear_w//2, y - ear_h//2, cx + w//2 + ear_w//2, y + ear_h//2]
    draw.ellipse(left, fill=skin, outline=outline)
    draw.ellipse(right, fill=skin, outline=outline)
    iw, ih = int(ear_w*0.5), int(ear_h*0.5)
    draw.ellipse([left[0]+ear_w//4, left[1]+ear_h//4, left[0]+ear_w//4+iw, left[1]+ear_h//4+ih], outline=(90,70,60))
    draw.ellipse([right[0]+ear_w//4, right[1]+ear_h//4, right[0]+ear_w//4+iw, right[1]+ear_h//4+ih], outline=(90,70,60))

def draw_eyes(draw, cx, cy, w, h, iris_color, expression, rng, detail=2):
    eye_y = cy - int(h*0.11) + rng.randint(-3, 3)
    eye_dx = int(w*0.22)
    eye_w, eye_h = int(w*0.17), int(h*0.095)

    lid_offset = int(h*0.045)
    for ex in (-eye_dx, eye_dx):
        draw.arc([cx+ex-eye_w//2, eye_y-lid_offset, cx+ex+eye_w//2, eye_y+lid_offset], start=180, end=360, fill=(80,60,60), width=1)

    for ex in (-eye_dx, eye_dx):
        draw.ellipse([cx+ex-eye_w//2, eye_y-eye_h//2, cx+ex+eye_w//2, eye_y+eye_h//2], fill=(250,250,250), outline=(0,0,0))
        ir = int(min(eye_w, eye_h) * 0.46)
        draw.ellipse([cx+ex-ir, eye_y-ir, cx+ex+ir, eye_y+ir], fill=iris_color, outline=(0,0,0))
        rim = ir + 1
        draw.ellipse([cx+ex-rim, eye_y-rim, cx+ex+rim, eye_y+rim], outline=(40,40,40))
        pr = int(ir*0.45)
        draw.ellipse([cx+ex-pr, eye_y-pr, cx+ex+pr, eye_y+pr], fill=(0,0,0))
        draw.ellipse([cx+ex-pr//2, eye_y-pr//2, cx+ex-pr//2+4, eye_y-pr//2+4], fill=(255,255,255))

    brow_y = eye_y - int(h*0.08)
    brow_len = int(w*0.23)
    slope = {"happy": -6, "angry": 6, "surprised": -2, "neutral": 0}.get(expression, 0)
    for ex, sign in ((-eye_dx, -1), (eye_dx, 1)):
        x1 = cx + ex - brow_len//2
        x2 = cx + ex + brow_len//2
        y1 = brow_y + slope * sign // 2
        y2 = brow_y - slope * sign // 2
        draw.line([x1, y1, x2, y2], fill=(30,30,30), width=4 if detail >= 1 else 3)

def draw_nose(draw, cx, cy, h, detail=2):
    ny = cy - int(h*0.02)
    draw.line([cx, ny, cx, ny + int(h*0.10)], fill=(90,70,60), width=1 if detail < 2 else 2)
    yb = ny + int(h*0.12)
    draw.arc([cx-18, yb-4, cx-6, yb+8], start=0, end=180, fill=(60,45,40), width=2)
    draw.arc([cx+6, yb-4, cx+18, yb+8], start=0, end=180, fill=(60,45,40), width=2)

def draw_mouth(draw, cx, cy, w, h, expression, detail=2):
    my = cy + int(h*0.15)
    mw = int(w*0.30)
    if expression == "happy":
        draw.arc([cx-mw, my-12, cx+mw, my+16], start=0, end=180, fill=(60,30,30), width=4)
        draw.line([cx-mw//2, my+5, cx+mw//2, my+5], fill=(120,60,70), width=2)
    elif expression == "angry":
        draw.arc([cx-mw, my-2, cx+mw, my+24], start=180, end=360, fill=(60,30,30), width=4)
        draw.line([cx-mw//2, my+9, cx+mw//2, my+9], fill=(120,60,70), width=2)
    elif expression == "surprised":
        draw.ellipse([cx-16, my-6, cx+16, my+26], outline=(60,30,30), width=4)
    else:
        draw.line([cx-mw//2, my+8, cx+mw//2, my+8], fill=(60,30,30), width=4)
        draw.line([cx-mw//3, my+10, cx+mw//3, my+10], fill=(120,60,70), width=2)

def draw_glasses(draw, cx, cy, w, h):
    gy = cy - int(h*0.11)
    gw = int(w*0.22)
    gh = int(h*0.10)
    for ex in (-gw, gw):
        draw.rectangle([cx+ex - gw//2, gy - gh//2, cx+ex + gw//2, gy + gh//2], outline=(30,30,30), width=3)
    draw.line([cx - gw//2, gy, cx + gw//2, gy], fill=(30,30,30), width=3)
    draw.line([cx - int(w*0.38), gy, cx - int(w*0.45), gy - 8], fill=(30,30,30), width=3)
    draw.line([cx + int(w*0.38), gy, cx + int(w*0.45), gy - 8], fill=(30,30,30), width=3)

def draw_facial_hair(draw, cx, cy, w, h, color=(40,30,25), opacity=90):
    beard = Image.new("RGBA", (w*2, h*2), (0,0,0,0))
    d = ImageDraw.Draw(beard)
    bx, by = w, h
    d.arc([bx-60, by+10, bx+60, by+70], start=200, end=340, fill=(*color, opacity), width=10)
    d.pieslice([bx-120, by+60, bx+120, by+180], start=10, end=170, fill=(*color, int(opacity*0.7)))
    return beard

def draw_freckles(draw, cx, cy, w, h, density=35):
    for _ in range(density):
        x = cx + random.randint(-int(w*0.25), int(w*0.25))
        y = cy + random.randint(0, int(h*0.2))
        r = random.randint(1, 2)
        draw.ellipse([x-r, y-r, x+r, y+r], fill=(110,80,70))

def draw_blush(draw, cx, cy, w, h):
    for sign in (-1, 1):
        bx = cx + sign * int(w*0.22)
        by = cy + int(h*0.08)
        for r, alpha in ((26, 30), (20, 50), (14, 70)):
            draw.ellipse([bx-r, by-r, bx+r, by+r], fill=(255,140,150, alpha))

def draw_neck_shoulders(draw, cx, cy, w, h, skin, cloth_color):
    neck_w = int(w*0.24)
    neck_h = int(h*0.18)
    nx1, nx2 = cx - neck_w//2, cx + neck_w//2
    ny1, ny2 = cy + int(h*0.2), cy + int(h*0.2) + neck_h
    draw.rectangle([nx1, ny1, nx2, ny2], fill=skin)
    sw = int(w*0.94)
    sy = ny2
    draw.pieslice([cx - sw, sy - 40, cx + sw, sy + 140], start=0, end=180, fill=cloth_color)

# ---- Hair split: back & front (voorkomt 'achterhoofden') ----
def draw_hair_back(draw, cx, cy, w, h, color, style, rng):
    top = cy - h//2
    left, right = cx - w//2, cx + w//2
    if style == "short":
        draw.ellipse([left-6, top-28, right+6, top + h*0.34], fill=color)
    elif style == "medium":
        draw.ellipse([left-12, top-36, right+12, top + h*0.44], fill=color)
        draw.polygon([(left-14, cy-8), (left+18, cy+70), (left-10, cy+84)], fill=color)
        draw.polygon([(right+14, cy-8), (right-18, cy+70), (right+10, cy+84)], fill=color)
    elif style == "curly":
        for _ in range(120):
            rx = rng.randint(-w//2, w//2)
            ry = rng.randint(-h//2-40, -h//2+90)
            r = rng.randint(8, 16)
            draw.ellipse([cx+rx-r, cy+ry-r, cx+rx+r, cy+ry+r], fill=color)
    elif style == "bun":
        draw.ellipse([left, top-18, right, top + h*0.38], fill=color)
        draw.ellipse([cx-32, top-56, cx+32, top-12], fill=color)
    elif style == "bangs":
        draw.ellipse([left-6, top-28, right+6, top + h*0.36], fill=color)
    else:
        draw.ellipse([left-6, top-28, right+6, top + h*0.36], fill=color)

def draw_hair_front(draw, cx, cy, w, h, color, style, rng, with_strands=True):
    top = cy - h//2
    left, right = cx - w//2, cx + w//2
    forehead_y = top + int(h*0.18)
    if style == "bangs":
        for i in range(12):
            x = left + i * (w // 12)
            seg_w = int(w/12.5)
            draw.rectangle([x, forehead_y-6, x+seg_w, forehead_y+10], fill=color)
    elif style in ("medium", "curly"):
        for i in range(6):
            x = rng.randint(left+20, right-20)
            y1 = forehead_y - rng.randint(2, 8)
            y2 = forehead_y + rng.randint(0, 6)
            draw.line([x, y1, x, y2], fill=color, width=2)
    else:
        draw.line([left+20, forehead_y, right-20, forehead_y], fill=color, width=2)

    if with_strands:
        strand_color = tuple(max(0, c-20) for c in color)
        for _ in range(28):
            x1 = rng.randint(left+24, right-24)
            y1 = forehead_y - rng.randint(0, 6)
            x2 = x1 + rng.randint(-8, 8)
            y2 = y1 + rng.randint(4, 12)
            draw.line([x1, y1, x2, y2], fill=strand_color, width=1)


# =====================================================
# Portrait generator (CPU-only, detailed, with hair fix)
# =====================================================
HAIR_STYLES = ["auto", "short", "medium", "curly", "bangs", "bun"]

def generate_portrait(age, gender, ethnicity, expression, lighting, background,
                      glasses=False, facial_hair=False, freckles=False, blush=False, detail=2,
                      hair_style="auto", seed=None) -> Image.Image:
    rnd = seeded_rng(age, gender, ethnicity, expression, lighting, background, glasses, facial_hair, freckles, blush, detail, hair_style, seed or random.random())
    W, H = 512, 512

    # Background gradient + light spot
    top_c, bot_c = bg_sets().get(lighting, ((235,235,240),(210,210,225)))
    bg = linear_gradient((W, H), top_c, bot_c)
    bg = radial_highlight(bg, (int(W*0.36), int(H*0.30)), 190, strength=0.16)

    # Base canvas
    img = bg.copy().convert("RGBA")
    draw = ImageDraw.Draw(img, "RGBA")

    # Colors (slight randomization)
    skin = tuple(min(255, max(0, c + rnd.randint(-6, 6))) for c in rnd.choice(skin_palette()))
    hair = rnd.choice(hair_palette())
    iris = rnd.choice(eye_palette())
    cloth = rnd.choice(clothing_palette())

    # Head geometry (age tweaks)
    cx, cy = W//2, int(H*0.48)
    head_w = rnd.randint(250, 290)
    head_h = rnd.randint(270, 310)
    jaw_ratio = 0.15 if age == "senior" else 0.13

    # 1) BACK HAIR (achterkant)
    style = hair_style if hair_style in HAIR_STYLES and hair_style != "auto" else rnd.choice(HAIR_STYLES[1:])
    draw_hair_back(draw, cx, cy, head_w, head_h, hair, style, rnd)

    # 2) NECK & SHOULDERS
    draw_neck_shoulders(draw, cx, cy, head_w, head_h, skin, cloth)

    # 3) EARS
    draw_ears(draw, cx, cy, head_w, head_h, skin, outline=(50,40,40))

    # 4) FACE BASE
    draw_face_base(draw, cx, cy, head_w, head_h, skin, outline=(50,40,40), jaw_ratio=jaw_ratio)

    # 5) FACIAL FEATURES
    draw_eyes(draw, cx, cy, head_w, head_h, iris, expression, rnd, detail=detail)
    draw_nose(draw, cx, cy, head_h, detail=detail)
    draw_mouth(draw, cx, cy, head_w, head_h, expression, detail=detail)

    # 6) FRONT HAIR (kleine lokjes, niet over gezicht)
    draw_hair_front(draw, cx, cy, head_w, head_h, hair, style, rnd, with_strands=True)

    # 7) BLUSH / FRECKLES / ACCESSOIRES
    if blush:
        blush_layer = Image.new("RGBA", img.size, (0,0,0,0))
        db = ImageDraw.Draw(blush_layer, "RGBA")
        draw_blush(db, cx, cy, head_w, head_h)
        img = Image.alpha_composite(img, blush_layer)

    if freckles:
        fr = Image.new("RGBA", img.size, (0,0,0,0))
        df = ImageDraw.Draw(fr, "RGBA")
        draw_freckles(df, cx, cy, head_w, head_h, density=35)
        img = Image.alpha_composite(img, fr)

    if glasses:
        draw_glasses(draw, cx, cy, head_w, head_h)

    if facial_hair:
        beard = draw_facial_hair(draw, cx, cy, 300, 300, color=tuple(max(0, c-10) for c in hair), opacity=90)
        img.alpha_composite(beard, (cx-300, cy-300))

    # 8) VIGNETTE
    vign = Image.new("L", (W, H), 0)
    dv = ImageDraw.Draw(vign)
    for r in range(420, 0, -24):
        alpha = int(255 * (r/420) ** 2 * 0.10)
        dv.ellipse([cx-r, cy-r, cx+r, cy+r], fill=alpha)
    img_rgb = Image.composite(img.convert("RGB"), Image.new("RGB", (W,H), (0,0,0)),
                              vign.filter(ImageFilter.GaussianBlur(12)).point(lambda p: 255 - p))
    return img_rgb


def generate_avatars(age, gender, ethnicity, expression, lighting, background,
                     glasses, facial_hair, freckles, blush, detail, hair_style, num_images=6):
    imgs = []
    for i in range(int(num_images)):
        im = generate_portrait(
            age, gender, ethnicity, expression, lighting, background,
            glasses=glasses, facial_hair=facial_hair,
            freckles=freckles, blush=blush, detail=int(detail),
            hair_style=hair_style, seed=i
        )
        imgs.append(im)
    return imgs


# =====================================================
# Dataset & Export
# =====================================================
def dataset_stats_markdown():
    if not STATE.samples:
        return "Dataset: **0** samples"
    df = pd.DataFrame([asdict(s) | {"image": None} for s in STATE.samples])
    lines = [f"**Dataset:** {len(df)} samples"]
    for col in ["age", "gender", "ethnicity", "expression"]:
        counts = df[col].value_counts().to_dict()
        top = ", ".join([f"{k}: {v}" for k, v in counts.items()])
        lines.append(f"- {col}: {top}")
    return "\n".join(lines)

def export_dataset_zip():
    if not STATE.samples:
        return None
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as zf:
        rows = []
        for idx, s in enumerate(STATE.samples):
            fname = f"images/{idx:06d}.png"
            img_bytes = io.BytesIO()
            s.image.save(img_bytes, format="PNG")
            zf.writestr(fname, img_bytes.getvalue())
            r = asdict(s).copy()
            r["image_path"] = fname
            rows.append(r)
        df = pd.DataFrame(rows)
        zf.writestr("metadata.csv", df.to_csv(index=False))
    buffer.seek(0)
    return buffer


# =====================================================
# Callbacks
# =====================================================
def cb_generate(age, gender, ethnicity, expression, lighting, background,
                glasses, facial_hair, freckles, blush, detail, hair_style, num_images):
    images = generate_avatars(age, gender, ethnicity, expression, lighting, background,
                              glasses, facial_hair, freckles, blush, detail, hair_style, int(num_images))
    for im in images:
        STATE.samples.append(Sample(
            im, age, gender, ethnicity, expression, lighting, background,
            glasses, facial_hair, freckles, blush, int(detail), hair_style
        ))
    return images, dataset_stats_markdown()

def cb_train(target_label):
    STATE.label_target = target_label
    if len(STATE.samples) < 6:
        return "**Need ≥ 6 samples to simulate training.**", go.Figure(), go.Figure()

    labels = [getattr(s, target_label) for s in STATE.samples]
    vc = pd.Series(labels).value_counts(normalize=True)
    balance_score = 1.0 - (vc.max() - 1.0 / len(vc)) if len(vc) > 0 else 0.2
    acc = 0.65 + 0.3 * max(0, min(balance_score, 1.0))

    epochs = list(range(1, 9))
    loss_vals = np.linspace(0.9, 0.2, num=len(epochs)) * (1.1 - (acc - 0.5))

    fig_loss = go.Figure()
    fig_loss.add_trace(go.Scatter(x=epochs, y=loss_vals, mode="lines+markers", name="loss"))
    fig_loss.update_layout(title="Training Curve (simulated)", xaxis_title="Epoch", yaxis_title="Loss")

    labels_unique = sorted(vc.index.tolist())
    n = len(labels_unique) if labels_unique else 2
    mat = np.random.rand(n, n) * (1.0 - acc) * 0.6
    np.fill_diagonal(mat, acc / max(n, 1) + 0.1)
    fig_cm = px.imshow(mat, x=labels_unique, y=labels_unique, title="Confusion Matrix (simulated)", aspect="auto", text_auto=True)

    acc_md = f"**Accuracy ({target_label})**: {acc:.3f} *(simulated)*"
    return acc_md, fig_loss, fig_cm

def cb_visualize():
    if len(STATE.samples) < 3:
        return go.Figure()
    labels = [getattr(s, STATE.label_target) for s in STATE.samples]
    uniq = sorted(pd.Series(labels).unique().tolist())
    xs, ys, cs = [], [], []
    for lab in labels:
        base = uniq.index(lab) * 3.0
        xs.append(np.random.normal(loc=base, scale=0.8))
        ys.append(np.random.normal(loc=base, scale=0.8))
        cs.append(lab)
    fig = px.scatter(x=xs, y=ys, color=cs, title="Embedding Visualization (simulated 2D)", labels={"x": "x", "y": "y", "color": "label"})
    return fig

def cb_export():
    buf = export_dataset_zip()
    if buf is None:
        return None
    return ("synthetic_avatar_dataset.zip", buf)

def cb_clear_dataset():
    STATE.samples = []
    return [], "Dataset: **0** samples", "Cleared."


# =====================================================
# Gradio UI
# =====================================================
with gr.Blocks(title="Synthetic Avatar Dataset Trainer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧠 Synthetic Avatar Dataset Trainer — CPU Portraits (Hair Fix + Style Choice)")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Controls")
            age = gr.Dropdown(["young", "adult", "senior"], value="adult", label="Age")
            gender = gr.Dropdown(["female", "male"], value="female", label="Gender")
            ethnicity = gr.Dropdown(["black", "white", "asian", "indian", "latino", "arab"], value="asian", label="Ethnicity")
            expression = gr.Dropdown(["neutral", "happy", "surprised", "angry"], value="neutral", label="Expression")
            lighting = gr.Dropdown(["soft", "dramatic", "bright", "moody"], value="soft", label="Lighting")
            background = gr.Dropdown(["light gray", "studio dark", "pastel", "bokeh"], value="light gray", label="Background")
            hair_style = gr.Dropdown(choices=["auto", "short", "medium", "curly", "bangs", "bun"], value="auto", label="Hair style")
            with gr.Row():
                glasses = gr.Checkbox(label="Glasses", value=False)
                facial_hair = gr.Checkbox(label="Facial hair", value=False)
            with gr.Row():
                freckles = gr.Checkbox(label="Freckles", value=False)
                blush = gr.Checkbox(label="Blush", value=False)
            detail = gr.Slider(0, 3, value=2, step=1, label="Detail level")
            num_images = gr.Slider(1, 12, value=6, step=1, label="Images per click")
            btn_generate = gr.Button("Generate Avatars")
            stats_md = gr.Markdown("Dataset: **0** samples")

            with gr.Row():
                btn_export = gr.Button("Download Dataset (.zip)")
                btn_clear = gr.Button("Clear Dataset")

        with gr.Column(scale=2):
            gr.Markdown("### Dataset")
            gallery = gr.Gallery(label="Generated Avatars", columns=3, height=420, type="pil", allow_preview=True)

            gr.Markdown("### Training")
            target_label = gr.Radio(choices=["expression", "age", "gender", "ethnicity"], value="expression", label="Target label")
            btn_train = gr.Button("Train Model (Simulated)")
            acc_md = gr.Markdown("")
            loss_plot = gr.Plot()
            cm_plot = gr.Plot()

        with gr.Column(scale=1):
            gr.Markdown("### Visualize")
            btn_visualize = gr.Button("Compute Embeddings & Plot (2D)")
            scatter = gr.Plot()
            gr.Markdown("> 💡 CPU-only — back & front hair layers, geen achterhoofden; met haarstijlkeuze.")

    btn_generate.click(
        cb_generate,
        inputs=[age, gender, ethnicity, expression, lighting, background, glasses, facial_hair, freckles, blush, detail, hair_style, num_images],
        outputs=[gallery, stats_md]
    )
    btn_train.click(cb_train, inputs=[target_label], outputs=[acc_md, loss_plot, cm_plot])
    btn_visualize.click(cb_visualize, inputs=[], outputs=[scatter])
    btn_export.click(cb_export, inputs=[], outputs=[gr.File(label="Download")])
    btn_clear.click(cb_clear_dataset, inputs=[], outputs=[gallery, stats_md, acc_md])

if __name__ == "__main__":
    print("✅ App loaded successfully — CPU Portraits (Hair Fix + Style Choice)")
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))