# -*- coding: utf-8 -*-
"""gui_api Colorway Lab: mood board in, ready-made colorways out.

Designers build colorways by hand: pick a color, open every fabric, type it in, render, compare, repeat.
Colorway Lab turns one idea into a set of 3D-rendered options in a few seconds:

1. Pick a palette source:
   - a mood board / photo (PNG): the dominant colors are extracted (k-means),
   - a base color plus a harmony rule (complementary, analogous, triadic, ...),
   - or one of the curated palettes.
2. "Generate" creates one colorway per variation. Colors are assigned by fabric area: the biggest fabric
   gets the dominant color, smaller ones the accents, and every variation shifts the assignment. Fabric
   textures are desaturated so the palette color shows true. Each variation is rendered and shown with
   its palette strip.
3. "Keep" marks the ones you like. "Finish" deletes the others and renames the kept colorways.
   "Lookbook" writes the kept ones side by side into one image.

Nothing touches the original colorway: every variation is a copy.
"""
import colorsys
import datetime
import os
import random
import re
import struct
import tempfile
import zlib

import export_api
import fabric_api
import gui_api
import pattern_api
import utility_api

if not gui_api.IsGuiAvailable():
    raise SystemExit("gui_api needs the CLO GUI (not available in headless mode)")

SLOTS = 6
WORK_DIR = os.path.join(tempfile.gettempdir(), "clo_colorway_lab").replace(os.sep, "/")
OUT_DIR = os.path.join(os.path.expanduser("~"), "Documents", "CLO", "ColorwayLab").replace(os.sep, "/")
os.makedirs(WORK_DIR, exist_ok=True)

CURATED = {
    "Tokyo Neon": ["#0d0221", "#ff2a6d", "#05d9e8", "#d1f7ff", "#7700a6"],
    "Nordic Winter": ["#e9eef2", "#9fb4c7", "#2e4057", "#c9a227", "#1b263b"],
    "Desert Sand": ["#d9b38c", "#a8653a", "#f2e3c6", "#5e3b2b", "#c27c4e"],
    "Ocean Deep": ["#03045e", "#0077b6", "#00b4d8", "#90e0ef", "#caf0f8"],
    "Forest Walk": ["#2d3a2e", "#5c7a52", "#a3b18a", "#dad7cd", "#8b5e34"],
    "Pastel Bloom": ["#ffd6e0", "#c1e1c1", "#fef9c3", "#cdb4db", "#a2d2ff"],
    "Monochrome Noir": ["#111111", "#3a3a3a", "#7a7a7a", "#c8c8c8", "#f5f5f5"],
    "Retro 70s": ["#e76f51", "#f4a261", "#e9c46a", "#2a9d8f", "#264653"],
}
HARMONIES = ["Complementary", "Analogous", "Triadic", "Split complementary", "Tetradic", "Monochrome"]
SOURCES = ["Mood board image (PNG)", "Base color + harmony", "Curated palette"]

gui_api.CloseAllWindows()  # a re-run must not leave the previous panel open

ui = {"slots": []}
state = {"palette": [], "palette_name": "", "variations": [], "kept": set()}


