#!/usr/bin/env python3
"""
Watch face vector reconstruction, v2.

Changes over v1:
  * ring ticks are annular sectors (constant angular width) instead of rects
  * every tick's colour is sampled directly from the source, no LUT interpolation
  * glyph alpha is normalised against the *local* fill colour, so vertically
    graded digits keep their true weight in the dark rows
  * time digits get multi-stop gradients sampled every 8 px
  * sub-pixel geometry throughout
"""
import json, math, os
import numpy as np
from PIL import Image
from skimage import measure

SRC = "/mnt/user-data/uploads"
OUT = "/mnt/user-data/outputs"
os.makedirs(OUT, exist_ok=True)

CX_IX, CY_IX = 250.038, 250.049      # dial centre, pixel-index space
CX, CY = CX_IX + 0.5, CY_IX + 0.5     # same point in SVG user space
PX = 0.5                              # index -> SVG offset

THEMES = [
    ("0fbde6715a76f357bcc5c20cff74c0bf2f156e89.png", "01", "Ember"),
    ("304a904845dbc44aa0c87eecad2dd3a2621d2ffc.png", "02", "Aurora"),
    ("a61e76acdc7f17edfda3b11d36ba5d1c25e0739b.png", "03", "Brass"),
    ("f43ad1228487aa0261183e7268277f0e9ab328f4.png", "04", "Voltage"),
    ("8fd4a2811dc0227b004a5832479cc0a5d4e26a9f.png", "05", "Reef"),
    ("f087f3345e5036bfd406719e8d7517c13109d998.png", "06", "Acid"),
    ("d3a77eea386551ff09b352a77f522bb3c68e925d.png", "07", "Kelp"),
]

# ---------------------------------------------------------- ring geometry ---
R_IN, R_OUT = 231.09, 249.10

TOP_PHASE, TOP_PITCH, TOP_W, TOP_N = 6.500, 4.000, 3.108, 21
BOT_PHASE, BOT_PITCH, BOT_W, BOT_N = 6.005, 3.000, 2.086, 21

DOT_R, DOT_RAD = 240.10, 9.00
DOT_ANGLES = [0.0, 93.15, 180.0, 266.85]

BAND_Y, BAND_H = 286.30, 44.19
COLON = [(237.88, 183.68, 24.24, 21.30), (237.88, 234.57, 24.24, 21.30)]


def tick_angles():
    """(group, signed centre angle) for all 86 ticks."""
    out = []
    for k in range(TOP_N):
        a = TOP_PHASE + TOP_PITCH * k
        out.append(("top-right", a, TOP_W))
        out.append(("top-left", -a, TOP_W))
    for k in range(BOT_N):
        a = BOT_PHASE + BOT_PITCH * k
        out.append(("bottom-left", 180.0 + a, BOT_W))
        out.append(("bottom-right", 180.0 - a, BOT_W))
    return out


def polar(ang, r):
    t = math.radians(ang - 90.0)
    return CX + math.cos(t) * r, CY + math.sin(t) * r


def polar_ix(ang, r):
    t = math.radians(ang - 90.0)
    return CX_IX + math.cos(t) * r, CY_IX + math.sin(t) * r


def sector_path(centre, width, r0=R_IN, r1=R_OUT):
    a0, a1 = centre - width / 2, centre + width / 2
    x1, y1 = polar(a0, r1)
    x2, y2 = polar(a1, r1)
    x3, y3 = polar(a1, r0)
    x4, y4 = polar(a0, r0)
    return (f"M{x1:.2f},{y1:.2f} A{r1:.2f},{r1:.2f} 0 0 1 {x2:.2f},{y2:.2f} "
            f"L{x3:.2f},{y3:.2f} A{r0:.2f},{r0:.2f} 0 0 0 {x4:.2f},{y4:.2f} Z")


def sample_tick(arr, centre, width):
    """Median colour of a tick's interior (AA edges excluded)."""
    px = []
    for r in np.arange(R_IN + 3.5, R_OUT - 3.5, 1.0):
        for d in np.arange(-width * 0.32, width * 0.32 + 1e-9, width * 0.16):
            x, y = polar_ix(centre + d, r)
            px.append(arr[int(round(y)), int(round(x))])
    return np.median(np.array(px), axis=0)


# --------------------------------------------------------------- tracing ---
def rdp(pts, eps):
    if len(pts) < 3:
        return pts
    a, b = pts[0], pts[-1]
    ab = b - a
    n = math.hypot(ab[0], ab[1])
    if n < 1e-9:
        d = np.hypot(*(pts - a).T)
    else:
        d = np.abs(ab[0] * (pts[:, 1] - a[1]) - ab[1] * (pts[:, 0] - a[0])) / n
    i = int(np.argmax(d))
    if d[i] > eps:
        return np.vstack([rdp(pts[: i + 1], eps)[:-1], rdp(pts[i:], eps)])
    return np.vstack([a, b])


