#!/usr/bin/env python3
"""
polyline2bezier.py - convert an SVG polyline path (M/L/H/V/Z) into a path of
cubic Bezier segments, with a user-chosen number of segments.

https://www.sqrt.ch/Buch/polyline2bezier.py (with claude.ai)

ODF arm-length fitting adapted to polylines according to
Alvin Penner, "Fitting a Cubic Bezier to a Parametric Function",
College Math. J. 50(3), 2019):

  * Each Bezier piece keeps its two endpoints and the endpoint tangent
    directions fixed, so only the two arm lengths r1, r2 are unknown.
  * r1, r2 are found by least-squares orthogonal distance fitting (ODF):
    alternate between (a) projecting every polyline sample onto the Bezier
    (Newton on the perpendicular-foot condition, Eq. 7) and (b) solving the
    resulting 2x2 linear least-squares problem for r1, r2.
  * Sharp corners (turn angle > --corner-angle) are kept as breaks; elsewhere
    neighbouring pieces share the tangent direction (G1 continuity).
  * Modes: 'count' spends N segments greedily (largest error is split first),
    'even' uses exactly N pieces of equal arc length, 'tolerance' adds
    segments until the max error is below a target.

Usage:
  python polyline2bezier.py --d "M0 0 L10 0 L10 10 L0 10 Z" -n 4
  python polyline2bezier.py in.svg -n 12 -o out.svg
  python polyline2bezier.py in.svg -m even -n 16
  python polyline2bezier.py in.svg -m tolerance --tol 0.01
  python polyline2bezier.py            (uses the settings block below)
Requires: numpy
"""

import argparse
import math
import re
import sys
import xml.etree.ElementTree as ET

import numpy as np

# ---------------------------------------------------------- user settings
# These apply when the script is run WITHOUT command-line arguments.
INPUT = "plot.svg"
OUTPUT = "out.svg"
MODE = "count"  # "count" | "even" | "tolerance"  (see fit_polyline)
N_SEG = 32  # count/even: number of segments; tolerance: upper limit (use e.g. 200)
TOL = 0.01  # tolerance mode: target max error in SVG units

NUM = r"[-+]?(?:\d+\.?\d*|\.\d+)(?:[eE][-+]?\d+)?"
TOKEN = re.compile(r"([A-Za-z])|(" + NUM + ")")


# --------------------------------------------------------------- parsing
def parse_path(d):
    """Parse a polyline-only path into a list of (points[k,2], closed)."""
    leftover = re.sub(NUM, "", d)
    bad = set(re.findall(r"[A-Za-z]", leftover)) - set("MmLlHhVvZz")
    if bad:
        raise ValueError(
            "Unsupported path commands %s: only M L H V Z (abs/rel) "
            "are accepted (input must be a piecewise-linear path)." % sorted(bad)
        )
    items = [m.group(1) or float(m.group(2)) for m in TOKEN.finditer(d)]
    subpaths, cur = [], None
    pos = np.zeros(2)
    start = pos.copy()
    cmd, i = None, 0

    def finish(closed):
        nonlocal cur
        if cur is not None and len(cur) > 1:
            subpaths.append((np.array(cur, float), closed))
        cur = None

    while i < len(items):
        if isinstance(items[i], str):
            cmd = items[i]
            i += 1
            if cmd in "Zz":
                if cur is not None:
                    finish(True)
                pos = start.copy()
                continue
        elif cmd is None:
            raise ValueError("Path data must start with a command")
        if cmd in "Zz":
            raise ValueError("Unexpected number after Z")
        n = {"M": 2, "L": 2, "H": 1, "V": 1}[cmd.upper()]
        args = items[i : i + n]
        if len(args) < n or any(isinstance(a, str) for a in args):
            raise ValueError("Bad arguments for command %s" % cmd)
        i += n
        rel = cmd.islower()
        c = cmd.upper()
        if c in "ML":
            new = np.array(args) + (pos if rel else 0)
        elif c == "H":
            new = np.array([args[0] + (pos[0] if rel else 0), pos[1]])
        else:
            new = np.array([pos[0], args[0] + (pos[1] if rel else 0)])
        if c == "M":
            finish(False)
            cur = [new]
            start = new.copy()
            cmd = "l" if rel else "L"  # implicit lineto after moveto
        else:
            if cur is None:
                cur = [pos.copy()]
                start = pos.copy()
            cur.append(new)
        pos = new
    finish(False)

    out = []
    for pts, closed in subpaths:
        keep = [0] + [
            k for k in range(1, len(pts)) if np.hypot(*(pts[k] - pts[k - 1])) > 1e-12
        ]
        pts = pts[keep]
        if len(pts) > 2 and np.hypot(*(pts[0] - pts[-1])) < 1e-12:
            pts, closed = pts[:-1], True  # explicit closure without Z
        elif closed and np.hypot(*(pts[0] - pts[-1])) < 1e-12:
            pts = pts[:-1]
        if len(pts) >= 2:
            out.append((pts, closed))
    return out