# ------------------------------------------------------------------ PNG in / out (pure Python, no PIL)
def png_read(path):
    """(width, height, rows of (r, g, b, a)) for 8-bit, non-interlaced PNG (gray, RGB, palette, +alpha)."""
    with open(path, "rb") as f:
        data = f.read()
    if data[:8] != b"\x89PNG\r\n\x1a\n":
        raise ValueError("not a PNG file")
    pos, idat, plte, trns = 8, [], None, None
    width = height = depth = ctype = interlace = None
    while pos < len(data):
        length, kind = struct.unpack(">I4s", data[pos:pos + 8])
        chunk = data[pos + 8:pos + 8 + length]
        pos += 12 + length
        if kind == b"IHDR":
            width, height, depth, ctype, _, _, interlace = struct.unpack(">IIBBBBB", chunk)
        elif kind == b"PLTE":
            plte = chunk
        elif kind == b"tRNS":
            trns = chunk
        elif kind == b"IDAT":
            idat.append(chunk)
        elif kind == b"IEND":
            break
    if depth != 8 or interlace:
        raise ValueError("only 8-bit non-interlaced PNG is supported (re-save the image as a plain PNG)")
    channels = {0: 1, 2: 3, 3: 1, 4: 2, 6: 4}[ctype]
    raw = zlib.decompress(b"".join(idat))
    stride = width * channels
    prev = bytearray(stride)
    rows, i = [], 0
    for _ in range(height):
        ftype = raw[i]
        line = bytearray(raw[i + 1:i + 1 + stride])
        i += 1 + stride
        if ftype == 1:
            for x in range(channels, stride):
                line[x] = (line[x] + line[x - channels]) & 255
        elif ftype == 2:
            for x in range(stride):
                line[x] = (line[x] + prev[x]) & 255
        elif ftype == 3:
            for x in range(stride):
                left = line[x - channels] if x >= channels else 0
                line[x] = (line[x] + ((left + prev[x]) >> 1)) & 255
        elif ftype == 4:
            for x in range(stride):
                a = line[x - channels] if x >= channels else 0
                b = prev[x]
                c = prev[x - channels] if x >= channels else 0
                p = a + b - c
                pa, pb, pc = abs(p - a), abs(p - b), abs(p - c)
                line[x] = (line[x] + (a if pa <= pb and pa <= pc else (b if pb <= pc else c))) & 255
        prev = line
        if ctype == 2:
            rows.append([(line[x], line[x + 1], line[x + 2], 255) for x in range(0, stride, 3)])
        elif ctype == 6:
            rows.append([(line[x], line[x + 1], line[x + 2], line[x + 3]) for x in range(0, stride, 4)])
        elif ctype == 0:
            rows.append([(v, v, v, 255) for v in line])
        elif ctype == 4:
            rows.append([(line[x], line[x], line[x], line[x + 1]) for x in range(0, stride, 2)])
        else:
            rows.append([(plte[3 * v], plte[3 * v + 1], plte[3 * v + 2], trns[v] if trns and v < len(trns) else 255) for v in line])
    return width, height, rows


def png_write(path, width, height, pixel):
    """Write an RGB PNG; pixel(x, y) -> (r, g, b)."""
    raw = bytearray()
    for y in range(height):
        raw.append(0)
        for x in range(width):
            raw.extend(pixel(x, y))

    def chunk(kind, body):
        return struct.pack(">I", len(body)) + kind + body + struct.pack(">I", zlib.crc32(kind + body) & 0xFFFFFFFF)

    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
                + chunk(b"IDAT", zlib.compress(bytes(raw), 6)) + chunk(b"IEND", b""))
    return path


# ------------------------------------------------------------------ color helpers
def hex_to_rgb(text):
    m = re.fullmatch(r"#?([0-9a-fA-F]{6})", text.strip())
    if not m:
        raise ValueError("use a hex color such as #3a7bd5")
    v = m.group(1)
    return tuple(int(v[i:i + 2], 16) for i in (0, 2, 4))


def rgb_to_hex(rgb):
    return "#%02x%02x%02x" % tuple(int(max(0, min(255, c))) for c in rgb)


def from_hsv(h, s, v):
    return tuple(int(round(c * 255)) for c in colorsys.hsv_to_rgb(h % 1.0, max(0.0, min(1.0, s)), max(0.0, min(1.0, v))))


def harmony(base, rule):
    """Five colors: the base, two accents on the harmony hues, a light neutral and a dark anchor.
    Each role has its own saturation / value target so a dark or pale base still gives a clean palette."""
    h, s, v = colorsys.rgb_to_hsv(*[c / 255.0 for c in base])
    hues = {"Complementary": (0.5, 0.5), "Analogous": (1 / 12, -1 / 12), "Triadic": (1 / 3, 2 / 3),
            "Split complementary": (5 / 12, 7 / 12), "Tetradic": (0.25, 0.5), "Monochrome": (0.0, 0.0)}[rule]
    vivid_s = max(0.55, s)
    if rule == "Monochrome":
        accents = [from_hsv(h, vivid_s * 0.7, 0.85), from_hsv(h, vivid_s, 0.55)]
    elif rule == "Complementary":
        accents = [from_hsv(h + 0.5, vivid_s, 0.85), from_hsv(h + 0.5, vivid_s * 0.6, 0.6)]
    else:
        accents = [from_hsv(h + hues[0], vivid_s, 0.85), from_hsv(h + hues[1], vivid_s * 0.85, 0.7)]
    light = from_hsv(h, 0.08, 0.96)
    dark = from_hsv(h, min(0.6, s), 0.2)
    return [tuple(base)] + accents + [light, dark]