def trace(alpha, x0, y0, eps=0.07, ss=6):
    half = 0.5 / ss   # supersampled index -> source continuous coordinate
    up = np.asarray(Image.fromarray((np.clip(alpha, 0, 1) * 255).astype(np.uint8))
                    .resize((alpha.shape[1] * ss, alpha.shape[0] * ss), Image.BICUBIC)) / 255.0
    up = np.pad(up, 1, mode="constant")
    parts = []
    for c in measure.find_contours(up, 0.5):
        c = (c - 1) / ss
        pts = np.column_stack([c[:, 1] + x0 + half, c[:, 0] + y0 + half])
        if len(pts) < 10:
            continue
        pts = rdp(pts, eps)
        if len(pts) < 4:
            continue
        parts.append("M" + " ".join(f"{x:.2f},{y:.2f}" for x, y in pts[:-1]) + "Z")
    return " ".join(parts)


def full_colour(reg, q=0.97):
    """Colour of a fully-opaque pixel in a solid-filled region."""
    flat = reg.reshape(-1, 3)
    l = flat.mean(axis=1)
    cut = np.quantile(l, q)
    sel = flat[l >= cut]
    return sel.mean(axis=0)


def alpha_solid(arr, box):
    x0, y0, x1, y1 = box
    reg = arr[y0:y1, x0:x1].astype(float)
    e = full_colour(reg)
    a = reg.mean(axis=2) / max(e.mean(), 1e-6)
    return np.clip(a, 0, 1), e, x0, y0


def row_ramp(reg):
    """Per-row fully-opaque colour for a vertically graded glyph."""
    h = reg.shape[0]
    ys, cols = [], []
    for y in range(h):
        row = reg[y]
        l = row.mean(axis=1)
        if l.max() < 30:
            continue
        sel = row[l >= l.max() * 0.92]
        if len(sel) < 3:
            continue
        ys.append(y)
        cols.append(sel.mean(axis=0))
    if len(ys) < 4:
        return None, None
    ys = np.array(ys)
    cols = np.array(cols)
    # smooth then extend to every row
    out = np.zeros((h, 3))
    for c in range(3):
        out[:, c] = np.interp(np.arange(h), ys, np.convolve(
            cols[:, c], np.ones(5) / 5, mode="same") if len(ys) > 8 else cols[:, c])
        out[: ys[0], c] = cols[0, c]
        out[ys[-1] + 1:, c] = cols[-1, c]
    return out, (ys[0], ys[-1])


def alpha_graded(arr, box):
    x0, y0, x1, y1 = box
    reg = arr[y0:y1, x0:x1].astype(float)
    ramp, span = row_ramp(reg)
    if ramp is None:
        return alpha_solid(arr, box) + (None,)
    denom = np.maximum(ramp.mean(axis=1), 1e-6)[:, None]
    a = np.clip(reg.mean(axis=2) / denom, 0, 1)
    return a, ramp, x0, y0, span


def alpha_on_fill(arr, box, fill):
    x0, y0, x1, y1 = box
    reg = arr[y0:y1, x0:x1].astype(float)
    fill = np.array(fill, float)
    d = reg - fill
    far = d.reshape(-1, 3)
    mag = np.linalg.norm(far, axis=1)
    e = far[mag >= np.quantile(mag, 0.97)].mean(axis=0)
    a = (d @ e) / max(e @ e, 1e-6)
    return np.clip(a, 0, 1), fill + e, x0, y0


def hx(c):
    return "#%02X%02X%02X" % tuple(int(round(min(255, max(0, v)))) for v in c)


# --------------------------------------------------------------- elements ---
SOLID = [
    ("battery-icon",  (188, 38, 228, 68)),
    ("battery-value", (228, 38, 312, 70)),
    ("sunrise-icon",  (136, 72, 178, 108)),
    ("sunrise-value", (106, 110, 208, 152)),
    ("weather-icon",  (322, 70, 364, 108)),
    ("weather-value", (258, 110, 428, 152)),
    ("dist-value",    (166, 345, 240, 388)),
    ("dist-label",    (143, 390, 238, 432)),
    ("steps-value",   (261, 345, 366, 388)),
    ("steps-label",   (263, 390, 386, 432)),
]
BAND_EL = [
    ("bt-icon",  (20, 289, 54, 326)),
    ("date-day", (132, 290, 223, 326)),
    ("date-md",  (225, 290, 370, 326)),
    ("meridiem", (404, 290, 472, 326)),
]
TIME_L = (76, 162, 218, 278)
TIME_R = (264, 162, 434, 278)