# ------------------------------------------------------------- polyline
class Polyline:
    def __init__(self, pts, closed):
        self.closed = closed
        P = np.vstack([pts, pts[0]]) if closed else pts
        seg = np.diff(P, axis=0)
        self.P = P
        self.seglen = np.hypot(seg[:, 0], seg[:, 1])
        self.n = len(self.seglen)
        self.cum = np.concatenate([[0], np.cumsum(self.seglen)])
        self.total = self.cum[-1]
        self.h = float(np.median(self.seglen))

    def at(self, s):
        s = np.atleast_1d(np.asarray(s, float))
        s = np.mod(s, self.total) if self.closed else np.clip(s, 0, self.total)
        idx = np.clip(np.searchsorted(self.cum, s, side="right") - 1, 0, self.n - 1)
        t = (s - self.cum[idx]) / self.seglen[idx]
        return self.P[idx] + t[:, None] * (self.P[idx + 1] - self.P[idx])

    def seg_dir(self, k):
        k = k % self.n if self.closed else min(max(k, 0), self.n - 1)
        v = self.P[k + 1] - self.P[k]
        return v / np.hypot(*v)

    def corners(self, thresh_deg):
        out = []
        rng = range(self.n) if self.closed else range(1, self.n)
        for i in rng:
            a = self.seg_dir(i - 1)
            b = self.seg_dir(i)
            ang = math.degrees(math.acos(max(-1.0, min(1.0, float(a @ b)))))
            if ang > thresh_deg:
                out.append(i)
        return out

    def smooth_tangent(self, s):
        # step = length of the local segment, so the estimate follows the
        # local sampling density (a global step is too coarse at sharp peaks)
        s0 = np.mod(s, self.total) if self.closed else min(max(s, 0.0), self.total)
        k = int(np.clip(np.searchsorted(self.cum, s0, side="right") - 1, 0, self.n - 1))
        h = self.seglen[k]
        v = self.at(s + h)[0] - self.at(s - h)[0]
        nv = np.hypot(*v)
        if nv < 1e-14:
            v = self.at(s + 1e-6 * self.total)[0] - self.at(s)[0]
            nv = np.hypot(*v) or 1.0
        return v / nv


# ----------------------------------------------------------- Bezier math
def bez_eval(c, u):
    u = u[:, None]
    w = 1 - u
    B = w**3 * c[0] + 3 * w**2 * u * c[1] + 3 * w * u**2 * c[2] + u**3 * c[3]
    d1 = 3 * (w**2 * (c[1] - c[0]) + 2 * w * u * (c[2] - c[1]) + u**2 * (c[3] - c[2]))
    d2 = 6 * (w * (c[2] - 2 * c[1] + c[0]) + u * (c[3] - 2 * c[2] + c[1]))
    return B, d1, d2


