"""Erzeugt die Körperteil-Masken („Rig“) für die freigestellten Figuren.

Für jede Figur ist unten ein Skelett aus Knochenlinien (in Bildpixeln) definiert.
Jedes Pixel der Figur wird genau einem Teil zugeordnet:
  1. Accessoires (Schal, Mütze, Fäustlinge, Stiefel …) über ihre Farbe,
  2. alles andere über den kürzesten Weg *innerhalb der Figur* zum nächsten
     Knochen. Der Weg führt nicht über Lücken, deshalb greift z. B. Lulus
     Schwanz nie ins Bein daneben.

Ausgabe je Figur unter public/figuren/<figur>/:
  teile/<teil>.png    harte Maske des Teils (weiß = gehört dazu)
  einfluss/<teil>.png weiche Maske für die Verformung: das Teil plus alles,
                      was daran hängt, mit weichem Übergang am Gelenk
  uebersicht.png      alle Teile farbig, zum Prüfen
Dazu src/rigs.ts mit Gelenkpunkten und Hierarchie.

Aufruf: python3 werkzeuge/rig.py   (braucht numpy, scipy, pillow)
"""

import colorsys
import json
import os
import sys

import numpy as np
from PIL import Image, ImageDraw, ImageFilter
from scipy import ndimage

HIER = os.path.dirname(os.path.abspath(__file__))
FIGUREN = os.path.join(HIER, "..", "public", "figuren")
RIGS_TS = os.path.join(HIER, "..", "src", "rigs.ts")

# Knochen: name, eltern, von (Gelenk), bis. weich = Länge des Übergangs am Gelenk in px.
# Accessoire-Gruppen: Farbe + Bereich; ihre Teile werden über Startlinien getrennt.