def grad_def(gid, ramp, y0, span):
    ya, yb = y0 + span[0] + 0.5, y0 + span[1] + 0.5
    n = max(2, int((span[1] - span[0]) // 8))
    stops = []
    for i in range(n + 1):
        f = i / n
        yy = int(round(span[0] + f * (span[1] - span[0])))
        stops.append(f'<stop offset="{f:.3f}" stop-color="{hx(ramp[yy])}"/>')
    return (f'<linearGradient id="{gid}" x1="0" y1="{ya}" x2="0" y2="{yb}" '
            f'gradientUnits="userSpaceOnUse">{"".join(stops)}</linearGradient>')


def build(path, tid, tname):
    arr = np.array(Image.open(os.path.join(SRC, path)).convert("RGB")).astype(int)
    defs, parts, tok = [], [], {}

    parts.append('<rect id="background" width="500" height="500" fill="#000000"/>')

    # ---- ring: every tick measured individually
    groups = {}
    for gname, centre, width in tick_angles():
        raw = sample_tick(arr, centre, width)
        if raw.mean() < 12:          # slot occupied by an anchor dot / not drawn
            continue
        col = hx(raw)
        groups.setdefault(gname, []).append(
            f'<path d="{sector_path(centre, width)}" fill="{col}"/>')
    ring = []
    for g in ("top-right", "top-left", "bottom-left", "bottom-right"):
        ring.append(f'<g id="ring-{g}">\n    ' + "\n    ".join(groups[g]) + "\n  </g>")
    dot_col = hx(np.median(arr[4:13, 245:255].reshape(-1, 3), axis=0))
    dots = []
    for a in DOT_ANGLES:
        x, y = polar(a, DOT_R)
        dots.append(f'<circle cx="{x:.2f}" cy="{y:.2f}" r="{DOT_RAD}" fill="{dot_col}"/>')
    ring.append('<g id="ring-anchors">\n    ' + "\n    ".join(dots) + "\n  </g>")
    parts.append('<g id="ring">\n  ' + "\n  ".join(ring) + "\n</g>")
    tok["anchorDot"] = dot_col

    # ---- date band
    fill = arr[300, 110].astype(float)
    band = [f'<rect x="0" y="{BAND_Y}" width="500" height="{BAND_H}" fill="{hx(fill)}"/>']
    band_text = None
    for name, box in BAND_EL:
        a, e, x0, y0 = alpha_on_fill(arr, box, fill)
        d = trace(a, x0, y0)
        if not d:
            continue
        band_text = hx(e)
        band.append(f'<path id="{name}" d="{d}" fill="{band_text}" fill-rule="evenodd"/>')
    parts.append('<g id="date-band">\n  ' + "\n  ".join(band) + "\n</g>")
    tok["bandFill"], tok["bandText"] = hx(fill), band_text

    # ---- time
    tg = []
    for gid, name, box in (("gHours", "time-hours", TIME_L),
                           ("gMinutes", "time-minutes", TIME_R)):
        a, ramp, x0, y0, span = alpha_graded(arr, box)
        if span is None:
            continue
        defs.append(grad_def(gid, ramp, y0, span))
        tok[gid] = [hx(ramp[span[0]]), hx(ramp[span[1]])]
        tg.append((name, f'<path id="{name}" d="{trace(a, x0, y0)}" '
                         f'fill="url(#{gid})" fill-rule="evenodd"/>'))
    colon_col = hx(np.median(arr[188:200, 242:258].reshape(-1, 3), axis=0))
    tok["colon"] = colon_col
    colon = '<g id="time-colon">' + "".join(
        f'<rect x="{x}" y="{y}" width="{w}" height="{h}" fill="{colon_col}"/>'
        for x, y, w, h in COLON) + "</g>"
    parts.append('<g id="time">\n  ' + tg[0][1] + "\n  " + colon + "\n  " + tg[1][1] + "\n</g>")

    # ---- complications
    comp = []
    for name, box in SOLID:
        a, e, x0, y0 = alpha_solid(arr, box)
        d = trace(a, x0, y0)
        if not d:
            continue
        comp.append(f'<path id="{name}" d="{d}" fill="{hx(e)}" fill-rule="evenodd"/>')
        tok[name] = hx(e)
    parts.append('<g id="complications">\n  ' + "\n  ".join(comp) + "\n</g>")

    svg = ('<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 500 500" width="500" '
           'height="500" shape-rendering="geometricPrecision">\n'
           f'<title>Watch face {tid} {tname}</title>\n'
           f'<defs>{"".join(defs)}</defs>\n' + "\n".join(parts) + "\n</svg>\n")
    p = os.path.join(OUT, f"watchface-{tid}-{tname.lower()}.svg")
    open(p, "w").write(svg)

    # ring colour table for the spec
    tok["ringTop"] = {f"{TOP_PHASE + TOP_PITCH * k:.2f}":
                      hx(sample_tick(arr, TOP_PHASE + TOP_PITCH * k, TOP_W))
                      for k in range(TOP_N)}
    tok["ringBottom"] = {f"{BOT_PHASE + BOT_PITCH * k:.2f}":
                         hx(sample_tick(arr, 180 + BOT_PHASE + BOT_PITCH * k, BOT_W))
                         for k in range(BOT_N)}
    return p, tok


if __name__ == "__main__":
    tokens = {}
    for p, tid, tname in THEMES:
        out, tok = build(p, tid, tname)
        tok["name"] = tname
        tokens[tid] = tok
        print("wrote", out)
    json.dump(tokens, open(os.path.join(OUT, "watchface-tokens.json"), "w"), indent=2)