def kmeans_palette(path, k=5, clusters=8):
    """Five distinct colors of an image. Plain k-means by pixel count picks the background (white walls,
    grey floors), so clusters are scored by size x chroma and picked greedily, skipping near duplicates."""
    width, height, rows = png_read(path)
    step = max(1, int((width * height / 6000.0) ** 0.5))
    pts = [px[:3] for y in range(0, height, step) for px in rows[y][::step] if px[3] > 128]
    if len(pts) < clusters:
        raise ValueError("the image has too few opaque pixels")
    rnd = random.Random(7)
    centers = rnd.sample(pts, clusters)
    for _ in range(12):
        groups = [[] for _ in range(clusters)]
        for p in pts:
            j = min(range(clusters), key=lambda c: (p[0] - centers[c][0]) ** 2 + (p[1] - centers[c][1]) ** 2 + (p[2] - centers[c][2]) ** 2)
            groups[j].append(p)
        centers = [tuple(sum(c[i] for c in g) / len(g) for i in range(3)) if g else centers[j] for j, g in enumerate(groups)]

    def chroma(c):
        return (max(c) - min(c)) / 255.0

    def far(a, b):
        return sum((a[i] - b[i]) ** 2 for i in range(3)) ** 0.5 > 48

    scored = sorted(((len(g) * (0.25 + chroma(c)), len(g), c) for g, c in zip(groups, centers) if g), reverse=True)
    picked = []
    for _, size, c in scored:
        if all(far(c, q) for _, q in picked):
            picked.append((size, c))
        if len(picked) == k:
            break
    for _, size, c in scored:  # top up with the remaining clusters if the image is nearly flat
        if len(picked) == k:
            break
        if all(c is not q for _, q in picked):
            picked.append((size, c))
    picked.sort(key=lambda sc: -sc[0])
    return [tuple(int(round(v)) for v in c) for _, c in picked]


def swatch(colors, path, width=300, height=36):
    colors = colors or [(80, 80, 80)]
    seg = float(width) / len(colors)
    return png_write(path, width, height, lambda x, y: colors[min(len(colors) - 1, int(x / seg))])


def crop_to_subject(src, dst, pad=12):
    """Crop a 3D thumbnail to the avatar and garment. The 3D background is a vertical gradient, so a pixel
    that differs from the left edge of its own row belongs to the subject."""
    w, h, rows = png_read(src)

    def differs(px, ref):
        return abs(px[0] - ref[0]) + abs(px[1] - ref[1]) + abs(px[2] - ref[2]) > 30

    ys = [y for y in range(0, h, 2) if any(differs(rows[y][x], rows[y][2]) for x in range(0, w, 3))]
    if not ys:
        return src
    xs = [x for x in range(0, w, 2) if any(differs(rows[y][x], rows[y][2]) for y in ys[::3])]
    top, bottom = max(0, ys[0] - pad), min(h, ys[-1] + pad)
    left, right = max(0, xs[0] - pad), min(w, xs[-1] + pad)
    return png_write(dst, right - left, bottom - top, lambda x, y: rows[top + y][left + x][:3])


# ------------------------------------------------------------------ garment model
def fabrics_by_area():
    """Fabric indices used by the garment, biggest total pattern area first."""
    area = {}
    for i in range(pattern_api.GetPatternCount()):
        f = pattern_api.GetPatternPieceFabricIndex(i)
        area[f] = area.get(f, 0.0) + max(0.0, pattern_api.GetPatternPieceArea(i))
    return [f for f, _ in sorted(area.items(), key=lambda kv: -kv[1])]