def project(c, G, u=None, iters=4, grid=64):
    """Foot points: coarse grid search (robust against local minima near
    tight turns), then Newton on (B(u)-G).B'(u)=0 (paper Eq. 7), u in [0,1]."""
    ug = np.linspace(0.0, 1.0, grid)
    Bg, _, _ = bez_eval(c, ug)
    d2 = ((G[:, None, :] - Bg[None, :, :]) ** 2).sum(axis=2)
    v = ug[np.argmin(d2, axis=1)]
    for _ in range(iters):
        B, d1, dd = bez_eval(c, v)
        diff = B - G
        f = np.sum(diff * d1, axis=1)
        fp = np.sum(d1 * d1, axis=1) + np.sum(diff * dd, axis=1)
        fp = np.where(fp > 1e-12, fp, np.sum(d1 * d1, axis=1) + 1e-12)
        v = np.clip(v - f / fp, 0.0, 1.0)
    return v


def fit_cubic(G, tA, tB, iters=80):
    """ODF fit of arm lengths r1, r2 with fixed endpoints and tangents."""
    p0, p3 = G[0], G[-1]
    dist = np.r_[0, np.cumsum(np.hypot(*np.diff(G, axis=0).T))]
    arc = dist[-1]
    u = dist / arc
    r = np.array([arc / 3, arc / 3])
    lo, hi = 1e-4 * arc, 3 * arc
    ctrl = None
    for _ in range(iters):
        b0, b1, b2, b3 = (1 - u) ** 3, 3 * u * (1 - u) ** 2, 3 * u**2 * (1 - u), u**3
        R = G - (np.outer(b0 + b1, p0) + np.outer(b2 + b3, p3))
        A1 = np.outer(b1, tA)
        A2 = -np.outer(b2, tB)
        a11, a12, a22 = np.sum(A1 * A1), np.sum(A1 * A2), np.sum(A2 * A2)
        det = a11 * a22 - a12**2
        if abs(det) > 1e-18 * max(a11 * a22, 1e-300):
            y1, y2 = np.sum(R * A1), np.sum(R * A2)
            rn = np.array([(a22 * y1 - a12 * y2) / det, (a11 * y2 - a12 * y1) / det])
            rn = np.clip(rn, lo, hi)
        else:
            rn = r
        done = np.max(np.abs(rn - r)) < 1e-7 * arc
        r = rn
        ctrl = np.array([p0, p0 + r[0] * tA, p3 - r[1] * tB, p3])
        u = project(ctrl, G, u)
        u[0], u[-1] = 0.0, 1.0
        if done:
            break
    B, _, _ = bez_eval(ctrl, u)
    err = np.hypot(*(B - G).T)
    return ctrl, float(err.max()), float(np.sqrt(np.mean(err**2)))


# ------------------------------------------------------------ the driver
class Piece:
    def __init__(self, sa, sb, ka=None, kb=None):
        self.sa, self.sb, self.ka, self.kb = sa, sb, ka, kb
        self.res = None


def fit_piece(poly, pc):
    nv = (pc.sb - pc.sa) / poly.total * poly.n
    ns = int(np.clip(2 * nv, 32, 400))
    s = np.linspace(pc.sa, pc.sb, ns)
    G = poly.at(s)
    if pc.ka is not None:
        tA = poly.seg_dir(pc.ka)
    else:
        tA = poly.smooth_tangent(pc.sa)
    if pc.kb is not None:
        tB = poly.seg_dir(pc.kb - 1)
    else:
        tB = poly.smooth_tangent(pc.sb)
    pc.res = fit_cubic(G, tA, tB)


def base_sections(poly, corner_deg):
    """Stretches of curve between mandatory breaks (corners / open ends)."""
    corners = poly.corners(corner_deg)
    cum, tot = poly.cum, poly.total
    if poly.closed:
        if not corners:
            return [Piece(0.0, tot)]
        secs = []
        for j, c in enumerate(corners):
            nxt = corners[(j + 1) % len(corners)]
            sb = cum[nxt] + (tot if nxt <= c else 0)
            secs.append(Piece(cum[c], sb, c, nxt))
        return secs
    vs = [0] + corners + [poly.n]
    return [Piece(cum[a], cum[b], a, b) for a, b in zip(vs[:-1], vs[1:])]


def split_even(sec, k):
    """Split a section into k pieces of equal arc length."""
    out = []
    for j in range(k):
        sa = sec.sa if j == 0 else sec.sa + j * (sec.sb - sec.sa) / k
        sb = sec.sb if j == k - 1 else sec.sa + (j + 1) * (sec.sb - sec.sa) / k
        out.append(
            Piece(sa, sb, sec.ka if j == 0 else None, sec.kb if j == k - 1 else None)
        )
    return out