FIGUREN_DEF = {
    "lulu": {
        "knochen": [
            ("rumpf", None, (130, 292), (232, 296), 0,
             [(120, 245), (170, 238), (230, 250), (248, 290), (245, 345), (200, 352), (150, 347), (105, 337), (95, 300)]),
            ("hals", "rumpf", (124, 244), (112, 156), 14,
             [(78, 150), (132, 140), (142, 178), (128, 205), (85, 205), (72, 176)]),
            ("kopf", "hals", (111, 150), (86, 72), 10,
             [(10, 60), (40, 45), (60, 47), (80, 25), (110, 30), (140, 52), (146, 90), (136, 135), (126, 152),
              (95, 160), (60, 156), (25, 140), (14, 110)]),
            ("ohr_links", "kopf", (44, 57), (12, 25), 5,
             [(0, 10), (30, 12), (52, 40), (52, 62), (40, 70), (20, 60), (4, 40)]),
            ("ohr_rechts", "kopf", (146, 72), (186, 56), 5,
             [(140, 56), (165, 42), (196, 45), (193, 62), (170, 78), (146, 88), (140, 78)]),
            ("horn_links", "kopf", (65, 44), (62, 8), 4,
             [(55, 0), (78, 0), (78, 30), (73, 44), (58, 44), (56, 30)]),
            ("horn_rechts", "kopf", (122, 46), (126, 10), 4,
             [(112, 8), (137, 8), (137, 28), (129, 48), (117, 48), (112, 30)]),
            ("bein_vorne_links", "rumpf", (92, 332), (80, 465), 12,
             [(95, 320), (106, 335), (100, 400), (95, 445), (93, 472), (62, 472), (70, 440), (72, 400), (78, 340)]),
            ("bein_vorne_rechts", "rumpf", (124, 340), (127, 475), 12,
             [(108, 338), (141, 335), (140, 400), (138, 470), (141, 487), (110, 487), (115, 440), (112, 380)]),
            ("bein_hinten_links", "rumpf", (186, 345), (193, 448), 12,
             [(168, 342), (200, 346), (205, 400), (206, 460), (178, 460), (182, 420), (170, 380)]),
            ("bein_hinten_rechts", "rumpf", (230, 335), (235, 470), 12,
             [(212, 332), (243, 324), (250, 360), (248, 420), (251, 481), (222, 481), (225, 430), (215, 380)]),
            ("schwanz", "rumpf", (246, 288), (266, 364), 6,
             [(244, 282), (256, 300), (262, 330), (277, 350), (273, 373), (256, 373), (254, 340), (244, 300)]),
        ],
        "accessoires": [
            {
                "farbe": "gruen",
                "bereich": None,
                "teile": [
                    ("schal_schlinge", "hals", [(70, 228), (160, 205)], (124, 244), 10),
                    ("schal_ende", "schal_schlinge", [(40, 292), (45, 372)], (62, 252), 10),
                    ("schal_boden", "rumpf", [(230, 420), (350, 440)], (160, 360), 10),
                    ("schal_boden", None, [(160, 362), (190, 388)], None, 0),
                ],
            },
            {
                "farbe": "rot",
                "bereich": (283, 378, 318, 404),
                "teile": [("schal_boden", None, [(300, 392)], None, 0)],
            },
        ],
    },
    "pino": {
        "knochen": [
            ("koerper", None, (230, 430), (400, 560), 0,
             [(20, 350), (150, 200), (245, 200), (262, 300), (275, 445), (300, 560), (200, 640), (60, 560)]),
            ("kopf", "koerper", (400, 465), (420, 320), 26,
             [(245, 150), (300, 115), (420, 120), (520, 180), (554, 240), (554, 455), (470, 465), (380, 462),
              (300, 450), (268, 400), (262, 300), (240, 240)]),
            ("ohr_links", "kopf", (290, 258), (256, 222), 6, "kreis", (272, 240, 27)),
            ("ohr_rechts", "kopf", (510, 270), (538, 248), 6, "kreis", (525, 258, 21)),
            ("bauch", "koerper", (330, 600), (430, 470), 20,
             [(300, 450), (470, 460), (485, 520), (480, 600), (470, 640), (400, 645), (360, 610), (300, 545)]),
            ("arm_links", "koerper", (322, 462), (370, 540), 10,
             [(322, 456), (345, 470), (380, 505), (392, 540), (372, 556), (345, 548), (330, 515), (318, 480)]),
            ("arm_rechts", "koerper", (458, 468), (455, 545), 10,
             [(445, 465), (472, 470), (478, 520), (470, 556), (442, 556), (440, 510)]),
            ("fuss_links", "koerper", (310, 568), (320, 645), 10,
             [(305, 565), (355, 563), (360, 600), (350, 655), (315, 658), (305, 620)]),
            ("fuss_rechts", "koerper", (495, 600), (540, 575), 10,
             [(495, 560), (540, 548), (554, 575), (545, 640), (505, 640), (490, 600)]),
        ],
        "accessoires": [
            {
                "farbe": "weiss",
                "bereich": (360, 0, 500, 92),
                "teile": [("bommel", "muetze", [(445, 45)], (445, 88), 8)],
            },
            {
                "farbe": "blau",
                "bereich": (250, 60, 554, 215),
                "teile": [("muetze", "kopf", [(330, 170), (480, 170)], (400, 200), 10)],
            },
        ],
    },
    "noa": {
        "knochen": [
            ("rumpf", None, (115, 165), (115, 320), 0,
             [(80, 150), (160, 150), (165, 200), (170, 320), (75, 325), (82, 200)]),
            ("kopf", "rumpf", (120, 156), (112, 40), 12,
             [(55, 15), (115, 3), (165, 15), (190, 45), (197, 90), (182, 108), (170, 135), (140, 152), (100, 155),
              (70, 140), (48, 118), (32, 95), (30, 60)]),
            ("zopf_links", "kopf", (48, 120), (15, 135), 6,
             [(0, 120), (20, 108), (45, 112), (52, 130), (45, 150), (20, 155), (5, 145)]),
            ("zopf_rechts", "kopf", (178, 108), (200, 126), 6,
             [(172, 100), (195, 95), (210, 108), (212, 130), (195, 148), (176, 140)]),
            ("arm_links", "rumpf", (64, 172), (52, 288), 12,
             [(35, 180), (62, 165), (78, 180), (80, 290), (60, 300), (35, 295), (28, 240)]),
            ("arm_rechts", "rumpf", (166, 172), (183, 272), 12,
             [(160, 170), (180, 180), (192, 230), (195, 275), (178, 280), (165, 240), (158, 190)]),
            ("bein_links", "rumpf", (92, 322), (92, 350), 6,
             [(78, 320), (108, 320), (108, 352), (80, 352)]),
            ("bein_rechts", "rumpf", (128, 322), (132, 350), 6,
             [(110, 320), (142, 320), (142, 352), (112, 352)]),
        ],
        "accessoires": [
            {
                "farbe": "rot",
                "bereich": (0, 260, 219, 340),
                "teile": [
                    ("hand_links", "arm_links", [(55, 310)], (54, 292), 6),
                    ("hand_rechts", "arm_rechts", [(188, 292)], (186, 275), 6),
                ],
            },
            {
                "farbe": "rot",
                "bereich": (0, 340, 219, 429),
                "teile": [
                    ("stiefel_links", "bein_links", [(92, 362), (95, 415)], (92, 350), 6),
                    ("stiefel_rechts", "bein_rechts", [(140, 362), (150, 410)], (132, 350), 6),
                ],
            },
        ],
    },
}