def assignment(palette, fabrics, variation):
    """fabric -> color. Variation 0 is the palette in order; later ones rotate it and swap the lead."""
    order = palette[variation % len(palette):] + palette[:variation % len(palette)]
    if variation // len(palette) % 2:
        order = order[::-1]
    return {f: order[i % len(order)] for i, f in enumerate(fabrics)}


# ------------------------------------------------------------------ actions
def status(text):
    gui_api.SetWidgetText(ui["status"], text)


def show_palette(colors, name):
    state["palette"], state["palette_name"] = colors, name
    gui_api.SetWidgetValue(ui["palette_img"], swatch(colors, WORK_DIR + "/palette.png", 420, 40))
    gui_api.SetWidgetText(ui["palette_text"], "%s: %s" % (name, "  ".join(rgb_to_hex(c) for c in colors)))


def build_palette():
    source = gui_api.GetWidgetSelectedIndex(ui["source"])
    try:
        if source == 0:
            path = gui_api.GetWidgetValue(ui["image"]).strip().strip('"')
            if not os.path.isfile(path):
                raise ValueError("image not found: %s" % path)
            status("extracting colors ...")
            show_palette(kmeans_palette(path), os.path.splitext(os.path.basename(path))[0])
        elif source == 1:
            rule = HARMONIES[gui_api.GetWidgetSelectedIndex(ui["harmony"])]
            base = gui_api.GetWidgetValue(ui["base"])
            show_palette(harmony(hex_to_rgb(base), rule), "%s %s" % (rule, rgb_to_hex(hex_to_rgb(base))))
        else:
            name = list(CURATED)[gui_api.GetWidgetSelectedIndex(ui["curated"])]
            show_palette([hex_to_rgb(c) for c in CURATED[name]], name)
        status("palette ready: press Generate")
        return True
    except (ValueError, OSError) as e:
        status(str(e))
        return False


def on_source(index):
    gui_api.SetWidgetVisible(ui["image_row"], index == 0)
    gui_api.SetWidgetVisible(ui["harmony_row"], index == 1)
    gui_api.SetWidgetVisible(ui["curated"], index == 2)
    if index != 0:
        build_palette()


def delete_unkept():
    for idx in sorted((v["colorway"] for n, v in enumerate(state["variations"]) if n not in state["kept"]), reverse=True):
        if idx > 0 and idx < utility_api.GetColorwayCount():
            utility_api.DeleteColorwayItem(idx)


def generate():
    try:
        _generate()
    except Exception as e:
        utility_api.SetCurrentColorwayIndex(0)
        status("Generate stopped: %s: %s" % (type(e).__name__, e))
        raise