def allocate(secs, n):
    """Distribute n segments over sections in proportion to their length."""
    if n <= len(secs):
        return [1] * len(secs)
    L = [s.sb - s.sa for s in secs]
    ideal = [n * l / sum(L) for l in L]
    k = [max(1, int(math.floor(x))) for x in ideal]
    while sum(k) < n:
        i = max(range(len(k)), key=lambda i: ideal[i] - k[i])
        k[i] += 1
    while sum(k) > n:
        i = max((i for i in range(len(k)) if k[i] > 1), key=lambda i: k[i] - ideal[i])
        k[i] -= 1
    return k


MODES = ("count", "even", "tolerance")


def _fit_all(poly, pieces):
    for pc in pieces:
        fit_piece(poly, pc)
    return pieces


def _max_err(pieces):
    return max(pc.res[1] for pc in pieces)


def build_even(poly, secs, n):
    pieces = []
    for sec, k in zip(secs, allocate(secs, n)):
        pieces += split_even(sec, k)
    return _fit_all(poly, pieces)


def build_greedy(poly, secs, n, stop):
    pieces = list(secs)
    if poly.closed and len(secs) == 1 and secs[0].ka is None:
        pieces = split_even(secs[0], max(1, min(n, 4)))
    _fit_all(poly, pieces)
    while len(pieces) < n:
        worst = max(range(len(pieces)), key=lambda i: pieces[i].res[1])
        pc = pieces[worst]
        if pc.res[1] <= stop:
            break  # target reached / already exact
        mid = 0.5 * (pc.sa + pc.sb)
        a, b = Piece(pc.sa, mid, pc.ka, None), Piece(mid, pc.sb, None, pc.kb)
        fit_piece(poly, a)
        fit_piece(poly, b)
        pieces[worst : worst + 1] = [a, b]
    return pieces


def fit_polyline(
    pts, closed, n_seg, corner_deg=35.0, mode="count", tol=0.01, rel_tol=1e-7
):
    """
    mode 'count'     : up to n_seg segments, split where the error is largest
    mode 'even'      : exactly n_seg segments of equal arc length
    mode 'tolerance' : error-driven splitting until max error <= tol (n_seg = limit)
    """
    if mode not in MODES:
        raise ValueError("mode must be one of %s" % (MODES,))
    poly = Polyline(pts, closed)
    secs = base_sections(poly, corner_deg)
    if len(secs) > n_seg:
        print(
            "warning: %d sharp corners force %d segments (> limit %d)"
            % (len(secs), len(secs), n_seg),
            file=sys.stderr,
        )
    diag = float(np.hypot(*(pts.max(0) - pts.min(0)))) or 1.0

    if mode == "even":
        pieces = build_even(poly, secs, n_seg)
    elif mode == "count":
        pieces = build_greedy(poly, secs, n_seg, rel_tol * diag)
    else:
        tol = max(tol, rel_tol * diag)
        pieces = build_greedy(poly, secs, n_seg, tol)
        if _max_err(pieces) > tol:
            print(
                "warning: tolerance %g not reached with the limit of %d "
                "segments (max err %.4g); raise the segment limit"
                % (tol, n_seg, _max_err(pieces)),
                file=sys.stderr,
            )

    segs = [pc.res[0] for pc in pieces]
    if closed:  # make the loop close exactly
        segs[-1][3] = segs[0][0]
    maxerr = _max_err(pieces)
    rms = math.sqrt(sum(pc.res[2] ** 2 for pc in pieces) / len(pieces))
    return segs, maxerr, rms


def fmt(x, p):
    s = ("%.*f" % (p, x)).rstrip("0").rstrip(".")
    return "0" if s in ("", "-0") else s