def farbmaske(rgb, art):
    h, w, _ = rgb.shape
    hsv = np.array([colorsys.rgb_to_hsv(*(p / 255)) for p in rgb.reshape(-1, 3)]).reshape(h, w, 3)
    hue, sat, val = hsv[..., 0] * 360, hsv[..., 1], hsv[..., 2]
    if art == "gruen":
        m = (hue > 52) & (hue < 180) & (sat > 0.06) & (val > 0.12)
    elif art == "blau":
        m = (hue > 185) & (hue < 260) & (sat > 0.12)
    elif art == "rot":
        m = ((hue < 22) | (hue > 335)) & (sat > 0.35) & (val > 0.3)
    elif art == "weiss":
        m = (sat < 0.18) & (val > 0.72)
    else:
        raise ValueError(art)
    return m


def linie_pixel(punkte, form):
    bild = Image.new("L", (form[1], form[0]), 0)
    d = ImageDraw.Draw(bild)
    if len(punkte) == 1:
        x, y = punkte[0]
        d.ellipse((x - 2, y - 2, x + 2, y + 2), fill=1)
    else:
        d.line(punkte, fill=1, width=3)
    return np.asarray(bild).astype(bool)


def flaechen_pixel(flaeche, form):
    bild = Image.new("L", (form[1], form[0]), 0)
    d = ImageDraw.Draw(bild)
    if flaeche[0] == "kreis":
        x, y, r = flaeche[1]
        d.ellipse((x - r, y - r, x + r, y + r), fill=1)
    else:
        d.polygon(flaeche[0], fill=1)
    return np.asarray(bild).astype(bool)


def ausbreiten(labels, erlaubt):
    """Wachsen der Labels innerhalb von `erlaubt`, Pixel für Pixel (Weg innerhalb der Figur)."""
    labels = labels.copy()
    kreuz = ndimage.generate_binary_structure(2, 2)
    while True:
        frei = erlaubt & (labels == 0)
        if not frei.any():
            break
        gewachsen = ndimage.grey_dilation(labels, footprint=kreuz)
        neu = frei & (gewachsen > 0)
        if not neu.any():
            break
        labels[neu] = gewachsen[neu]
    return labels


def rampe_entlang(form, von, bis, weich):
    """0 hinter dem Gelenk, 1 nach `weich` px in Knochenrichtung."""
    if weich <= 0:
        return np.ones(form)
    yy, xx = np.mgrid[0 : form[0], 0 : form[1]].astype(float)
    dx, dy = bis[0] - von[0], bis[1] - von[1]
    l = np.hypot(dx, dy)
    t = ((xx - von[0]) * dx + (yy - von[1]) * dy) / l
    s = np.clip((t + weich * 0.4) / weich, 0, 1)
    return s * s * (3 - 2 * s)