def _generate():
    if not state["palette"] and not build_palette():
        return
    fabrics = fabrics_by_area()
    if not fabrics:
        status("open a garment first")
        return
    # replace the previous, not kept, round
    delete_unkept()
    state["variations"], state["kept"] = [], set()
    base = utility_api.GetCurrentColorwayIndex()
    count = int(float(gui_api.GetWidgetValue(ui["count"]) or SLOTS))
    true_color = gui_api.GetWidgetValue(ui["true_color"]).lower() in ("true", "1")
    stamp = datetime.datetime.now().strftime("%H%M%S")
    for n in range(min(SLOTS, max(1, count))):
        status("rendering variation %d / %d ..." % (n + 1, count))
        colors = assignment(state["palette"], fabrics, n)
        idx = utility_api.CopyColorway(base, 1)
        if idx <= 0:
            status("CLO refused to copy the colorway")
            break
        utility_api.SetColorwayName(idx, "Lab %s v%d" % (state["palette_name"][:24], n + 1))
        for f, rgb in colors.items():
            fabric_api.SetFabricPBRMaterialBaseColor(idx, f, 0, rgb[0] / 255.0, rgb[1] / 255.0, rgb[2] / 255.0, 1.0)
            if true_color:
                utility_api.SetBaseTextureMapImageDesaturation(f, idx, True)
        utility_api.SetCurrentColorwayIndex(idx)
        raw = export_api.ExportThumbnail3D("%s/v%d_%s_raw.png" % (WORK_DIR, n + 1, stamp))
        shot = raw
        if raw:
            try:  # cropping only makes the gallery nicer: never let it stop the batch
                shot = crop_to_subject(raw, "%s/v%d_%s.png" % (WORK_DIR, n + 1, stamp))
            except Exception as e:
                state["last_error"] = "crop skipped: %s: %s" % (type(e).__name__, e)
        strip = swatch([colors[f] for f in fabrics], "%s/v%d_%s_strip.png" % (WORK_DIR, n + 1, stamp), 220, 18)
        state["variations"].append({"colorway": idx, "colors": colors, "image": shot, "strip": strip})
        slot = ui["slots"][n]
        gui_api.SetWidgetValue(slot["image"], shot)
        gui_api.SetWidgetValue(slot["strip"], strip)
        gui_api.SetWidgetText(slot["label"], "v%d  %s" % (n + 1, "  ".join(rgb_to_hex(colors[f]) for f in fabrics)))
        gui_api.SetWidgetText(slot["keep"], "Keep")
        gui_api.SetWidgetVisible(slot["group"], True)
    for n in range(len(state["variations"]), SLOTS):
        gui_api.SetWidgetVisible(ui["slots"][n]["group"], False)
    utility_api.SetCurrentColorwayIndex(base)
    note = ("  (" + state.pop("last_error") + ")") if state.get("last_error") else ""
    status("%d variation(s) of %s: Keep the ones you like, then Finish%s" % (len(state["variations"]), state["palette_name"], note))


def toggle_keep(n):
    if n >= len(state["variations"]):
        return
    if n in state["kept"]:
        state["kept"].discard(n)
    else:
        state["kept"].add(n)
    gui_api.SetWidgetText(ui["slots"][n]["keep"], "Kept" if n in state["kept"] else "Keep")
    status("%d kept" % len(state["kept"]))


def view(n):
    if n < len(state["variations"]):
        utility_api.SetCurrentColorwayIndex(state["variations"][n]["colorway"])


def finish():
    if not state["variations"]:
        status("generate first")
        return
    if not state["kept"]:
        if not gui_api.AskConfirm("Colorway Lab", "Nothing is kept. Delete all generated colorways?", "Delete", "Cancel"):
            return
    kept = sorted(state["kept"])
    delete_unkept()
    # indices of kept colorways shifted down by the deletions in front of them
    survivors = sorted(state["variations"][n]["colorway"] for n in kept)
    for rank, idx in enumerate(survivors):
        current = idx - sum(1 for n, v in enumerate(state["variations"]) if n not in state["kept"] and v["colorway"] < idx)
        utility_api.SetColorwayName(current, "%s %d" % (state["palette_name"][:28], rank + 1))
    utility_api.SetCurrentColorwayIndex(0)
    for slot in ui["slots"]:
        gui_api.SetWidgetVisible(slot["group"], False)
    state["finished"] = [state["variations"][n] for n in kept]
    state["variations"], state["kept"] = [], set()
    status("%d colorway(s) added to the project" % len(kept))