def polyline_to_bezier_path(
    d, n_seg=8, corner_deg=35.0, precision=3, verbose=True, mode="count", tol=0.01
):
    """Convert polyline path data `d` to a cubic Bezier path string."""
    out = []
    for pts, closed in parse_path(d):
        segs, maxerr, rms = fit_polyline(pts, closed, n_seg, corner_deg, mode, tol)
        if verbose:
            print(
                "subpath: %d points -> %d cubic segments, max err %.4g, rms %.4g"
                % (len(pts), len(segs), maxerr, rms),
                file=sys.stderr,
            )
        f = lambda v: fmt(v, precision)
        s = "M%s %s" % (f(segs[0][0][0]), f(segs[0][0][1]))
        for c in segs:
            s += " C%s %s %s %s %s %s" % tuple(f(v) for v in c[1:].ravel())
        out.append(s + (" Z" if closed else ""))
    return " ".join(out)


# ------------------------------------------------------------ SVG files
def convert_svg(src, dst, **kw):
    for pfx, uri in [
        ("", "http://www.w3.org/2000/svg"),
        ("xlink", "http://www.w3.org/1999/xlink"),
        ("inkscape", "http://www.inkscape.org/namespaces/inkscape"),
        ("sodipodi", "http://sodipodi.sourceforge.net/DTD/sodipodi-0.dtd"),
    ]:
        ET.register_namespace(pfx, uri)
    tree = ET.parse(src)
    ns = "http://www.w3.org/2000/svg"
    count = 0
    for el in tree.iter():
        tag = el.tag.rsplit("}", 1)[-1] if isinstance(el.tag, str) else ""
        if tag == "path" and "d" in el.attrib:
            el.set("d", polyline_to_bezier_path(el.get("d"), **kw))
            count += 1
        elif tag in ("polyline", "polygon") and "points" in el.attrib:
            v = [float(x) for x in re.findall(NUM, el.get("points"))]
            pts = list(zip(v[0::2], v[1::2]))
            d = (
                "M"
                + " L".join("%r %r" % p for p in pts)
                + (" Z" if tag == "polygon" else "")
            )
            el.tag = "{%s}path" % ns
            del el.attrib["points"]
            el.set("d", polyline_to_bezier_path(d, **kw))
            count += 1
    tree.write(dst, encoding="utf-8", xml_declaration=True)
    return count


def main():
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("input", nargs="?", help="input SVG file")
    ap.add_argument("--d", help="path data string instead of a file")
    ap.add_argument(
        "-m",
        "--mode",
        choices=MODES,
        default=MODE,
        help="count: up to N segments, error-driven splitting; "
        "even: exactly N equal-length segments; "
        "tolerance: fewest segments with max error <= --tol",
    )
    ap.add_argument(
        "-n",
        "--segments",
        type=int,
        default=None,
        help="segments per subpath (in tolerance mode: upper limit); "
        "default 8, or 200 in tolerance mode",
    )
    ap.add_argument(
        "--tol",
        type=float,
        default=TOL,
        help="target max error for tolerance mode (SVG units)",
    )
    ap.add_argument(
        "-o", "--output", help="output SVG file (default: <input>_bezier.svg)"
    )
    ap.add_argument(
        "--corner-angle",
        type=float,
        default=35.0,
        help="turn angle in degrees above which a vertex is a sharp corner",
    )
    ap.add_argument("--precision", type=int, default=3, help="decimals in output")
    a = ap.parse_args()
    if a.segments is None:
        a.segments = 200 if a.mode == "tolerance" else 8
    if a.segments < 1:
        ap.error("--segments must be >= 1")
    kw = dict(
        n_seg=a.segments,
        corner_deg=a.corner_angle,
        precision=a.precision,
        mode=a.mode,
        tol=a.tol,
    )
    if a.d:
        print(polyline_to_bezier_path(a.d, **kw))
    elif a.input:
        dst = a.output or re.sub(r"(\.svg)?$", "_bezier.svg", a.input, count=1)
        k = convert_svg(a.input, dst, **kw)
        print("converted %d element(s) -> %s" % (k, dst), file=sys.stderr)
    else:
        ap.error("give an input SVG or --d")


if __name__ == "__main__":
    if len(sys.argv) > 1:
        main()  # command-line use
    else:
        convert_svg(INPUT, OUTPUT, n_seg=N_SEG, mode=MODE, tol=TOL)