def baue(figur, d):
    bild = Image.open(os.path.join(FIGUREN, f"{figur}.png")).convert("RGBA")
    arr = np.asarray(bild).astype(float)
    rgb, alpha = arr[..., :3], arr[..., 3]
    figur_maske = alpha > 40
    form = figur_maske.shape

    namen = []
    info = {}  # name -> (eltern, gelenk, bis, weich)
    for name, eltern, von, bis, weich, *_ in d["knochen"]:
        namen.append(name)
        info[name] = (eltern, von, bis, weich)
    for gruppe in d["accessoires"]:
        for name, eltern, linie, gelenk, weich in gruppe["teile"]:
            if name not in info:
                namen.append(name)
                info[name] = (eltern, gelenk, linie[-1], weich)
    nummer = {n: i + 1 for i, n in enumerate(namen)}

    labels = np.zeros(form, dtype=np.int32)

    # 1. Accessoires über Farbe
    for gruppe in d["accessoires"]:
        m = figur_maske & farbmaske(rgb, gruppe["farbe"])
        if gruppe["bereich"]:
            x0, y0, x1, y1 = gruppe["bereich"]
            b = np.zeros(form, bool)
            b[y0:y1, x0:x1] = True
            m &= b
        m = ndimage.binary_opening(m, iterations=1)
        m = ndimage.binary_closing(m, iterations=2) & figur_maske
        start = np.zeros(form, np.int32)
        for name, _, linie, _, _ in gruppe["teile"]:
            start[linie_pixel(linie, form) & m] = nummer[name]
        # Flächen ohne Startlinie (z. B. Schalstück zwischen zwei Beinen) bekommen ab einer
        # Mindestgröße das nächstgelegene Teil der Gruppe, kleine Farbflecken fallen weg.
        lab, n = ndimage.label(m)
        verbunden = set(np.unique(lab[start > 0])) - {0}
        _, (iy, ix) = ndimage.distance_transform_edt(start == 0, return_indices=True)
        for k in range(1, n + 1):
            if k in verbunden:
                continue
            flaeche = lab == k
            if flaeche.sum() >= 40:
                ys, xs = np.where(flaeche)
                y0, x0 = ys[len(ys) // 2], xs[len(xs) // 2]
                start[flaeche] = start[iy[y0, x0], ix[y0, x0]]
                verbunden.add(k)
        behalten = np.isin(lab, list(verbunden)) & (lab > 0)
        behalten = ndimage.binary_fill_holes(behalten) & figur_maske  # z. B. Blümchen im Schal
        gruppe_labels = ausbreiten(start, behalten)
        frei = labels == 0
        labels[frei & (gruppe_labels > 0)] = gruppe_labels[frei & (gruppe_labels > 0)]

    # 2. Körper: Kernflächen der Teile (spätere Teile überschreiben frühere), dazu die Knochenlinie
    start = np.zeros(form, np.int32)
    rest = figur_maske & (labels == 0)
    for name, _, von, bis, _, *flaeche in d["knochen"]:
        kern = linie_pixel([von, bis], form)
        if flaeche:
            kern |= flaechen_pixel(flaeche, form)
        start[kern & rest] = nummer[name]
    for name in [k[0] for k in d["knochen"]]:
        if not (start == nummer[name]).any():
            print(f"  Warnung: Knochen {name} liegt nicht auf der Figur", file=sys.stderr)
    koerper = ausbreiten(start, rest)
    labels[rest] = koerper[rest]
    # 3. abgeschnittene Inseln über die ganze Figur auffüllen
    labels = ausbreiten(labels, figur_maske)

    ziel = os.path.join(FIGUREN, figur)
    os.makedirs(os.path.join(ziel, "teile"), exist_ok=True)
    os.makedirs(os.path.join(ziel, "einfluss"), exist_ok=True)

    kinder = {n: [m for m in namen if info[m][0] == n] for n in namen}

    def unterbaum(n):
        return [n] + [x for k in kinder[n] for x in unterbaum(k)]

    # Für die Verformung reicht jedes Teil über die Kontur hinaus: Punkte außerhalb der Figur
    # bekommen das Teil des nächsten Figurpixels. Sonst bleiben Gitterpunkte am Rand stehen,
    # während die Figur sich dreht, und die Randpixel verschmieren.
    _, (iy, ix) = ndimage.distance_transform_edt(~figur_maske, return_indices=True)
    labels_aussen = labels[iy, ix]

    for name in namen:
        teil = labels == nummer[name]
        Image.fromarray((teil * 255).astype("uint8")).save(os.path.join(ziel, "teile", f"{name}.png"))
        eltern, gelenk, bis, weich = info[name]
        baum = np.isin(labels_aussen, [nummer[x] for x in unterbaum(name)])
        weich_bild = Image.fromarray((baum * 255).astype("uint8")).filter(ImageFilter.GaussianBlur(1.5))
        e = np.asarray(weich_bild).astype(float) / 255
        if gelenk is not None:
            e *= rampe_entlang(form, gelenk, bis, weich)
        Image.fromarray((e * 255).astype("uint8")).save(os.path.join(ziel, "einfluss", f"{name}.png"))

    # Übersicht
    rng = np.random.default_rng(3)
    farben = rng.integers(60, 255, size=(len(namen) + 1, 3))
    farben[0] = (30, 30, 60)
    uebersicht = farben[labels].astype(float)
    uebersicht = uebersicht * 0.55 + rgb * 0.45
    uebersicht[~figur_maske] = (30, 30, 60)
    ueb = Image.fromarray(uebersicht.astype("uint8"))
    dr = ImageDraw.Draw(ueb)
    for name in namen:
        ys, xs = np.where(labels == nummer[name])
        if len(xs):
            dr.text((int(xs.mean()) - 10, int(ys.mean()) - 4), name, fill=(255, 255, 255))
    ueb.save(os.path.join(ziel, "uebersicht.png"))

    groessen = {n: int((labels == nummer[n]).sum()) for n in namen}
    print(figur, groessen)
    return {
        n: {"eltern": info[n][0], "um": list(info[n][1]) if info[n][1] else None}
        for n in namen
    }


def main():
    rigs = {figur: baue(figur, d) for figur, d in FIGUREN_DEF.items()}
    with open(RIGS_TS, "w") as f:
        f.write("// Automatisch erzeugt von werkzeuge/rig.py – nicht von Hand ändern.\n")
        f.write("// Gelenkpunkte (um) in Bildpixeln der Figur, eltern = woran das Teil hängt.\n")
        f.write("// Masken: figuren/<figur>/teile/<teil>.png und figuren/<figur>/einfluss/<teil>.png\n\n")
        f.write("export type RigTeil = {eltern: string | null; um: [number, number] | null};\n\n")
        daten = json.dumps(rigs, indent="\t", ensure_ascii=False)
        f.write(f"export const RIGS: Record<string, Record<string, RigTeil>> = {daten};\n")


if __name__ == "__main__":
    main()