def lookbook():
    picks = [state["variations"][n] for n in sorted(state["kept"])] or state.get("finished") or state["variations"]
    if not picks:
        status("generate (and keep) some variations first")
        return
    status("composing the lookbook ...")
    tiles = []
    for v in picks:
        w, h, rows = png_read(v["image"])  # already cropped to the subject
        tiles.append((w, rows, [v["colors"][f] for f in fabrics_by_area()]))
    tile_w = max(t[0] for t in tiles)
    tile_h = max(len(t[1]) for t in tiles)
    strip_h, gap = 24, 12
    width = len(tiles) * tile_w + (len(tiles) + 1) * gap
    height = tile_h + strip_h + 3 * gap

    def pixel(x, y):
        col = (x - gap) // (tile_w + gap)
        tx = (x - gap) % (tile_w + gap)
        if col < 0 or col >= len(tiles) or tx >= tile_w or y < gap:
            return (245, 245, 245)
        w, rows, colors = tiles[col]
        ty = y - gap
        if ty < len(rows) and tx < w:
            r, g, b, a = rows[ty][tx]
            k = a / 255.0
            return (int(r * k + 245 * (1 - k)), int(g * k + 245 * (1 - k)), int(b * k + 245 * (1 - k)))
        sy = y - (gap * 2 + tile_h)
        if 0 <= sy < strip_h:
            return colors[min(len(colors) - 1, tx * len(colors) // tile_w)]
        return (245, 245, 245)

    os.makedirs(OUT_DIR, exist_ok=True)
    path = "%s/%s_lookbook_%s.png" % (OUT_DIR, (utility_api.GetProjectName() or "untitled").replace(" ", "_"),
                                      datetime.datetime.now().strftime("%Y%m%d_%H%M%S"))
    png_write(path, width, height, pixel)
    gui_api.SetWidgetValue(ui["lookbook_img"], path)
    gui_api.SetWidgetVisible(ui["lookbook_img"], True)
    status("lookbook: %s" % path)


# ------------------------------------------------------------------ layout
blank = swatch([(60, 60, 60)], WORK_DIR + "/blank.png", 220, 18)
panel = gui_api.CreateDockPanel("Colorway Lab", "left", "colorway_lab_left")

src = gui_api.AddGroup(panel, "Palette")
ui["source"] = gui_api.AddDropdown(src, SOURCES, 2, on_source)
ui["image_row"] = gui_api.AddRow(src)
ui["image"] = gui_api.AddTextInput(ui["image_row"], "", "mood board PNG path")
gui_api.AddButton(ui["image_row"], "Extract", build_palette)
ui["harmony_row"] = gui_api.AddRow(src)
ui["base"] = gui_api.AddTextInput(ui["harmony_row"], "#3a7bd5", "#rrggbb", lambda v: build_palette() if re.fullmatch(r"#?[0-9a-fA-F]{6}", v.strip()) else None)
ui["harmony"] = gui_api.AddDropdown(ui["harmony_row"], HARMONIES, 0, lambda i: build_palette())
ui["curated"] = gui_api.AddDropdown(src, list(CURATED), 0, lambda i: build_palette())
ui["palette_img"] = gui_api.AddImage(src, swatch([(60, 60, 60)], WORK_DIR + "/palette.png", 420, 40), 420, 40)
ui["palette_text"] = gui_api.AddLabel(src, "")

opts = gui_api.AddRow(panel)
gui_api.AddLabel(opts, "Variations")
ui["count"] = gui_api.AddNumberInput(opts, 6, 1, SLOTS, 0)
ui["true_color"] = gui_api.AddCheckBox(opts, "True color (desaturate textures)", True)

actions = gui_api.AddRow(panel)
gui_api.AddButton(actions, "Generate", generate)
gui_api.AddButton(actions, "Finish", finish)
gui_api.AddButton(actions, "Lookbook", lookbook)
ui["status"] = gui_api.AddLabel(panel, "pick a palette, then Generate")

gallery = gui_api.AddGroup(panel, "Variations")
row = None
for n in range(SLOTS):
    if n % 3 == 0:
        row = gui_api.AddRow(gallery)
    group = gui_api.AddGroup(row, "v%d" % (n + 1))
    image = gui_api.AddImage(group, blank, 200, 300)
    strip = gui_api.AddImage(group, blank, 200, 16)
    label = gui_api.AddLabel(group, "")
    buttons = gui_api.AddRow(group)
    keep = gui_api.AddButton(buttons, "Keep", lambda n=n: toggle_keep(n))
    gui_api.AddButton(buttons, "View", lambda n=n: view(n))
    ui["slots"].append({"group": group, "image": image, "strip": strip, "label": label, "keep": keep})
    gui_api.SetWidgetVisible(group, False)
ui["lookbook_img"] = gui_api.AddImage(panel, blank, 640, 300)
gui_api.SetWidgetVisible(ui["lookbook_img"], False)

on_source(2)
gui_api.ShowWindow(panel)
# The script returns here. The panel stays docked and all buttons keep working.